DebTik/internal/firewall/iptables.go

package firewall

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"hash/fnv"
	"net"
	"os"
	"os/exec"
	"sort"
	"strconv"
	"strings"
)

// MarkNameToID converts a string mark name into a deterministic 32-bit fwmark integer
func MarkNameToID(name string) int {
	if name == "" {
		return 0
	}
	if id, err := strconv.Atoi(name); err == nil && id > 0 {
		return id
	}
	h := fnv.New32a()
	h.Write([]byte(name))
	val := int(h.Sum32() % 65000)
	if val < 100 {
		val += 100
	}
	return val
}

// ReadRuleCounters reads `iptables-save -c` and populates the Packets and Bytes fields for all rules
func ReadRuleCounters(filterRules []FilterRule, natRules []NATRule, mangleRules []MangleRule) (map[string][2]uint64, error) {
	out, err := exec.Command("iptables-save", "-c").Output()
	if err != nil {
		return nil, err
	}

	if v6, err := exec.Command("ip6tables-save", "-c").Output(); err == nil {
		out = append(out, '\n')
		out = append(out, v6...)
	}
	counters := make(map[string][2]uint64)
	scanner := bufio.NewScanner(bytes.NewReader(out))

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if !strings.HasPrefix(line, "[") {
			continue
		}

		closeBracket := strings.Index(line, "]")
		if closeBracket == -1 {
			continue
		}

		countsStr := line[1:closeBracket]
		parts := strings.Split(countsStr, ":")
		if len(parts) != 2 {
			continue
		}

		pkts, _ := strconv.ParseUint(parts[0], 10, 64)
		bts, _ := strconv.ParseUint(parts[1], 10, 64)

		commentIdx := strings.Index(line, "--comment \"id:")
		if commentIdx != -1 {
			sub := line[commentIdx+len("--comment \"id:"):]
			endQuote := strings.Index(sub, "\"")
			if endQuote != -1 {
				ruleID := sub[:endQuote]
				existing := counters[ruleID]
				counters[ruleID] = [2]uint64{existing[0] + pkts, existing[1] + bts}
			}
		}
	}

	return counters, nil
}

// resolveInterfaceList resolves an interface list or role name (e.g. "WAN", "LAN") to interface names.
func resolveInterfaceList(listName string) []string {
	listName = strings.TrimSpace(listName)
	if listName == "" {
		return nil
	}
	data, err := os.ReadFile("/etc/proxmox-router/metadata.json")
	if err != nil {
		return nil
	}
	var meta struct {
		Interfaces map[string]struct {
			Role string `json:"role"`
		} `json:"interfaces"`
	}
	if err := json.Unmarshal(data, &meta); err != nil {
		return nil
	}
	var matched []string
	for ifaceName, info := range meta.Interfaces {
		if strings.EqualFold(info.Role, listName) {
			matched = append(matched, ifaceName)
		}
	}
	sort.Strings(matched)
	return matched
}

func getInterfaceCombinations(inIface, inList, outIface, outList string) [][2]string {
	var inListResolved []string
	if inIface != "" {
		inListResolved = []string{inIface}
	} else if inList != "" {
		inListResolved = resolveInterfaceList(inList)
	} else {
		inListResolved = []string{""}
	}

	var outListResolved []string
	if outIface != "" {
		outListResolved = []string{outIface}
	} else if outList != "" {
		outListResolved = resolveInterfaceList(outList)
	} else {
		outListResolved = []string{""}
	}

	if inList != "" && inIface == "" && len(inListResolved) == 0 {
		return nil
	}
	if outList != "" && outIface == "" && len(outListResolved) == 0 {
		return nil
	}

	var combos [][2]string
	for _, in := range inListResolved {
		for _, out := range outListResolved {
			combos = append(combos, [2]string{in, out})
		}
	}
	return combos
}

