DebTik/internal/routing/manager.go

package routing

import (
	"bufio"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"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 (m *RoutingManager) ListRoutes() []RouteItem {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 1. Get live routes from netlink across all tables
	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

			liveRoutes = append(liveRoutes, RouteItem{
				ID:           fmt.Sprintf("kr-%d-%s-%s", r.Table, dst, gw),
				DstAddress:   dst,
				Gateway:      gw,
				Interface:    ifName,
				Distance:     r.Priority,
				RoutingTable: tableName,
				TableID:      r.Table,
				Active:       true,
				Connect:      isDirect,
				Dynamic:      !isDirect && (r.Protocol == netlink.RouteProtocol(2) || r.Protocol == netlink.RouteProtocol(3)),
				Static:       !isDirect && r.Protocol != netlink.RouteProtocol(2) && r.Protocol != netlink.RouteProtocol(3),
				Disabled:     false,
			})
		}
	}

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

	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.RoutingTable == "" {
		r.RoutingTable = "main"
	}
	if r.ID == "" {
		r.ID = fmt.Sprintf("route-%d", time.Now().UnixNano())
	}
	r.Static = true

	found := false
	for i, existing := range m.staticRoutes {
		if existing.ID == r.ID {
			m.staticRoutes[i] = r
			found = true
			break
		}
	}
	if !found {
		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()
}

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
		}
	}
	return fmt.Errorf("route not found")
}

func (m *RoutingManager) applyRouteToKernel(r RouteItem) error {
	args := []string{"route", "replace"}

	if r.DstAddress == "default" || r.DstAddress == "0.0.0.0/0" {
		args = append(args, "default")
	} else {
		args = append(args, r.DstAddress)
	}

	if r.Gateway != "" {
		// check if gateway is IP or device
		if net.ParseIP(r.Gateway) != nil {
			args = append(args, "via", r.Gateway)
		} else {
			args = append(args, "dev", r.Gateway)
		}
	}

	if r.Interface != "" && r.Gateway != r.Interface {
		args = append(args, "dev", r.Interface)
	}

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

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

	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 {
	args := []string{"route", "del"}
	if r.DstAddress == "default" || r.DstAddress == "0.0.0.0/0" {
		args = append(args, "default")
	} else {
		args = append(args, r.DstAddress)
	}

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

	_ = exec.Command("ip", args...).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

	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)

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

	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 {
			target = &pr
		} else {
			updated = append(updated, pr)
		}
	}

	if target != nil {
		_ = m.removeRuleFromKernel(*target)
	}
	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 {
	args := []string{"rule", "del"}
	if pr.Priority > 0 {
		args = append(args, "pref", strconv.Itoa(pr.Priority))
	}
	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
}