DebTik/internal/firewall/manager.go

package firewall

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"
)

const (
	FilterStorePath       = "/etc/proxmox-router/firewall/filter.json"
	NATStorePath          = "/etc/proxmox-router/firewall/nat.json"
	MangleStorePath       = "/etc/proxmox-router/firewall/mangle.json"
	AddressListsStorePath = "/etc/proxmox-router/firewall/address_lists.json"
)

type FirewallManager struct {
	mu           sync.Mutex
	committed    []byte
	filterRules  []FilterRule
	natRules     []NATRule
	mangleRules  []MangleRule
	addressLists []AddressListEntry
	geoIPMgr     *GeoIPManager
}

func NewFirewallManager() *FirewallManager {
	m := &FirewallManager{
		filterRules:  make([]FilterRule, 0),
		natRules:     make([]NATRule, 0),
		mangleRules:  make([]MangleRule, 0),
		addressLists: make([]AddressListEntry, 0),
		geoIPMgr:     NewGeoIPManager(),
	}
	m.load()
	m.remember()

	// If clean installation and no rules exist, initialize default sensible router rules (like RouterOS default firewall)
	if len(m.filterRules) == 0 || len(m.natRules) == 0 {
		m.initDefaultRules()
	}

	_ = m.Apply()
	m.remember()
	return m
}

func (m *FirewallManager) initDefaultRules() {
	if len(m.natRules) == 0 {
		// Default Masquerade NAT only on WAN
		m.natRules = append(m.natRules, NATRule{
			ID:               fmt.Sprintf("nat-%d", time.Now().UnixNano()),
			Order:            1,
			Chain:            "srcnat",
			Action:           "masquerade",
			OutInterfaceList: "WAN",
			Comment:          "defconf: masquerade outgoing traffic to WAN",
			Disabled:         false,
		})
	}

	if len(m.filterRules) == 0 {
		// Default Fasttrack / Established,Related Filter
		m.filterRules = append(m.filterRules, FilterRule{
			ID:              fmt.Sprintf("filter-%d-1", time.Now().UnixNano()),
			Order:           1,
			Chain:           "input",
			Action:          "accept",
			ConnectionState: []string{"established", "related"},
			Comment:         "defconf: accept established,related",
			Disabled:        false,
		})

		m.filterRules = append(m.filterRules, FilterRule{
			ID:              fmt.Sprintf("filter-%d-2", time.Now().UnixNano()),
			Order:           2,
			Chain:           "forward",
			Action:          "accept",
			ConnectionState: []string{"established", "related"},
			Comment:         "defconf: accept established,related forward",
			Disabled:        false,
		})
	}

	_ = m.save()
}

func (m *FirewallManager) load() {
	if data, err := os.ReadFile(FilterStorePath); err == nil {
		_ = json.Unmarshal(data, &m.filterRules)
	}
	if data, err := os.ReadFile(NATStorePath); err == nil {
		_ = json.Unmarshal(data, &m.natRules)
	}
	if data, err := os.ReadFile(MangleStorePath); err == nil {
		_ = json.Unmarshal(data, &m.mangleRules)
	}
	if data, err := os.ReadFile(AddressListsStorePath); err == nil {
		_ = json.Unmarshal(data, &m.addressLists)
	}

	if m.filterRules == nil {
		m.filterRules = make([]FilterRule, 0)
	}
	if m.natRules == nil {
		m.natRules = make([]NATRule, 0)
	}
	if m.mangleRules == nil {
		m.mangleRules = make([]MangleRule, 0)
	}
	if m.addressLists == nil {
		m.addressLists = make([]AddressListEntry, 0)
	}
}

