DebTik/internal/routing/manager.go

package routing

import (
	"bufio"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/vishvananda/netlink"
)

const (
	StaticRoutesStorePath = "/etc/proxmox-router/routing/static_routes.json"
	PolicyRulesStorePath  = "/etc/proxmox-router/routing/policy_rules.json"
	CustomTablesStorePath = "/etc/proxmox-router/routing/custom_tables.json"
	RTTableFilePath       = "/etc/iproute2/rt_tables"
)

type RoutingManager struct {
	mu           sync.Mutex
	staticRoutes []RouteItem
	policyRules  []PolicyRuleItem
	customTables []RoutingTableItem
}

func NewRoutingManager() *RoutingManager {
	m := &RoutingManager{
		staticRoutes: make([]RouteItem, 0),
		policyRules:  make([]PolicyRuleItem, 0),
		customTables: make([]RoutingTableItem, 0),
	}
	m.initRTTableFile()
	m.load()
	m.syncCustomTablesToSystem()
	m.applyAll()
	return m
}

func (m *RoutingManager) initRTTableFile() {
	_ = os.MkdirAll(filepath.Dir(RTTableFilePath), 0755)
	if _, err := os.Stat(RTTableFilePath); os.IsNotExist(err) {
		// Copy from /usr/share/iproute2/rt_tables or seed default
		content, err := os.ReadFile("/usr/share/iproute2/rt_tables")
		if err == nil {
			_ = os.WriteFile(RTTableFilePath, content, 0644)
		} else {
			defaultContent := "# reserved values\n255\tlocal\n254\tmain\n253\tdefault\n0\tunspec\n"
			_ = os.WriteFile(RTTableFilePath, []byte(defaultContent), 0644)
		}
	}
}

func (m *RoutingManager) load() {
	if data, err := os.ReadFile(StaticRoutesStorePath); err == nil {
		_ = json.Unmarshal(data, &m.staticRoutes)
	}
	if data, err := os.ReadFile(PolicyRulesStorePath); err == nil {
		_ = json.Unmarshal(data, &m.policyRules)
	}
	if data, err := os.ReadFile(CustomTablesStorePath); err == nil {
		_ = json.Unmarshal(data, &m.customTables)
	}
}

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

	if d, err := json.MarshalIndent(m.staticRoutes, "", "  "); err == nil {
		_ = os.WriteFile(StaticRoutesStorePath, d, 0600)
	}
	if d, err := json.MarshalIndent(m.policyRules, "", "  "); err == nil {
		_ = os.WriteFile(PolicyRulesStorePath, d, 0600)
	}
	if d, err := json.MarshalIndent(m.customTables, "", "  "); err == nil {
		_ = os.WriteFile(CustomTablesStorePath, d, 0600)
	}
	return nil
}

func (m *RoutingManager) syncCustomTablesToSystem() {
	content, _ := os.ReadFile(RTTableFilePath)
	str := string(content)

	var sb strings.Builder
	sb.WriteString(str)

	for _, t := range m.customTables {
		check := fmt.Sprintf("\t%s", t.Name)
		if !strings.Contains(str, check) {
			sb.WriteString(fmt.Sprintf("%d\t%s\n", t.TableID, t.Name))
		}
	}

	_ = os.WriteFile(RTTableFilePath, []byte(sb.String()), 0644)
}

func (m *RoutingManager) applyAll() {
	// Apply all static routes
	for _, r := range m.staticRoutes {
		if !r.Disabled {
			_ = m.applyRouteToKernel(r)
		}
	}

	// Apply policy rules
	for _, pr := range m.policyRules {
		_ = m.applyRuleToKernel(pr)
	}
}

// ==================== Routes (/ip route) ====================

func isDynamicProtocol(proto netlink.RouteProtocol) bool {
	switch int(proto) {
	case 11, // RTPROT_ZEBRA
		16,  // RTPROT_DHCP
		42,  // RTPROT_BABEL
		186, // RTPROT_BGP
		187, // RTPROT_ISIS
		188, // RTPROT_OSPF
		189, // RTPROT_RIP
		192: // RTPROT_EIGRP
		return true
	default:
		return false
	}
}

