DebTik/internal/firewall/geoip.go

package firewall

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	GeoIPDir        = "/etc/proxmox-router/firewall/geoip"
	GeoIPConfigFile = "/etc/proxmox-router/firewall/geoip/config.json"
	GeoIPZonesDir   = "/etc/proxmox-router/firewall/geoip/zones"
)

// GeoIPConfig defines interval and enabled country persistence
type GeoIPConfig struct {
	AutoUpdate       bool      `json:"auto_update"`
	IntervalHours    int       `json:"interval_hours"` // e.g. 24
	LastUpdated      time.Time `json:"last_updated"`
	NextUpdate       time.Time `json:"next_update"`
	EnabledCountries []string  `json:"enabled_countries"`
}

// GeoIPManager handles country IP databases, ipset sync, and auto-update scheduling
type GeoIPManager struct {
	mu          sync.RWMutex
	config      GeoIPConfig
	countries   map[string]*GeoCountryInfo
	timer       *time.Timer
	stopCh      chan struct{}
	restoreDone chan struct{}
	isUpdating  bool
	lastError   string
	referenced  func(string) bool
	syncMu      sync.Mutex
}

func NewGeoIPManager() *GeoIPManager {
	_ = os.MkdirAll(GeoIPZonesDir, 0755)

	m := &GeoIPManager{
		countries:   make(map[string]*GeoCountryInfo),
		stopCh:      make(chan struct{}),
		restoreDone: make(chan struct{}),
	}

	// 1. Initialize country metadata
	for _, c := range StandardISOCountries {
		info := c
		info.IPSetName = "mk_geo_" + c.Code
		m.countries[c.Code] = &info
	}

	// 2. Load Config
	m.loadConfig()

	// 3. Scan cached zone files
	m.scanCachedZones()

	// Restore cached enabled countries without hiding a failed restore.
	go func() {
		defer close(m.restoreDone)
		m.mu.RLock()
		var enabled []string
		for code, c := range m.countries {
			if c.Enabled {
				enabled = append(enabled, code)
			}
		}
		m.mu.RUnlock()
		var failures []error
		for _, code := range enabled {
			zone := filepath.Join(GeoIPZonesDir, strings.ToLower(code)+".zone")
			if _, err := os.Stat(zone); os.IsNotExist(err) {
				if _, err = m.SyncCountryZone(code); err != nil {
					failures = append(failures, fmt.Errorf("%s: %w", code, err))
				}
			} else if err := m.ApplyCountryIPSet(code); err != nil {
				failures = append(failures, fmt.Errorf("%s: %w", code, err))
			}
		}
		if err := errors.Join(failures...); err != nil {
			m.mu.Lock()
			m.lastError = err.Error()
			m.mu.Unlock()
			fmt.Fprintf(os.Stderr, "GeoIP restore: %v\n", err)
		}
	}()

	// 5. Start background auto-update scheduler
	m.startScheduler()

	return m
}

func (m *GeoIPManager) loadConfig() {
	m.config = GeoIPConfig{
		AutoUpdate:       true,
		IntervalHours:    24,
		EnabledCountries: []string{"GR"},
	}

	if data, err := os.ReadFile(GeoIPConfigFile); err == nil {
		if err := json.Unmarshal(data, &m.config); err != nil {
			m.lastError = "Invalid GeoIP configuration: " + err.Error()
			m.config.AutoUpdate = false
			m.config.EnabledCountries = nil
		}
	}

	if m.config.IntervalHours <= 0 {
		m.config.IntervalHours = 24
	}

	// Mark enabled countries
	for _, code := range m.config.EnabledCountries {
		if c, ok := m.countries[strings.ToUpper(code)]; ok {
			c.Enabled = true
		}
	}
}

func (m *GeoIPManager) saveConfig() error {
	if err := os.MkdirAll(GeoIPDir, 0755); err != nil {
		return err
	}

	m.config.EnabledCountries = make([]string, 0)
	for code, c := range m.countries {
		if c.Enabled {
			m.config.EnabledCountries = append(m.config.EnabledCountries, code)
		}
	}
	sort.Strings(m.config.EnabledCountries)

	data, err := json.MarshalIndent(m.config, "", "  ")
	if err != nil {
		return err
	}
	return writeAtomic(GeoIPConfigFile, data)
}