func (m *FirewallManager) snapshot() []byte {
	data, _ := json.Marshal([]any{m.filterRules, m.natRules, m.mangleRules, m.addressLists})
	return data
}
func (m *FirewallManager) remember() { m.committed = m.snapshot() }
func (m *FirewallManager) restore(data []byte) {
	var fields []json.RawMessage
	if json.Unmarshal(data, &fields) == nil && len(fields) == 4 {
		_ = json.Unmarshal(fields[0], &m.filterRules)
		_ = json.Unmarshal(fields[1], &m.natRules)
		_ = json.Unmarshal(fields[2], &m.mangleRules)
		_ = json.Unmarshal(fields[3], &m.addressLists)
	}
}
func (m *FirewallManager) save() error {
	if err := os.MkdirAll("/etc/proxmox-router/firewall", 0755); err != nil {
		return err
	}
	for path, value := range map[string]any{FilterStorePath: m.filterRules, NATStorePath: m.natRules, MangleStorePath: m.mangleRules, AddressListsStorePath: m.addressLists} {
		data, err := json.MarshalIndent(value, "", "  ")
		if err != nil {
			return err
		}
		f, err := os.CreateTemp("/etc/proxmox-router/firewall", ".config-*")
		if err != nil {
			return err
		}
		_, err = f.Write(data)
		closeErr := f.Close()
		if err == nil {
			err = closeErr
		}
		if err == nil {
			err = os.Rename(f.Name(), path)
		}
		_ = os.Remove(f.Name())
		if err != nil {
			return err
		}
	}
	return nil
}
func (m *FirewallManager) commitChange() error {
	before := append([]byte(nil), m.committed...)
	err := m.Apply()
	if err == nil {
		err = m.save()
	}
	if err != nil {
		m.restore(before)
		rollbackErr := m.Apply()
		saveErr := m.save()
		if rollbackErr != nil || saveErr != nil {
			return fmt.Errorf("firewall change failed: %v; rollback apply: %v; rollback save: %v", err, rollbackErr, saveErr)
		}
		return err
	}
	m.remember()
	return nil
}
func (m *FirewallManager) Reapply() error { m.mu.Lock(); defer m.mu.Unlock(); return m.Apply() }

// Apply re-synchronizes AddressLists & GeoIP sets with ipset and applies iptables
func (m *FirewallManager) Apply() error {
	// Sync GeoIP country sets
	if m.geoIPMgr != nil {

		// Also ensure any GeoIP lists explicitly referenced in active rules are loaded into ipset
		for _, r := range m.filterRules {
			if !r.Disabled {
				if r.SrcGeoIP != "" {
					if err := m.geoIPMgr.EnsureCountryOrPreset(r.SrcGeoIP); err != nil {
						return err
					}
				}
				if r.DstGeoIP != "" {
					if err := m.geoIPMgr.EnsureCountryOrPreset(r.DstGeoIP); err != nil {
						return err
					}
				}
			}
		}
		for _, r := range m.natRules {
			if !r.Disabled {
				if r.SrcGeoIP != "" {
					if err := m.geoIPMgr.EnsureCountryOrPreset(r.SrcGeoIP); err != nil {
						return err
					}
				}
				if r.DstGeoIP != "" {
					if err := m.geoIPMgr.EnsureCountryOrPreset(r.DstGeoIP); err != nil {
						return err
					}
				}
			}
		}
		for _, r := range m.mangleRules {
			if !r.Disabled {
				if r.SrcGeoIP != "" {
					if err := m.geoIPMgr.EnsureCountryOrPreset(r.SrcGeoIP); err != nil {
						return err
					}
				}
				if r.DstGeoIP != "" {
					if err := m.geoIPMgr.EnsureCountryOrPreset(r.DstGeoIP); err != nil {
						return err
					}
				}
			}
		}
	}

	// Collect any custom address lists referenced in active rules
	var refLists []string
	for _, r := range m.filterRules {
		if !r.Disabled {
			if r.SrcAddressList != "" {
				refLists = append(refLists, r.SrcAddressList)
			}
			if r.DstAddressList != "" {
				refLists = append(refLists, r.DstAddressList)
			}
			if r.AddressList != "" {
				refLists = append(refLists, r.AddressList)
			}
		}
	}
	for _, r := range m.natRules {
		if !r.Disabled {
			if r.SrcAddressList != "" {
				refLists = append(refLists, r.SrcAddressList)
			}
			if r.DstAddressList != "" {
				refLists = append(refLists, r.DstAddressList)
			}
			if r.AddressList != "" {
				refLists = append(refLists, r.AddressList)
			}
		}
	}
	for _, r := range m.mangleRules {
		if !r.Disabled {
			if r.SrcAddressList != "" {
				refLists = append(refLists, r.SrcAddressList)
			}
			if r.DstAddressList != "" {
				refLists = append(refLists, r.DstAddressList)
			}
			if r.AddressList != "" {
				refLists = append(refLists, r.AddressList)
			}
		}
	}

	// Sync custom ipset address lists
	if err := SyncAddressLists(m.addressLists, refLists...); err != nil {
		return err
	}

	return ApplyDualStack(m.filterRules, m.natRules, m.mangleRules)
}