// BuildIPTablesRestore generates the atomic configuration for `iptables-restore`
func BuildIPTablesRestore(filterRules []FilterRule, natRules []NATRule, mangleRules []MangleRule) string {
	var sb strings.Builder

	// ================= MANGLE TABLE =================
	sb.WriteString("*mangle\n")
	sb.WriteString(":PREROUTING ACCEPT [0:0]\n")
	sb.WriteString(":INPUT ACCEPT [0:0]\n")
	sb.WriteString(":FORWARD ACCEPT [0:0]\n")
	sb.WriteString(":OUTPUT ACCEPT [0:0]\n")
	sb.WriteString(":POSTROUTING ACCEPT [0:0]\n")

	// Custom user chains in mangle
	userMangleChains := make(map[string]bool)
	for _, r := range mangleRules {
		ch := strings.ToUpper(r.Chain)
		if ch != "PREROUTING" && ch != "POSTROUTING" && ch != "INPUT" && ch != "OUTPUT" && ch != "FORWARD" && ch != "" {
			if !userMangleChains[ch] {
				sb.WriteString(fmt.Sprintf(":%s - [0:0]\n", ch))
				userMangleChains[ch] = true
			}
		}
		if r.Action == "jump" && r.JumpTarget != "" {
			jt := strings.ToUpper(r.JumpTarget)
			if jt != "ACCEPT" && jt != "DROP" && jt != "RETURN" && !userMangleChains[jt] {
				sb.WriteString(fmt.Sprintf(":%s - [0:0]\n", jt))
				userMangleChains[jt] = true
			}
		}
	}

	for _, r := range mangleRules {
		if r.Disabled {
			continue
		}
		for _, line := range buildMangleRuleLines(r) {
			if line != "" {
				sb.WriteString(line + "\n")
			}
		}
	}
	sb.WriteString("COMMIT\n")

	// ================= NAT TABLE =================
	sb.WriteString("*nat\n")
	sb.WriteString(":PREROUTING ACCEPT [0:0]\n")
	sb.WriteString(":INPUT ACCEPT [0:0]\n")
	sb.WriteString(":OUTPUT ACCEPT [0:0]\n")
	sb.WriteString(":POSTROUTING ACCEPT [0:0]\n")

	userNATChains := make(map[string]bool)
	for _, r := range natRules {
		ch := strings.ToUpper(r.Chain)
		if ch != "PREROUTING" && ch != "POSTROUTING" && ch != "OUTPUT" && ch != "INPUT" && ch != "SRCNAT" && ch != "DSTNAT" && ch != "" {
			if !userNATChains[ch] {
				sb.WriteString(fmt.Sprintf(":%s - [0:0]\n", ch))
				userNATChains[ch] = true
			}
		}
		if r.Action == "jump" && r.JumpTarget != "" {
			jt := strings.ToUpper(r.JumpTarget)
			if jt != "ACCEPT" && jt != "RETURN" && !userNATChains[jt] {
				sb.WriteString(fmt.Sprintf(":%s - [0:0]\n", jt))
				userNATChains[jt] = true
			}
		}
	}

	for _, r := range natRules {
		if r.Disabled {
			continue
		}
		for _, line := range buildNATRuleLines(r) {
			if line != "" {
				sb.WriteString(line + "\n")
			}
		}
	}
	sb.WriteString("COMMIT\n")

	// ================= FILTER TABLE =================
	sb.WriteString("*filter\n")
	sb.WriteString(":INPUT ACCEPT [0:0]\n")
	sb.WriteString(":FORWARD ACCEPT [0:0]\n")
	sb.WriteString(":OUTPUT ACCEPT [0:0]\n")

	userFilterChains := make(map[string]bool)
	for _, r := range filterRules {
		ch := strings.ToUpper(r.Chain)
		if ch != "INPUT" && ch != "FORWARD" && ch != "OUTPUT" && ch != "" {
			if !userFilterChains[ch] {
				sb.WriteString(fmt.Sprintf(":%s - [0:0]\n", ch))
				userFilterChains[ch] = true
			}
		}
		if r.Action == "jump" && r.JumpTarget != "" {
			jt := strings.ToUpper(r.JumpTarget)
			if jt != "ACCEPT" && jt != "DROP" && jt != "REJECT" && jt != "RETURN" && !userFilterChains[jt] {
				sb.WriteString(fmt.Sprintf(":%s - [0:0]\n", jt))
				userFilterChains[jt] = true
			}
		}
	}

	for _, r := range filterRules {
		if r.Disabled {
			continue
		}
		for _, line := range buildFilterRuleLines(r) {
			if line != "" {
				sb.WriteString(line + "\n")
			}
		}
	}
	sb.WriteString("COMMIT\n")

	return sb.String()
}