func (m *GeoIPManager) scanCachedZones() {
	for code, c := range m.countries {
		zoneFile := filepath.Join(GeoIPZonesDir, strings.ToLower(code)+".zone")
		if stat, err := os.Stat(zoneFile); err == nil && stat.Size() > 0 {
			c.LastUpdated = stat.ModTime()
			if count, err := countLines(zoneFile); err == nil {
				c.SubnetCount = count
			}
		}
	}
}

func countLines(path string) (int, error) {
	f, err := os.Open(path)
	if err != nil {
		return 0, err
	}
	defer f.Close()

	count := 0
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && !strings.HasPrefix(line, "#") {
			count++
		}
	}
	return count, scanner.Err()
}

func (m *GeoIPManager) startScheduler() {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.timer != nil {
		m.timer.Stop()
	}

	if !m.config.AutoUpdate {
		return
	}

	interval := time.Duration(m.config.IntervalHours) * time.Hour
	now := time.Now()

	// Calculate delay until next run
	var delay time.Duration
	if m.config.LastUpdated.IsZero() {
		delay = 1 * time.Minute
	} else {
		next := m.config.LastUpdated.Add(interval)
		if next.After(now) {
			delay = next.Sub(now)
		} else {
			delay = 10 * time.Second
		}
	}

	m.config.NextUpdate = now.Add(delay)

	m.timer = time.AfterFunc(delay, func() {
		_ = m.UpdateNow()
		// Re-arm for next full interval
		m.startScheduler()
	})
}

// ListCountries returns all countries sorted by Name with Greece & Cyprus first
func (m *GeoIPManager) ListCountries() []GeoCountryInfo {
	m.mu.RLock()
	defer m.mu.RUnlock()

	res := make([]GeoCountryInfo, 0, len(m.countries))
	for _, c := range m.countries {
		res = append(res, *c)
	}

	sort.Slice(res, func(i, j int) bool {
		// Put enabled countries at the top, then alphabetically
		if res[i].Enabled != res[j].Enabled {
			return res[i].Enabled
		}
		if res[i].Code == "GR" {
			return true
		}
		if res[j].Code == "GR" {
			return false
		}
		return res[i].Name < res[j].Name
	})

	return res
}

func (m *GeoIPManager) GetStatus() map[string]any {
	m.mu.RLock()
	defer m.mu.RUnlock()

	enabledCount := 0
	totalSubnets := 0
	for _, c := range m.countries {
		if c.Enabled {
			enabledCount++
			totalSubnets += c.SubnetCount
		}
	}

	restoring := false
	if m.restoreDone != nil {
		select {
		case <-m.restoreDone:
		default:
			restoring = true
		}
	}

	return map[string]any{
		"auto_update":       m.config.AutoUpdate,
		"interval_hours":    m.config.IntervalHours,
		"last_updated":      m.config.LastUpdated,
		"next_update":       m.config.NextUpdate,
		"enabled_countries": m.config.EnabledCountries,
		"enabled_count":     enabledCount,
		"total_subnets":     totalSubnets,
		"is_updating":       m.isUpdating,
		"is_restoring":      restoring,
		"last_error":        m.lastError,
	}
}