func (m *FirewallManager) GeoIP() *GeoIPManager {
	return m.geoIPMgr
}

// ==================== Filter Rules ====================

func (m *FirewallManager) ListFilterRules() []FilterRule {
	m.mu.Lock()
	defer m.mu.Unlock()

	counters, _ := ReadRuleCounters(m.filterRules, m.natRules, m.mangleRules)

	res := make([]FilterRule, len(m.filterRules))
	for i, r := range m.filterRules {
		res[i] = r
		if c, ok := counters[r.ID]; ok {
			res[i].Packets = c[0]
			res[i].Bytes = c[1]
		}
	}
	return res
}

func (m *FirewallManager) SaveFilterRule(r FilterRule) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if r.ID == "" {
		r.ID = fmt.Sprintf("filter-%d", time.Now().UnixNano())
	}

	found := false
	for i, existing := range m.filterRules {
		if existing.ID == r.ID {
			r.Order = existing.Order
			m.filterRules[i] = r
			found = true
			break
		}
	}
	if !found {
		r.Order = len(m.filterRules) + 1
		m.filterRules = append(m.filterRules, r)
	}

	return m.commitChange()
}

func (m *FirewallManager) DeleteFilterRule(id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	var updated []FilterRule
	for _, r := range m.filterRules {
		if r.ID != id {
			updated = append(updated, r)
		}
	}
	for i := range updated {
		updated[i].Order = i + 1
	}
	m.filterRules = updated

	return m.commitChange()
}

func (m *FirewallManager) MoveFilterRule(id string, dir string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	idx := -1
	for i, r := range m.filterRules {
		if r.ID == id {
			idx = i
			break
		}
	}
	if idx == -1 {
		return fmt.Errorf("rule not found")
	}

	if dir == "up" && idx > 0 {
		m.filterRules[idx], m.filterRules[idx-1] = m.filterRules[idx-1], m.filterRules[idx]
	} else if dir == "down" && idx < len(m.filterRules)-1 {
		m.filterRules[idx], m.filterRules[idx+1] = m.filterRules[idx+1], m.filterRules[idx]
	}

	for i := range m.filterRules {
		m.filterRules[i].Order = i + 1
	}

	return m.commitChange()
}

func (m *FirewallManager) SetFilterRuleState(id string, disabled bool) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	for i := range m.filterRules {
		if m.filterRules[i].ID == id {
			m.filterRules[i].Disabled = disabled
			return m.commitChange()
		}
	}
	return fmt.Errorf("rule not found")
}

// ==================== NAT Rules ====================

func (m *FirewallManager) ListNATRules() []NATRule {
	m.mu.Lock()
	defer m.mu.Unlock()

	counters, _ := ReadRuleCounters(m.filterRules, m.natRules, m.mangleRules)

	res := make([]NATRule, len(m.natRules))
	for i, r := range m.natRules {
		res[i] = r
		if c, ok := counters[r.ID]; ok {
			res[i].Packets = c[0]
			res[i].Bytes = c[1]
		}
	}
	return res
}