func (m *RoutingManager) listLiveRoutesLocked() []RouteItem {
	tableNames := m.getTableMap()
	links, _ := netlink.LinkList()
	linkMap := make(map[int]string)
	for _, l := range links {
		linkMap[l.Attrs().Index] = l.Attrs().Name
	}

	var liveRoutes []RouteItem
	routes, err := netlink.RouteList(nil, netlink.FAMILY_V4)
	if err == nil {
		for _, r := range routes {
			dst := "0.0.0.0/0"
			if r.Dst != nil {
				dst = r.Dst.String()
			}
			gw := ""
			if r.Gw != nil {
				gw = r.Gw.String()
			}

			tableName := "main"
			if name, ok := tableNames[r.Table]; ok {
				tableName = name
			} else if r.Table != 254 && r.Table != 0 {
				tableName = strconv.Itoa(r.Table)
			}

			ifName := linkMap[r.LinkIndex]
			isDirect := r.Scope == netlink.SCOPE_LINK || r.Protocol == netlink.RouteProtocol(2)
			isDyn := isDynamicProtocol(r.Protocol)
			isStatic := !isDirect && !isDyn

			cleanDst := strings.ReplaceAll(dst, "/", "_")
			cleanGw := gw
			if cleanGw == "" {
				cleanGw = ifName
			}
			routeID := fmt.Sprintf("kr-%d-%s-%s", r.Table, cleanDst, cleanGw)

			liveRoutes = append(liveRoutes, RouteItem{
				ID:           routeID,
				DstAddress:   dst,
				Gateway:      gw,
				Interface:    ifName,
				Distance:     r.Priority,
				RoutingTable: tableName,
				TableID:      r.Table,
				Active:       true,
				Connect:      isDirect,
				Dynamic:      isDyn,
				Static:       isStatic,
				Disabled:     false,
			})
		}
	}
	return liveRoutes
}

func (m *RoutingManager) ListRoutes() []RouteItem {
	m.mu.Lock()
	defer m.mu.Unlock()

	liveRoutes := m.listLiveRoutesLocked()

	// Merge static routes that might be disabled or have custom metadata
	for _, sr := range m.staticRoutes {
		found := false
		for i, lr := range liveRoutes {
			if lr.DstAddress == sr.DstAddress && lr.RoutingTable == sr.RoutingTable && (sr.Gateway == "" || lr.Gateway == sr.Gateway) {
				liveRoutes[i].Comment = sr.Comment
				liveRoutes[i].Static = true
				liveRoutes[i].Dynamic = false
				liveRoutes[i].ID = sr.ID
				if sr.Distance > 0 && liveRoutes[i].Distance == 0 {
					liveRoutes[i].Distance = sr.Distance
				}
				found = true
				break
			}
		}
		if !found {
			// Inactive or disabled static route
			srCopy := sr
			srCopy.Active = false
			srCopy.Static = true
			srCopy.Dynamic = false
			liveRoutes = append(liveRoutes, srCopy)
		}
	}

	sort.Slice(liveRoutes, func(i, j int) bool {
		if liveRoutes[i].RoutingTable != liveRoutes[j].RoutingTable {
			return liveRoutes[i].RoutingTable < liveRoutes[j].RoutingTable
		}
		if liveRoutes[i].DstAddress != liveRoutes[j].DstAddress {
			return liveRoutes[i].DstAddress < liveRoutes[j].DstAddress
		}
		return liveRoutes[i].Gateway < liveRoutes[j].Gateway
	})

	return liveRoutes
}