func buildFilterRuleLines(r FilterRule) []string {
	combos := getInterfaceCombinations(r.InInterface, r.InInterfaceList, r.OutInterface, r.OutInterfaceList)
	if len(combos) == 0 {
		return nil
	}

	chain := strings.ToUpper(r.Chain)
	if chain == "" {
		chain = "FORWARD"
	}

	var lines []string
	for _, combo := range combos {
		inIface, outIface := combo[0], combo[1]

		var args []string
		args = append(args, fmt.Sprintf("-A %s", chain))

		// Match criteria
		args = appendMatchCriteria(args, r.SrcAddress, r.DstAddress, r.Protocol, r.SrcPort, r.DstPort, inIface, outIface, r.ConnectionState, r.ConnectionMark, r.RoutingMark, r.SrcAddressList, r.DstAddressList, r.SrcGeoIP, r.DstGeoIP)

		// Comment ID for counters mapping & user comment
		args = append(args, fmt.Sprintf("-m comment --comment \"id:%s\"", r.ID))

		// Action
		act := strings.ToLower(r.Action)
		switch act {
		case "accept":
			args = append(args, "-j ACCEPT")
		case "drop":
			args = append(args, "-j DROP")
		case "reject":
			args = append(args, "-j REJECT")
			if r.RejectWith != "" {
				args = append(args, fmt.Sprintf("--reject-with %s", r.RejectWith))
			}
		case "return":
			args = append(args, "-j RETURN")
		case "jump":
			if r.JumpTarget != "" {
				args = append(args, fmt.Sprintf("-j %s", strings.ToUpper(r.JumpTarget)))
			}
		case "log":
			args = append(args, "-j LOG")
			if r.LogPrefix != "" {
				args = append(args, fmt.Sprintf("--log-prefix \"%s: \"", r.LogPrefix))
			}
		case "add-src-to-address-list":
			if r.AddressList != "" {
				args = append(args, fmt.Sprintf("-j SET --add-set %s src", GetIPSetName(r.AddressList)))
			}
		case "add-dst-to-address-list":
			if r.AddressList != "" {
				args = append(args, fmt.Sprintf("-j SET --add-set %s dst", GetIPSetName(r.AddressList)))
			}
		case "passthrough":
			// in iptables, passthrough continues evaluation
			continue
		default:
			args = append(args, "-j ACCEPT")
		}

		lines = append(lines, strings.Join(args, " "))
	}

	return lines
}