func (m *FirewallManager) SaveNATRule(r NATRule) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if r.ID == "" {
		r.ID = fmt.Sprintf("nat-%d", time.Now().UnixNano())
	}

	found := false
	for i, existing := range m.natRules {
		if existing.ID == r.ID {
			r.Order = existing.Order
			m.natRules[i] = r
			found = true
			break
		}
	}
	if !found {
		r.Order = len(m.natRules) + 1
		m.natRules = append(m.natRules, r)
	}

	return m.commitChange()
}

func (m *FirewallManager) DeleteNATRule(id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	var updated []NATRule
	for _, r := range m.natRules {
		if r.ID != id {
			updated = append(updated, r)
		}
	}
	for i := range updated {
		updated[i].Order = i + 1
	}
	m.natRules = updated

	return m.commitChange()
}

func (m *FirewallManager) MoveNATRule(id string, dir string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	idx := -1
	for i, r := range m.natRules {
		if r.ID == id {
			idx = i
			break
		}
	}
	if idx == -1 {
		return fmt.Errorf("rule not found")
	}

	if dir == "up" && idx > 0 {
		m.natRules[idx], m.natRules[idx-1] = m.natRules[idx-1], m.natRules[idx]
	} else if dir == "down" && idx < len(m.natRules)-1 {
		m.natRules[idx], m.natRules[idx+1] = m.natRules[idx+1], m.natRules[idx]
	}

	for i := range m.natRules {
		m.natRules[i].Order = i + 1
	}

	return m.commitChange()
}

func (m *FirewallManager) SetNATRuleState(id string, disabled bool) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	for i := range m.natRules {
		if m.natRules[i].ID == id {
			m.natRules[i].Disabled = disabled
			return m.commitChange()
		}
	}
	return fmt.Errorf("rule not found")
}

// ==================== Mangle Rules ====================

func (m *FirewallManager) ListMangleRules() []MangleRule {
	m.mu.Lock()
	defer m.mu.Unlock()

	counters, _ := ReadRuleCounters(m.filterRules, m.natRules, m.mangleRules)

	res := make([]MangleRule, len(m.mangleRules))
	for i, r := range m.mangleRules {
		res[i] = r
		if c, ok := counters[r.ID]; ok {
			res[i].Packets = c[0]
			res[i].Bytes = c[1]
		}
	}
	return res
}

func (m *FirewallManager) SaveMangleRule(r MangleRule) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if r.ID == "" {
		r.ID = fmt.Sprintf("mangle-%d", time.Now().UnixNano())
	}

	found := false
	for i, existing := range m.mangleRules {
		if existing.ID == r.ID {
			r.Order = existing.Order
			m.mangleRules[i] = r
			found = true
			break
		}
	}
	if !found {
		r.Order = len(m.mangleRules) + 1
		m.mangleRules = append(m.mangleRules, r)
	}

	return m.commitChange()
}

func (m *FirewallManager) DeleteMangleRule(id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	var updated []MangleRule
	for _, r := range m.mangleRules {
		if r.ID != id {
			updated = append(updated, r)
		}
	}
	for i := range updated {
		updated[i].Order = i + 1
	}
	m.mangleRules = updated

	return m.commitChange()
}

func (m *FirewallManager) MoveMangleRule(id string, dir string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	idx := -1
	for i, r := range m.mangleRules {
		if r.ID == id {
			idx = i
			break
		}
	}
	if idx == -1 {
		return fmt.Errorf("rule not found")
	}

	if dir == "up" && idx > 0 {
		m.mangleRules[idx], m.mangleRules[idx-1] = m.mangleRules[idx-1], m.mangleRules[idx]
	} else if dir == "down" && idx < len(m.mangleRules)-1 {
		m.mangleRules[idx], m.mangleRules[idx+1] = m.mangleRules[idx+1], m.mangleRules[idx]
	}

	for i := range m.mangleRules {
		m.mangleRules[i].Order = i + 1
	}

	return m.commitChange()
}

func (m *FirewallManager) SetMangleRuleState(id string, disabled bool) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	for i := range m.mangleRules {
		if m.mangleRules[i].ID == id {
			m.mangleRules[i].Disabled = disabled
			return m.commitChange()
		}
	}
	return fmt.Errorf("rule not found")
}