func (m *RoutingManager) SaveRoute(r RouteItem) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if r.DstAddress == "" {
		return fmt.Errorf("destination address is required")
	}
	if r.DstAddress == "default" {
		r.DstAddress = "0.0.0.0/0"
	}
	if r.RoutingTable == "" {
		r.RoutingTable = "main"
	}
	r.Static = true

	// 1. If old route details are provided from edit modal, remove old route first
	if r.OldDstAddress != "" {
		oldDist := 0
		if r.OldDistance != nil {
			oldDist = *r.OldDistance
		}
		oldRoute := RouteItem{
			DstAddress:   r.OldDstAddress,
			Gateway:      r.OldGateway,
			RoutingTable: r.OldRoutingTable,
			Distance:     oldDist,
			Interface:    r.OldInterface,
		}
		_ = m.removeRouteFromKernel(oldRoute)
	}

	// 2. If existing entry in m.staticRoutes exists, remove its old route from kernel
	foundIndex := -1
	for i, existing := range m.staticRoutes {
		if existing.ID == r.ID || (r.ID == "" && existing.DstAddress == r.DstAddress && existing.RoutingTable == r.RoutingTable) {
			foundIndex = i
			_ = m.removeRouteFromKernel(existing)
			break
		}
	}

	// 3. If r.ID starts with kr- (kernel static route being edited/adopted), remove it from kernel
	if strings.HasPrefix(r.ID, "kr-") {
		for _, lr := range m.listLiveRoutesLocked() {
			if lr.ID == r.ID {
				_ = m.removeRouteFromKernel(lr)
				break
			}
		}
		r.ID = fmt.Sprintf("route-%d", time.Now().UnixNano())
	}

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

	if foundIndex >= 0 {
		m.staticRoutes[foundIndex] = r
	} else {
		m.staticRoutes = append(m.staticRoutes, r)
	}

	_ = m.save()

	if r.Disabled {
		_ = m.removeRouteFromKernel(r)
	} else {
		if err := m.applyRouteToKernel(r); err != nil {
			return err
		}
	}

	return nil
}

func (m *RoutingManager) DeleteRoute(id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	var updated []RouteItem
	var target *RouteItem
	for _, r := range m.staticRoutes {
		if r.ID == id {
			target = &r
		} else {
			updated = append(updated, r)
		}
	}

	if target != nil {
		_ = m.removeRouteFromKernel(*target)
		m.staticRoutes = updated
		return m.save()
	}

	// If not in staticRoutes, could be a live route (e.g. kr-...)
	for _, lr := range m.listLiveRoutesLocked() {
		if lr.ID == id {
			_ = m.removeRouteFromKernel(lr)
			return nil
		}
	}

	return fmt.Errorf("route not found")
}

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

	for i := range m.staticRoutes {
		if m.staticRoutes[i].ID == id {
			m.staticRoutes[i].Disabled = disabled
			_ = m.save()
			if disabled {
				_ = m.removeRouteFromKernel(m.staticRoutes[i])
			} else {
				_ = m.applyRouteToKernel(m.staticRoutes[i])
			}
			return nil
		}
	}

	// If not in staticRoutes, could be a kernel static route to disable
	if disabled {
		for _, lr := range m.listLiveRoutesLocked() {
			if lr.ID == id {
				lr.Disabled = true
				lr.Static = true
				lr.Active = false
				lr.ID = fmt.Sprintf("route-%d", time.Now().UnixNano())
				_ = m.removeRouteFromKernel(lr)
				m.staticRoutes = append(m.staticRoutes, lr)
				return m.save()
			}
		}
	}

	return fmt.Errorf("route not found")
}

func (m *RoutingManager) applyRouteToKernel(r RouteItem) error {
	dst := strings.TrimSpace(r.DstAddress)
	if dst == "default" || dst == "0.0.0.0/0" {
		dst = "default"
	} else if !strings.Contains(dst, "/") && net.ParseIP(dst) != nil {
		dst = dst + "/32"
	}

	args := []string{"route", "replace", dst}

	gw := strings.TrimSpace(r.Gateway)
	if gw != "" {
		if net.ParseIP(gw) != nil {
			args = append(args, "via", gw)
		} else {
			args = append(args, "dev", gw)
		}
	}

	iface := strings.TrimSpace(r.Interface)
	if iface != "" && gw != iface && (gw == "" || net.ParseIP(gw) != nil) {
		args = append(args, "dev", iface)
	}

	if r.Distance > 0 {
		args = append(args, "metric", strconv.Itoa(r.Distance))
	}

	if r.RoutingTable != "" && r.RoutingTable != "main" {
		args = append(args, "table", r.RoutingTable)
	}

	args = append(args, "proto", "static")

	cmd := exec.Command("ip", args...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to apply route (%s): %s: %w", strings.Join(args, " "), string(out), err)
	}
	return nil
}

