DebTik/internal/firewall/iptables.go

package firewall

import (
	"bufio"
	"bytes"
	"fmt"
	"hash/fnv"
	"os/exec"
	"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
	}

	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]
				counters[ruleID] = [2]uint64{pkts, bts}
			}
		}
	}

	return counters, nil
}

// 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
		}
		ruleStr := buildMangleRuleLine(r)
		if ruleStr != "" {
			sb.WriteString(ruleStr + "\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
		}
		ruleStr := buildNATRuleLine(r)
		if ruleStr != "" {
			sb.WriteString(ruleStr + "\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
		}
		ruleStr := buildFilterRuleLine(r)
		if ruleStr != "" {
			sb.WriteString(ruleStr + "\n")
		}
	}
	sb.WriteString("COMMIT\n")

	return sb.String()
}

func buildFilterRuleLine(r FilterRule) string {
	chain := strings.ToUpper(r.Chain)
	if chain == "" {
		chain = "FORWARD"
	}

	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, r.InInterface, r.OutInterface, 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
		return ""
	default:
		args = append(args, "-j ACCEPT")
	}

	return strings.Join(args, " ")
}

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

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

	args = appendMatchCriteria(args, r.SrcAddress, r.DstAddress, r.Protocol, r.SrcPort, r.DstPort, r.InInterface, r.OutInterface, 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 = fmt.Sprintf("%s:%s", 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 = fmt.Sprintf("%s:%s", 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")
	}

	return strings.Join(args, " ")
}

func buildMangleRuleLine(r MangleRule) string {
	chain := strings.ToUpper(r.Chain)
	if chain == "" {
		chain = "PREROUTING"
	}

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

	args = appendMatchCriteria(args, r.SrcAddress, r.DstAddress, r.Protocol, r.SrcPort, r.DstPort, r.InInterface, r.OutInterface, 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":
		return ""
	default:
		return ""
	}

	return strings.Join(args, " ")
}

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 {
	cmd := exec.Command("iptables-restore", "-c")
	cmd.Stdin = strings.NewReader(payload)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to apply iptables: %s: %w", string(out), err)
	}
	return nil
}