// ==================== Address Lists ====================

func (m *FirewallManager) ListAddressLists() []AddressListEntry {
	m.mu.Lock()
	defer m.mu.Unlock()

	res := make([]AddressListEntry, len(m.addressLists))
	copy(res, m.addressLists)
	return res
}

func (m *FirewallManager) SaveAddressListEntry(e AddressListEntry) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if e.ID == "" {
		e.ID = fmt.Sprintf("addrlist-%d", time.Now().UnixNano())
	}

	found := false
	for i, existing := range m.addressLists {
		if existing.ID == e.ID {
			m.addressLists[i] = e
			found = true
			break
		}
	}
	if !found {
		m.addressLists = append(m.addressLists, e)
	}

	return m.commitChange()
}

func (m *FirewallManager) DeleteAddressListEntry(id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	var updated []AddressListEntry
	for _, e := range m.addressLists {
		if e.ID != id {
			updated = append(updated, e)
		}
	}
	m.addressLists = updated

	return m.commitChange()
}

func (m *FirewallManager) SetAddressListEntryState(id string, disabled bool) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	for i := range m.addressLists {
		if m.addressLists[i].ID == id {
			m.addressLists[i].Disabled = disabled
			return m.commitChange()
		}
	}
	return fmt.Errorf("address list entry not found")
}

func (m *FirewallManager) ResetCounters() error {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, binary := range []string{"iptables", "ip6tables"} {
		out, err := exec.Command(binary + "-save").Output()
		if err != nil {
			if binary == "iptables" {
				return err
			}
			continue
		}
		table := ""
		for _, line := range strings.Split(string(out), "\n") {
			if strings.HasPrefix(line, "*") {
				table = line[1:]
			}
			if strings.HasPrefix(line, ":DT_") {
				chain := strings.Fields(line)[0][1:]
				if err := exec.Command(binary, "-t", table, "-Z", chain).Run(); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func (m *FirewallManager) GetSummary() FirewallSummary {
	m.mu.Lock()
	defer m.mu.Unlock()

	conns, _ := ListActiveConnections()

	geoCount := 0
	geoSubnets := 0
	if m.geoIPMgr != nil {
		status := m.geoIPMgr.GetStatus()
		if c, ok := status["enabled_count"].(int); ok {
			geoCount = c
		}
		if s, ok := status["total_subnets"].(int); ok {
			geoSubnets = s
		}
	}

	return FirewallSummary{
		FilterRulesCount:  len(m.filterRules),
		NATRulesCount:     len(m.natRules),
		MangleRulesCount:  len(m.mangleRules),
		AddressListsCount: len(m.addressLists),
		GeoCountriesCount: geoCount,
		GeoSubnetsCount:   geoSubnets,
		ActiveConnections: len(conns),
	}
}

// RenameInterface updates firewall rule interface names when an interface is renamed.
func (m *FirewallManager) RenameInterface(oldName, newName string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	changed := false
	for i := range m.filterRules {
		if m.filterRules[i].InInterface == oldName {
			m.filterRules[i].InInterface = newName
			changed = true
		}
		if m.filterRules[i].OutInterface == oldName {
			m.filterRules[i].OutInterface = newName
			changed = true
		}
	}
	for i := range m.natRules {
		if m.natRules[i].InInterface == oldName {
			m.natRules[i].InInterface = newName
			changed = true
		}
		if m.natRules[i].OutInterface == oldName {
			m.natRules[i].OutInterface = newName
			changed = true
		}
	}
	for i := range m.mangleRules {
		if m.mangleRules[i].InInterface == oldName {
			m.mangleRules[i].InInterface = newName
			changed = true
		}
		if m.mangleRules[i].OutInterface == oldName {
			m.mangleRules[i].OutInterface = newName
			changed = true
		}
	}
	if changed {
		_ = m.save()
		_ = m.Apply()
	}
}