func buildNATRuleLines(r NATRule) []string {
	combos := getInterfaceCombinations(r.InInterface, r.InInterfaceList, r.OutInterface, r.OutInterfaceList)
	if len(combos) == 0 {
		return nil
	}

	chain := strings.ToUpper(r.Chain)
	if chain == "SRCNAT" {
		chain = "POSTROUTING"
	} else if chain == "DSTNAT" {
		chain = "PREROUTING"
	}

	var lines []string
	for _, combo := range combos {
		inIface, outIface := combo[0], combo[1]

		var args []string
		args = append(args, fmt.Sprintf("-A %s", chain))

		args = appendMatchCriteria(args, r.SrcAddress, r.DstAddress, r.Protocol, r.SrcPort, r.DstPort, inIface, outIface, nil, "", "", r.SrcAddressList, r.DstAddressList, r.SrcGeoIP, r.DstGeoIP)
		args = append(args, fmt.Sprintf("-m comment --comment \"id:%s\"", r.ID))

		act := strings.ToLower(r.Action)
		switch act {
		case "masquerade":
			args = append(args, "-j MASQUERADE")
		case "src-nat":
			if r.ToAddresses != "" {
				target := r.ToAddresses
				if r.ToPorts != "" {
					target = net.JoinHostPort(r.ToAddresses, r.ToPorts)
				}
				args = append(args, fmt.Sprintf("-j SNAT --to-source %s", target))
			} else {
				args = append(args, "-j MASQUERADE")
			}
		case "dst-nat":
			if r.ToAddresses != "" {
				target := r.ToAddresses
				if r.ToPorts != "" {
					target = net.JoinHostPort(r.ToAddresses, r.ToPorts)
				}
				args = append(args, fmt.Sprintf("-j DNAT --to-destination %s", target))
			}
		case "redirect":
			if r.ToPorts != "" {
				args = append(args, fmt.Sprintf("-j REDIRECT --to-ports %s", r.ToPorts))
			} else {
				args = append(args, "-j REDIRECT")
			}
		case "jump":
			if r.JumpTarget != "" {
				args = append(args, fmt.Sprintf("-j %s", strings.ToUpper(r.JumpTarget)))
			}
		case "log":
			args = append(args, "-j LOG")
			if r.LogPrefix != "" {
				args = append(args, fmt.Sprintf("--log-prefix \"%s: \"", r.LogPrefix))
			}
		case "add-src-to-address-list":
			if r.AddressList != "" {
				args = append(args, fmt.Sprintf("-j SET --add-set %s src", GetIPSetName(r.AddressList)))
			}
		case "add-dst-to-address-list":
			if r.AddressList != "" {
				args = append(args, fmt.Sprintf("-j SET --add-set %s dst", GetIPSetName(r.AddressList)))
			}
		case "accept":
			args = append(args, "-j ACCEPT")
		case "return":
			args = append(args, "-j RETURN")
		default:
			args = append(args, "-j MASQUERADE")
		}

		lines = append(lines, strings.Join(args, " "))
	}

	return lines
}