func (m *GeoIPManager) SetConfig(auto bool, hours int) error {
	m.mu.Lock()
	old := m.config
	if hours < 1 || hours > 8760 {
		m.mu.Unlock()
		return fmt.Errorf("interval must be 1–8760 hours")
	}
	m.config.AutoUpdate = auto
	m.config.IntervalHours = hours
	if err := m.saveConfig(); err != nil {
		m.config = old
		m.mu.Unlock()
		return err
	}
	m.mu.Unlock()
	m.startScheduler()
	return nil
}
func (m *GeoIPManager) ToggleCountry(code string, enable bool) error {
	code = strings.ToUpper(code)
	m.mu.RLock()
	_, ok := m.countries[code]
	m.mu.RUnlock()
	if !ok {
		return fmt.Errorf("unknown country %q", code)
	}
	if !enable && m.referenced != nil && m.referenced(code) {
		return fmt.Errorf("country %s is used by an active firewall rule; change that rule first", code)
	}
	if enable {
		if _, err := os.Stat(filepath.Join(GeoIPZonesDir, strings.ToLower(code)+".zone")); os.IsNotExist(err) {
			if _, err = m.SyncCountryZone(code); err != nil {
				return err
			}
		}
		if err := m.ApplyCountryIPSet(code); err != nil {
			return err
		}
	} else {
		name := GetGeoIPSetName(code)
		if output, err := exec.Command("ipset", "list", name, "-terse").CombinedOutput(); err == nil {
			if err := runIPSet("destroy", name); err != nil {
				return fmt.Errorf("GeoIP set is referenced; it was preserved: %w", err)
			}
		} else if !strings.Contains(string(output), "does not exist") {
			return fmt.Errorf("cannot inspect GeoIP set: %s: %w", output, err)
		}
	}
	m.mu.Lock()
	c := m.countries[code]
	old := c.Enabled
	oldConfig := m.config
	c.Enabled = enable
	if err := m.saveConfig(); err != nil {
		c.Enabled = old
		m.config = oldConfig
		m.mu.Unlock()
		if !enable && old {
			if restoreErr := m.ApplyCountryIPSet(code); restoreErr != nil {
				return fmt.Errorf("persist: %v; restore set: %w", err, restoreErr)
			}
		}
		return err
	}
	m.mu.Unlock()
	return nil
}

func (m *GeoIPManager) BatchToggle(preset string, enable bool, customCodes []string) error {
	var targetCodes []string

	switch strings.ToUpper(preset) {
	case "EU":
		targetCodes = []string{"AT", "BE", "BG", "HR", "CY", "CZ", "DK", "EE", "FI", "FR", "DE", "GR", "HU", "IE", "IT", "LV", "LT", "LU", "MT", "NL", "PL", "PT", "RO", "SK", "SI", "ES", "SE"}
	case "BALKANS":
		targetCodes = []string{"GR", "CY", "AL", "BG", "RO", "RS", "MK", "ME", "BA", "HR", "SI", "TR"}
	case "BLOCKLIST", "HIGHRISK":
		targetCodes = []string{"CN", "RU", "IR", "KP", "BY"}
	case "ALL":
		m.mu.RLock()
		for code := range m.countries {
			targetCodes = append(targetCodes, code)
		}
		m.mu.RUnlock()
	case "NONE":
		m.mu.RLock()
		for code, c := range m.countries {
			if c.Enabled {
				targetCodes = append(targetCodes, code)
			}
		}
		m.mu.RUnlock()
		enable = false
	default:
		targetCodes = customCodes
	}

	if len(targetCodes) == 0 && strings.ToUpper(preset) != "NONE" {
		return fmt.Errorf("no countries selected or unknown preset")
	}
	for _, code := range targetCodes {
		m.mu.RLock()
		_, ok := m.countries[strings.ToUpper(code)]
		m.mu.RUnlock()
		if !ok {
			return fmt.Errorf("unknown country %q", code)
		}
		if !enable && m.referenced != nil && m.referenced(strings.ToUpper(code)) {
			return fmt.Errorf("country %s is used by an active firewall rule", code)
		}
	}
	var failures []error
	for _, code := range targetCodes {
		if err := m.ToggleCountry(code, enable); err != nil {
			failures = append(failures, fmt.Errorf("%s: %w", code, err))
		}
	}
	return errors.Join(failures...)
}