func (m *RoutingManager) removeRouteFromKernel(r RouteItem) error {
	dst := strings.TrimSpace(r.DstAddress)
	if dst == "default" || dst == "0.0.0.0/0" {
		dst = "default"
	} else if !strings.Contains(dst, "/") && net.ParseIP(dst) != nil {
		dst = dst + "/32"
	}

	args := []string{"route", "del", dst}

	gw := strings.TrimSpace(r.Gateway)
	if gw != "" {
		if net.ParseIP(gw) != nil {
			args = append(args, "via", gw)
		} else {
			args = append(args, "dev", gw)
		}
	}

	if r.Distance > 0 {
		args = append(args, "metric", strconv.Itoa(r.Distance))
	}

	if r.RoutingTable != "" && r.RoutingTable != "main" {
		args = append(args, "table", r.RoutingTable)
	}

	_ = exec.Command("ip", args...).Run()

	// Fallback deletion without metric
	fallbackArgs := []string{"route", "del", dst}
	if r.RoutingTable != "" && r.RoutingTable != "main" {
		fallbackArgs = append(fallbackArgs, "table", r.RoutingTable)
	}
	_ = exec.Command("ip", fallbackArgs...).Run()

	return nil
}

// ==================== Policy Rules (/ip route rule) ====================

func (m *RoutingManager) ListPolicyRules() []PolicyRuleItem {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Query live rules from kernel
	rules, err := netlink.RuleList(netlink.FAMILY_V4)
	if err != nil {
		return m.policyRules
	}

	tableNames := m.getTableMap()
	var res []PolicyRuleItem

	userRulesByPrio := make(map[int]PolicyRuleItem)
	for _, pr := range m.policyRules {
		userRulesByPrio[pr.Priority] = pr
	}

	for _, r := range rules {
		tableName := "main"
		if name, ok := tableNames[r.Table]; ok {
			tableName = name
		} else if r.Table != 254 && r.Table != 0 {
			tableName = strconv.Itoa(r.Table)
		}

		src := ""
		if r.Src != nil {
			src = r.Src.String()
		}
		dst := ""
		if r.Dst != nil {
			dst = r.Dst.String()
		}

		markStr := ""
		if r.Mark > 0 {
			markStr = fmt.Sprintf("0x%x (%d)", r.Mark, r.Mark)
		}

		id := fmt.Sprintf("rule-%d-%d", r.Priority, r.Table)
		comment := ""
		isDefault := false
		isSystem := false

		if r.Priority == 0 || r.Priority == 32766 || r.Priority == 32767 {
			isDefault = true
			isSystem = true
		} else if userRule, ok := userRulesByPrio[r.Priority]; ok {
			if userRule.ID != "" {
				id = userRule.ID
			}
			comment = userRule.Comment
			isDefault = false
			isSystem = false
		} else {
			// System/daemon rule not added by user in proxmox-router (e.g. WireGuard, IPsec)
			isSystem = true
		}

		res = append(res, PolicyRuleItem{
			ID:          id,
			Priority:    r.Priority,
			SrcAddress:  src,
			DstAddress:  dst,
			RoutingMark: markStr,
			Table:       tableName,
			TableID:     r.Table,
			Action:      "lookup",
			Comment:     comment,
			IsDefault:   isDefault,
			IsSystem:    isSystem,
		})
	}

	return res
}

func (m *RoutingManager) SavePolicyRule(pr PolicyRuleItem) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if pr.Table == "" {
		pr.Table = "main"
	}
	if pr.ID == "" {
		pr.ID = fmt.Sprintf("prule-%d", time.Now().UnixNano())
	}
	if pr.Priority == 0 {
		pr.Priority = 1000 + len(m.policyRules)*10
	}

	if err := m.applyRuleToKernel(pr); err != nil {
		return err
	}

	m.policyRules = append(m.policyRules, pr)
	return m.save()
}