func buildMangleRuleLines(r MangleRule) []string {
	combos := getInterfaceCombinations(r.InInterface, r.InInterfaceList, r.OutInterface, r.OutInterfaceList)
	if len(combos) == 0 {
		return nil
	}

	chain := strings.ToUpper(r.Chain)
	if chain == "" {
		chain = "PREROUTING"
	}

	var lines []string
	for _, combo := range combos {
		inIface, outIface := combo[0], combo[1]

		var args []string
		args = append(args, fmt.Sprintf("-A %s", chain))

		args = appendMatchCriteria(args, r.SrcAddress, r.DstAddress, r.Protocol, r.SrcPort, r.DstPort, inIface, outIface, r.ConnectionState, r.ConnectionMark, r.RoutingMark, r.SrcAddressList, r.DstAddressList, r.SrcGeoIP, r.DstGeoIP)
		args = append(args, fmt.Sprintf("-m comment --comment \"id:%s\"", r.ID))

		act := strings.ToLower(r.Action)
		switch act {
		case "mark-connection":
			if r.NewConnectionMark != "" {
				markID := MarkNameToID(r.NewConnectionMark)
				args = append(args, fmt.Sprintf("-j CONNMARK --set-mark %d", markID))
			}
		case "mark-routing":
			if r.NewRoutingMark != "" {
				markID := MarkNameToID(r.NewRoutingMark)
				args = append(args, fmt.Sprintf("-j MARK --set-mark %d", markID))
			}
		case "mark-packet":
			if r.NewPacketMark != "" {
				markID := MarkNameToID(r.NewPacketMark)
				args = append(args, fmt.Sprintf("-j MARK --set-mark %d", markID))
			}
		case "change-mss":
			args = append(args, "-p tcp --tcp-flags SYN,RST SYN -j TCPMSS")
			if r.NewMSS == "clamp-to-pmtu" || r.NewMSS == "" {
				args = append(args, "--clamp-mss-to-pmtu")
			} else {
				args = append(args, fmt.Sprintf("--set-mss %s", r.NewMSS))
			}
		case "set-dscp", "change-dscp":
			if r.NewDSCP > 0 {
				args = append(args, fmt.Sprintf("-j DSCP --set-dscp %d", r.NewDSCP))
			}
		case "change-ttl":
			if r.NewTTL != "" {
				if strings.HasPrefix(r.NewTTL, "+") {
					args = append(args, fmt.Sprintf("-j TTL --ttl-inc %s", strings.TrimPrefix(r.NewTTL, "+")))
				} else if strings.HasPrefix(r.NewTTL, "-") {
					args = append(args, fmt.Sprintf("-j TTL --ttl-dec %s", strings.TrimPrefix(r.NewTTL, "-")))
				} else {
					args = append(args, fmt.Sprintf("-j TTL --ttl-set %s", r.NewTTL))
				}
			}
		case "set-priority":
			if r.NewPriority != "" && chain == "POSTROUTING" {
				args = append(args, fmt.Sprintf("-j CLASSIFY --set-class 0:%s", r.NewPriority))
			}
		case "drop":
			args = append(args, "-j DROP")
		case "jump":
			if r.JumpTarget != "" {
				args = append(args, fmt.Sprintf("-j %s", strings.ToUpper(r.JumpTarget)))
			}
		case "log":
			args = append(args, "-j LOG")
			if r.LogPrefix != "" {
				args = append(args, fmt.Sprintf("--log-prefix \"%s: \"", r.LogPrefix))
			}
		case "add-src-to-address-list":
			if r.AddressList != "" {
				args = append(args, fmt.Sprintf("-j SET --add-set %s src", GetIPSetName(r.AddressList)))
			}
		case "add-dst-to-address-list":
			if r.AddressList != "" {
				args = append(args, fmt.Sprintf("-j SET --add-set %s dst", GetIPSetName(r.AddressList)))
			}
		case "accept":
			args = append(args, "-j ACCEPT")
		case "return":
			args = append(args, "-j RETURN")
		case "clear-df", "strip-ipv4-options", "route", "sniff-pc", "sniff-tzsp", "passthrough":
			continue
		default:
			continue
		}

		lines = append(lines, strings.Join(args, " "))
	}

	return lines
}

