DebTik/internal/firewall/manager.go

package firewall

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"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
	lastError    string
	applied      bool
	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.geoIPMgr.referenced = m.geoReferenced
	loadErr := m.load()
	m.remember()
	if loadErr != nil {
		m.lastError = loadErr.Error()
		log.Printf("Firewall configuration not loaded: %v", loadErr)
		return m
	}
	// Defaults are installed only once, never because a user intentionally saved
	// an empty list or deleted all rules.
	fresh := true
	for _, path := range []string{FilterStorePath, NATStorePath, MangleStorePath, AddressListsStorePath} {
		if _, err := os.Stat(path); !os.IsNotExist(err) {
			fresh = false
		}
	}
	if fresh {
		if err := m.initDefaultRules(); err != nil {
			m.lastError = err.Error()
			return m
		}
	}
	if err := m.Apply(); err != nil {
		log.Printf("Firewall not applied: %v", err)
	}
	m.remember()
	return m
}

func (m *FirewallManager) initDefaultRules() error {
	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,
		})
	}

	return m.save()
}

func (m *FirewallManager) load() error {
	var failures []error
	for path, target := range map[string]any{FilterStorePath: &m.filterRules, NATStorePath: &m.natRules, MangleStorePath: &m.mangleRules, AddressListsStorePath: &m.addressLists} {
		data, err := os.ReadFile(path)
		if os.IsNotExist(err) {
			continue
		}
		if err == nil {
			err = json.Unmarshal(data, target)
		}
		if err != nil {
			failures = append(failures, fmt.Errorf("%s: %w", path, err))
		}
	}
	if m.filterRules == nil {
		m.filterRules = []FilterRule{}
	}
	if m.natRules == nil {
		m.natRules = []NATRule{}
	}
	if m.mangleRules == nil {
		m.mangleRules = []MangleRule{}
	}
	if m.addressLists == nil {
		m.addressLists = []AddressListEntry{}
	}
	return errors.Join(failures...)
}

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
	}
	files := map[string][]byte{}
	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
		}
		files[path] = data
	}
	return saveFileGroup(files)
}
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 {
			m.lastError = fmt.Sprintf("firewall change failed: %v; rollback apply: %v; rollback save: %v", err, rollbackErr, saveErr)
			m.applied = false
			return errors.New(m.lastError)
		}
		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() (result error) {
	defer func() {
		m.applied = result == nil
		m.lastError = ""
		if result != nil {
			m.lastError = result.Error()
		}
	}()
	for _, r := range m.filterRules {
		if err := validateRule("filter", r); err != nil {
			return err
		}
	}
	for _, r := range m.natRules {
		if err := validateRule("nat", r); err != nil {
			return err
		}
	}
	for _, r := range m.mangleRules {
		if err := validateRule("mangle", r); err != nil {
			return err
		}
	}

	// 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
	}

	if err := ApplyDualStack(m.filterRules, m.natRules, m.mangleRules); err != nil {
		return err
	}
	collectUnusedAddressLists(m.addressLists, refLists)
	return nil
}

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 append(res, m.dynamicEntries()...)
}

func (m *FirewallManager) SaveAddressListEntry(e AddressListEntry) error {
	if e.Dynamic || e.Timeout != 0 {
		return fmt.Errorf("dynamic entries are managed by packet rules")
	}
	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 {
	if strings.HasPrefix(id, "dynamic-") {
		m.mu.Lock()
		defer m.mu.Unlock()
		return deleteDynamicEntry(id)
	}
	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, connErr := ListActiveConnections()
	connectionError := ""
	if connErr != nil {
		connectionError = connErr.Error()
	}

	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),
		ConnectionsError:  connectionError,
		Applied:           m.applied, LastError: m.lastError,
	}
}

// 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 {
		if err := m.commitChange(); err != nil {
			m.lastError = err.Error()
			log.Printf("Firewall rename failed: %v", err)
		}
	}
}