// SyncCountryZone downloads the country zone file from ipdeny and refreshes ipset
func (m *GeoIPManager) SyncCountryZone(code string) (int, error) {
	m.syncMu.Lock()
	defer m.syncMu.Unlock()
	m.mu.RLock()
	_, valid := m.countries[strings.ToUpper(code)]
	m.mu.RUnlock()
	if !valid {
		return 0, fmt.Errorf("unknown country %q", code)
	}
	code = strings.ToUpper(code)
	url := fmt.Sprintf("https://www.ipdeny.com/ipblocks/data/countries/%s.zone", strings.ToLower(code))

	client := &http.Client{Timeout: 15 * time.Second}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return 0, err
	}
	req.Header.Set("User-Agent", "Proxmox-Router/1.0")

	resp, err := client.Do(req)
	if err != nil {
		return 0, fmt.Errorf("failed to fetch zone for %s: %w", code, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return 0, fmt.Errorf("ipdeny returned HTTP %d for %s", resp.StatusCode, code)
	}

	zoneFile := filepath.Join(GeoIPZonesDir, strings.ToLower(code)+".zone")
	out, err := os.CreateTemp(GeoIPZonesDir, ".zone-*")
	tmpFile := ""
	if out != nil {
		tmpFile = out.Name()
	}
	defer os.Remove(tmpFile)
	if err != nil {
		return 0, err
	}

	if _, err := io.Copy(out, io.LimitReader(resp.Body, 32*1024*1024+1)); err != nil {
		out.Close()
		_ = os.Remove(tmpFile)
		return 0, err
	}
	if err := out.Close(); err != nil {
		return 0, err
	}
	data, err := os.ReadFile(tmpFile)
	if err != nil {
		return 0, err
	}
	if len(data) > 32*1024*1024 {
		return 0, fmt.Errorf("GeoIP download too large")
	}
	countValid := 0
	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		ip, _, e := net.ParseCIDR(line)
		if e != nil || ip.To4() == nil {
			return 0, fmt.Errorf("invalid IPv4 CIDR in country download")
		}
		countValid++
	}
	if countValid == 0 {
		return 0, fmt.Errorf("empty country download")
	}
	if err := os.Rename(tmpFile, zoneFile); err != nil {
		return 0, err
	}

	count, _ := countLines(zoneFile)

	m.mu.Lock()
	if c, ok := m.countries[code]; ok {
		c.SubnetCount = count
		c.LastUpdated = time.Now()
	}
	m.mu.Unlock()

	// If enabled, load into kernel ipset immediately
	m.mu.RLock()
	enabled := m.countries[code] != nil && m.countries[code].Enabled
	m.mu.RUnlock()

	if enabled {
		if err := m.ApplyCountryIPSet(code); err != nil {
			return 0, err
		}
	}

	return count, nil
}

// ApplyCountryIPSet uses high-speed `ipset restore` to load the country's subnets
func (m *GeoIPManager) ApplyCountryIPSet(code string) error {
	m.mu.RLock()
	_, valid := m.countries[strings.ToUpper(code)]
	m.mu.RUnlock()
	if !valid {
		return fmt.Errorf("unknown country %q", code)
	}
	code = strings.ToUpper(code)
	zoneFile := filepath.Join(GeoIPZonesDir, strings.ToLower(code)+".zone")

	f, err := os.Open(zoneFile)
	if err != nil {
		return err
	}
	defer f.Close()

	setName := "mk_geo_" + code

	var addresses []string
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && !strings.HasPrefix(line, "#") {
			addresses = append(addresses, line)
		}
	}
	if err := scanner.Err(); err != nil {
		return err
	}
	return replaceSet(setName, "inet", addresses)

}

// ApplyAllEnabledSets ensures all enabled countries are present in kernel ipset
func (m *GeoIPManager) ApplyAllEnabledSets() error {
	m.mu.RLock()
	var enabled []string
	for code, c := range m.countries {
		if c.Enabled {
			enabled = append(enabled, code)
		}
	}
	m.mu.RUnlock()

	for _, code := range enabled {
		if err := m.ApplyCountryIPSet(code); err != nil {
			return err
		}
	}
	return nil
}

