DebTik/internal/firewall/geoip.go

package firewall

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"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{}
	isUpdating bool
}

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

	m := &GeoIPManager{
		countries: make(map[string]*GeoCountryInfo),
		stopCh:    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()

	// 4. Pre-download / sync initial enabled countries if missing
	go func() {
		m.mu.RLock()
		enabled := make([]string, len(m.config.EnabledCountries))
		copy(enabled, m.config.EnabledCountries)
		m.mu.RUnlock()

		for _, code := range enabled {
			zonePath := filepath.Join(GeoIPZonesDir, strings.ToLower(code)+".zone")
			if _, err := os.Stat(zonePath); os.IsNotExist(err) {
				_, _ = m.SyncCountryZone(code)
			} else {
				if err := m.ApplyCountryIPSet(code); err != nil {
					fmt.Fprintf(os.Stderr, "GeoIP restore failed for %s: %v\n", code, 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 {
		_ = json.Unmarshal(data, &m.config)
	}

	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 {
	_ = os.MkdirAll(GeoIPDir, 0755)

	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 os.WriteFile(GeoIPConfigFile, data, 0644)
}

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

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

func (m *GeoIPManager) SetConfig(autoUpdate bool, intervalHours int) error {
	m.mu.Lock()
	if intervalHours <= 0 {
		intervalHours = 24
	}
	m.config.AutoUpdate = autoUpdate
	m.config.IntervalHours = intervalHours
	_ = m.saveConfig()
	m.mu.Unlock()

	m.startScheduler()
	return nil
}

func (m *GeoIPManager) ToggleCountry(code string, enable bool) error {
	code = strings.ToUpper(code)

	m.mu.Lock()
	c, ok := m.countries[code]
	if !ok {
		m.mu.Unlock()
		return fmt.Errorf("country %q not found", code)
	}
	c.Enabled = enable
	_ = m.saveConfig()
	m.mu.Unlock()

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

	// Disable -> destroy ipset
	setName := "mk_geo_" + code
	_ = exec.Command("ipset", "flush", setName).Run()
	_ = exec.Command("ipset", "destroy", setName).Run()
	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
	}

	for _, code := range targetCodes {
		_ = m.ToggleCountry(code, enable)
	}

	return nil
}

// SyncCountryZone downloads the country zone file from ipdeny and refreshes ipset
func (m *GeoIPManager) SyncCountryZone(code string) (int, error) {
	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")
	tmpFile := zoneFile + ".tmp"

	out, err := os.Create(tmpFile)
	if err != nil {
		return 0, err
	}

	if _, err := io.Copy(out, resp.Body); err != nil {
		out.Close()
		_ = os.Remove(tmpFile)
		return 0, err
	}
	out.Close()

	_ = os.Rename(tmpFile, zoneFile)

	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 {
	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.mu.Unlock()

	go func() {
		defer func() {
			m.mu.Lock()
			m.isUpdating = false
			m.config.LastUpdated = time.Now()
			m.config.NextUpdate = time.Now().Add(time.Duration(m.config.IntervalHours) * time.Hour)
			_ = m.saveConfig()
			m.mu.Unlock()
		}()

		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 {
			_, _ = m.SyncCountryZone(code)
		}
	}()

	return nil
}

// GetCountrySubnets returns first N subnets from cached zone file for preview
func (m *GeoIPManager) GetCountrySubnets(code string, limit int) []string {
	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
		_ = m.saveConfig()
	}
	m.mu.Unlock()

	return m.ApplyCountryIPSet(code)
}