func (m *RoutingManager) DeletePolicyRule(id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	var updated []PolicyRuleItem
	var target *PolicyRuleItem
	for _, pr := range m.policyRules {
		if pr.ID == id || fmt.Sprintf("rule-%d-%d", pr.Priority, pr.TableID) == id {
			target = &pr
		} else {
			updated = append(updated, pr)
		}
	}

	if target != nil {
		_ = m.removeRuleFromKernel(*target)
	} else if strings.HasPrefix(id, "rule-") {
		parts := strings.Split(id, "-")
		if len(parts) >= 2 {
			if prio, err := strconv.Atoi(parts[1]); err == nil && prio > 0 {
				_ = exec.Command("ip", "rule", "del", "pref", strconv.Itoa(prio)).Run()
			}
		}
	}
	m.policyRules = updated
	return m.save()
}

func (m *RoutingManager) applyRuleToKernel(pr PolicyRuleItem) error {
	args := []string{"rule", "add"}
	if pr.Priority > 0 {
		args = append(args, "pref", strconv.Itoa(pr.Priority))
	}
	if pr.SrcAddress != "" {
		args = append(args, "from", pr.SrcAddress)
	}
	if pr.DstAddress != "" {
		args = append(args, "to", pr.DstAddress)
	}
	if pr.RoutingMark != "" {
		// Clean mark string if it has (hex)
		markClean := pr.RoutingMark
		if strings.Contains(markClean, "(") {
			markClean = strings.TrimSpace(markClean[:strings.Index(markClean, "(")])
		}
		args = append(args, "fwmark", markClean)
	}
	if pr.Table != "" {
		args = append(args, "lookup", pr.Table)
	}

	cmd := exec.Command("ip", args...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to add rule (%s): %s: %w", strings.Join(args, " "), string(out), err)
	}
	return nil
}

func (m *RoutingManager) removeRuleFromKernel(pr PolicyRuleItem) error {
	if pr.Priority > 0 {
		_ = exec.Command("ip", "rule", "del", "pref", strconv.Itoa(pr.Priority)).Run()
		return nil
	}
	args := []string{"rule", "del"}
	if pr.RoutingMark != "" {
		markClean := pr.RoutingMark
		if strings.Contains(markClean, "(") {
			markClean = strings.TrimSpace(markClean[:strings.Index(markClean, "(")])
		}
		args = append(args, "fwmark", markClean)
	}
	if pr.Table != "" {
		args = append(args, "lookup", pr.Table)
	}

	_ = exec.Command("ip", args...).Run()
	return nil
}

// ==================== Routing Tables ====================

func (m *RoutingManager) ListTables() []RoutingTableItem {
	m.mu.Lock()
	defer m.mu.Unlock()

	tables := []RoutingTableItem{
		{Name: "main", TableID: 254, Comment: "Default system routing table"},
		{Name: "default", TableID: 253, Comment: "Default fallback table"},
		{Name: "local", TableID: 255, Comment: "Local kernel device table"},
	}

	tables = append(tables, m.customTables...)
	return tables
}

func (m *RoutingManager) SaveTable(t RoutingTableItem) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if t.Name == "" {
		return fmt.Errorf("table name is required")
	}
	if t.TableID == 0 {
		// Auto assign FIB ID starting at 100
		t.TableID = 100 + len(m.customTables)
	}

	for _, existing := range m.customTables {
		if existing.Name == t.Name || existing.TableID == t.TableID {
			return fmt.Errorf("table with name %q or ID %d already exists", t.Name, t.TableID)
		}
	}

	m.customTables = append(m.customTables, t)
	_ = m.save()
	m.syncCustomTablesToSystem()
	return nil
}

func (m *RoutingManager) getTableMap() map[int]string {
	tableMap := map[int]string{
		254: "main",
		253: "default",
		255: "local",
	}

	// Read /etc/iproute2/rt_tables
	if f, err := os.Open(RTTableFilePath); err == nil {
		defer f.Close()
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line == "" || strings.HasPrefix(line, "#") {
				continue
			}
			parts := strings.Fields(line)
			if len(parts) >= 2 {
				if id, err := strconv.Atoi(parts[0]); err == nil {
					tableMap[id] = parts[1]
				}
			}
		}
	}

	for _, t := range m.customTables {
		tableMap[t.TableID] = t.Name
	}

	return tableMap
}