// UpdateNow updates zone files for all currently enabled countries
func (m *GeoIPManager) UpdateNow() error {
	m.mu.Lock()
	if m.isUpdating {
		m.mu.Unlock()
		return fmt.Errorf("update already in progress")
	}
	m.isUpdating = true
	m.lastError = ""
	var enabled []string
	for code, c := range m.countries {
		if c.Enabled {
			enabled = append(enabled, code)
		}
	}
	m.mu.Unlock()
	go func() {
		var failures []error
		for _, code := range enabled {
			if _, err := m.SyncCountryZone(code); err != nil {
				failures = append(failures, fmt.Errorf("%s: %w", code, err))
			}
		}
		m.mu.Lock()
		defer m.mu.Unlock()
		m.isUpdating = false
		m.config.NextUpdate = time.Now().Add(time.Duration(m.config.IntervalHours) * time.Hour)
		if len(failures) == 0 {
			m.config.LastUpdated = time.Now()
		}
		if err := m.saveConfig(); err != nil {
			failures = append(failures, err)
		}
		if err := errors.Join(failures...); err != nil {
			m.lastError = err.Error()
		}
	}()
	return nil
}

// GetCountrySubnets returns first N subnets from cached zone file for preview
func (m *GeoIPManager) GetCountrySubnets(code string, limit int) []string {
	m.mu.RLock()
	_, valid := m.countries[strings.ToUpper(code)]
	m.mu.RUnlock()
	if !valid {
		return nil
	}
	code = strings.ToUpper(code)
	zoneFile := filepath.Join(GeoIPZonesDir, strings.ToLower(code)+".zone")

	f, err := os.Open(zoneFile)
	if err != nil {
		return nil
	}
	defer f.Close()

	if limit <= 0 {
		limit = 100
	}

	var res []string
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && !strings.HasPrefix(line, "#") {
			res = append(res, line)
			if len(res) >= limit {
				break
			}
		}
	}
	return res
}

// EnsureCountryOrPreset guarantees that the ipset mk_geo_<CODE> exists and has subnets loaded
func (m *GeoIPManager) EnsureCountryOrPreset(code string) error {
	code = strings.TrimPrefix(code, "!")
	code = strings.TrimPrefix(code, "geo:")
	code = strings.TrimPrefix(code, "geo_")
	code = strings.ToUpper(strings.TrimSpace(code))
	if code == "" {
		return nil
	}

	// Check if it is a preset
	var targetCodes []string
	switch code {
	case "EU":
		targetCodes = []string{"AT", "BE", "BG", "HR", "CY", "CZ", "DK", "EE", "FI", "FR", "DE", "GR", "HU", "IE", "IT", "LV", "LT", "LU", "MT", "NL", "PL", "PT", "RO", "SK", "SI", "ES", "SE"}
	case "BALKANS":
		targetCodes = []string{"GR", "CY", "AL", "BG", "RO", "RS", "MK", "ME", "BA", "HR", "SI", "TR"}
	case "BLOCKLIST", "HIGHRISK":
		targetCodes = []string{"CN", "RU", "IR", "KP", "BY"}
	}

	if len(targetCodes) > 0 {
		setName := "mk_geo_" + code
		var addresses []string
		for _, country := range targetCodes {
			zoneFile := filepath.Join(GeoIPZonesDir, strings.ToLower(country)+".zone")
			if _, err := os.Stat(zoneFile); os.IsNotExist(err) {
				if _, err = m.SyncCountryZone(country); err != nil {
					return err
				}
			}
			data, err := os.ReadFile(zoneFile)
			if err != nil {
				return err
			}
			for _, line := range strings.Split(string(data), "\n") {
				line = strings.TrimSpace(line)
				if line != "" && !strings.HasPrefix(line, "#") {
					addresses = append(addresses, line)
				}
			}
		}
		return replaceSet(setName, "inet", addresses)
	}

	// Single country
	zoneFile := filepath.Join(GeoIPZonesDir, strings.ToLower(code)+".zone")
	if _, err := os.Stat(zoneFile); os.IsNotExist(err) {
		if _, err := m.SyncCountryZone(code); err != nil {
			return err
		}
	}

	m.mu.Lock()
	if c, ok := m.countries[code]; ok {
		c.Enabled = true
		if err := m.saveConfig(); err != nil {
			m.mu.Unlock()
			return err
		}
	}
	m.mu.Unlock()

	return m.ApplyCountryIPSet(code)
}