DebTik/internal/firewall/manager.go

package firewall

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"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
	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()

	// 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()
	return m
}

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

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

func (m *FirewallManager) save() error {
	_ = os.MkdirAll("/etc/proxmox-router/firewall", 0755)

	if d, err := json.MarshalIndent(m.filterRules, "", "  "); err == nil {
		_ = os.WriteFile(FilterStorePath, d, 0600)
	}
	if d, err := json.MarshalIndent(m.natRules, "", "  "); err == nil {
		_ = os.WriteFile(NATStorePath, d, 0600)
	}
	if d, err := json.MarshalIndent(m.mangleRules, "", "  "); err == nil {
		_ = os.WriteFile(MangleStorePath, d, 0600)
	}
	if d, err := json.MarshalIndent(m.addressLists, "", "  "); err == nil {
		_ = os.WriteFile(AddressListsStorePath, d, 0600)
	}
	return nil
}

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

		// 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 != "" { _ = m.geoIPMgr.EnsureCountryOrPreset(r.SrcGeoIP) }
				if r.DstGeoIP != "" { _ = m.geoIPMgr.EnsureCountryOrPreset(r.DstGeoIP) }
			}
		}
		for _, r := range m.natRules {
			if !r.Disabled {
				if r.SrcGeoIP != "" { _ = m.geoIPMgr.EnsureCountryOrPreset(r.SrcGeoIP) }
				if r.DstGeoIP != "" { _ = m.geoIPMgr.EnsureCountryOrPreset(r.DstGeoIP) }
			}
		}
		for _, r := range m.mangleRules {
			if !r.Disabled {
				if r.SrcGeoIP != "" { _ = m.geoIPMgr.EnsureCountryOrPreset(r.SrcGeoIP) }
				if r.DstGeoIP != "" { _ = m.geoIPMgr.EnsureCountryOrPreset(r.DstGeoIP) }
			}
		}
	}

	// 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
	_ = SyncAddressLists(m.addressLists, refLists...)

	payload := BuildIPTablesRestore(m.filterRules, m.natRules, m.mangleRules)
	return ApplyIPTables(payload)
}

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

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

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

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

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
	}

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

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
			_ = m.save()
			return m.Apply()
		}
	}
	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)
	}

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

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

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

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
	}

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

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
			_ = m.save()
			return m.Apply()
		}
	}
	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)
	}

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

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

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

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
	}

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

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
			_ = m.save()
			return m.Apply()
		}
	}
	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)
	}

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

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

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

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
			_ = m.save()
			return m.Apply()
		}
	}
	return fmt.Errorf("address list entry not found")
}

func (m *FirewallManager) ResetCounters() error {
	_ = exec.Command("iptables", "-Z").Run()
	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),
	}
}