func appendMatchCriteria(args []string, srcAddr, dstAddr, proto, srcPort, dstPort, inIface, outIface string, ctStates []string, connMark, routingMark, srcList, dstList, srcGeoIP, dstGeoIP string) []string {
	if srcAddr != "" {
		if strings.HasPrefix(srcAddr, "!") {
			args = append(args, fmt.Sprintf("! -s %s", strings.TrimPrefix(srcAddr, "!")))
		} else {
			args = append(args, fmt.Sprintf("-s %s", srcAddr))
		}
	}

	if dstAddr != "" {
		if strings.HasPrefix(dstAddr, "!") {
			args = append(args, fmt.Sprintf("! -d %s", strings.TrimPrefix(dstAddr, "!")))
		} else {
			args = append(args, fmt.Sprintf("-d %s", dstAddr))
		}
	}

	effectiveProto := strings.ToLower(proto)
	if effectiveProto != "" && effectiveProto != "all" {
		args = append(args, fmt.Sprintf("-p %s", effectiveProto))
	} else if srcPort != "" || dstPort != "" {
		effectiveProto = "tcp"
		args = append(args, "-p tcp")
	}

	if srcPort != "" && (effectiveProto == "tcp" || effectiveProto == "udp") {
		if strings.Contains(srcPort, ",") {
			args = append(args, fmt.Sprintf("-m multiport --sports %s", srcPort))
		} else {
			args = append(args, fmt.Sprintf("--sport %s", srcPort))
		}
	}

	if dstPort != "" && (effectiveProto == "tcp" || effectiveProto == "udp") {
		if strings.Contains(dstPort, ",") {
			args = append(args, fmt.Sprintf("-m multiport --dports %s", dstPort))
		} else {
			args = append(args, fmt.Sprintf("--dport %s", dstPort))
		}
	}

	if inIface != "" {
		if strings.HasPrefix(inIface, "!") {
			args = append(args, fmt.Sprintf("! -i %s", strings.TrimPrefix(inIface, "!")))
		} else {
			args = append(args, fmt.Sprintf("-i %s", inIface))
		}
	}

	if outIface != "" {
		if strings.HasPrefix(outIface, "!") {
			args = append(args, fmt.Sprintf("! -o %s", strings.TrimPrefix(outIface, "!")))
		} else {
			args = append(args, fmt.Sprintf("-o %s", outIface))
		}
	}

	if len(ctStates) > 0 {
		var upperStates []string
		for _, s := range ctStates {
			upperStates = append(upperStates, strings.ToUpper(s))
		}
		args = append(args, fmt.Sprintf("-m conntrack --ctstate %s", strings.Join(upperStates, ",")))
	}

	if connMark != "" {
		markID := MarkNameToID(connMark)
		args = append(args, fmt.Sprintf("-m connmark --mark %d", markID))
	}

	if routingMark != "" {
		markID := MarkNameToID(routingMark)
		args = append(args, fmt.Sprintf("-m mark --mark %d", markID))
	}

	if srcList != "" {
		invert := ""
		cleanList := srcList
		if strings.HasPrefix(cleanList, "!") {
			invert = "! "
			cleanList = strings.TrimPrefix(cleanList, "!")
		}
		setName := GetIPSetName(cleanList)
		args = append(args, fmt.Sprintf("-m set %s--match-set %s src", invert, setName))
	}

	if dstList != "" {
		invert := ""
		cleanList := dstList
		if strings.HasPrefix(cleanList, "!") {
			invert = "! "
			cleanList = strings.TrimPrefix(cleanList, "!")
		}
		setName := GetIPSetName(cleanList)
		args = append(args, fmt.Sprintf("-m set %s--match-set %s dst", invert, setName))
	}

	if srcGeoIP != "" {
		invert := ""
		cleanGeo := srcGeoIP
		if strings.HasPrefix(cleanGeo, "!") {
			invert = "! "
			cleanGeo = strings.TrimPrefix(cleanGeo, "!")
		}
		setName := GetGeoIPSetName(cleanGeo)
		args = append(args, fmt.Sprintf("-m set %s--match-set %s src", invert, setName))
	}

	if dstGeoIP != "" {
		invert := ""
		cleanGeo := dstGeoIP
		if strings.HasPrefix(cleanGeo, "!") {
			invert = "! "
			cleanGeo = strings.TrimPrefix(cleanGeo, "!")
		}
		setName := GetGeoIPSetName(cleanGeo)
		args = append(args, fmt.Sprintf("-m set %s--match-set %s dst", invert, setName))
	}

	return args
}

// ApplyIPTables executes `iptables-restore` with the given rule set payload
func ApplyIPTables(payload string) error {
	live, err := exec.Command("iptables-save").Output()
	if err != nil {
		return fmt.Errorf("cannot snapshot existing firewall: %w", err)
	}
	payload = ownedPayload(payload, string(live))
	for _, test := range []bool{true, false} {
		args := []string{"--noflush", "-c"}
		if test {
			args = append(args, "--test")
		}
		cmd := exec.Command("iptables-restore", args...)
		cmd.Stdin = strings.NewReader(payload)
		if out, err := cmd.CombinedOutput(); err != nil {
			return fmt.Errorf("iptables apply (test=%v) failed: %s: %w", test, out, err)
		}
	}
	return nil
}