package firewall

import (
	"strings"
	"testing"
)

func TestChainPlacement(t *testing.T) {
	cases := []struct {
		kind  string
		rule  any
		valid bool
	}{
		{"nat", NATRule{Chain: "dstnat", Action: "dst-nat"}, true},
		{"nat", NATRule{Chain: "prerouting", Action: "dst-nat"}, true},
		{"nat", NATRule{Chain: "output", Action: "dst-nat"}, true},
		{"nat", NATRule{Chain: "custom", Action: "dst-nat"}, true},
		{"nat", NATRule{Chain: "srcnat", Action: "dst-nat"}, false},
		{"nat", NATRule{Chain: "forward", Action: "accept"}, false},
		{"nat", NATRule{Chain: "srcnat", Action: "masquerade", OutInterface: "eth0"}, true},
		{"nat", NATRule{Chain: "input", Action: "masquerade"}, false},
		{"nat", NATRule{Chain: "prerouting", Action: "accept", OutInterface: "eth0"}, false},
		{"filter", FilterRule{Chain: "prerouting", Action: "accept"}, false},
		{"filter", FilterRule{Chain: "input", Action: "accept", OutInterface: "eth0"}, false},
		{"mangle", MangleRule{Chain: "output", Action: "accept", InInterface: "eth0"}, false},
		{"mangle", MangleRule{Chain: "forward", Action: "accept", InInterface: "eth0", OutInterface: "eth1"}, true},
	}
	for _, c := range cases {
		err := validateChainPlacement(c.kind, c.rule)
		if (err == nil) != c.valid {
			t.Errorf("%s %#v: %v", c.kind, c.rule, err)
		}
	}
}

func TestNATChainHooks(t *testing.T) {
	rules := []NATRule{{ID: "incoming", Chain: "dstnat", Action: "dst-nat", ToAddresses: "192.0.2.1"}, {ID: "local", Chain: "output", Action: "dst-nat", ToAddresses: "192.0.2.2"}, {ID: "outgoing", Chain: "srcnat", Action: "masquerade"}}
	payload := ownedPayload(BuildIPTablesRestore(nil, rules, nil), "")
	for _, want := range []string{"-A DT_nat_PREROUTING ", "-A DT_nat_OUTPUT ", "-A DT_nat_POSTROUTING ", "-I PREROUTING 1 -m comment --comment debtik:hook -j DT_nat_PREROUTING", "-I OUTPUT 1 -m comment --comment debtik:hook -j DT_nat_OUTPUT", "-I POSTROUTING 1 -m comment --comment debtik:hook -j DT_nat_POSTROUTING"} {
		if !strings.Contains(payload, want) {
			t.Errorf("missing %s", want)
		}
	}
}

func TestNativeTargetsAndPortValidation(t *testing.T) {
	for _, action := range []string{"DNAT", "dnat", "dst-nat"} {
		r := NATRule{Chain: "prerouting", Action: action, Protocol: "udp", DstPort: "53", ToAddresses: "192.0.2.53"}
		if err := ValidateRule("nat", r); err != nil {
			t.Fatalf("%s: %v", action, err)
		}
		normal := normalizeNativeRule("nat", r).(NATRule)
		if !strings.Contains(BuildIPTablesRestore(nil, []NATRule{normal}, nil), "-j DNAT --to-destination 192.0.2.53") {
			t.Fatalf("%s did not produce DNAT", action)
		}
	}
	for target, want := range map[string]string{"MARK": "mark-packet", "CONNMARK": "mark-connection", "TCPMSS": "change-mss", "DSCP": "change-dscp", "TTL": "change-ttl", "HL": "change-ttl", "CLASSIFY": "set-priority"} {
		got := normalizeNativeRule("mangle", MangleRule{Action: target}).(MangleRule).Action
		if got != want {
			t.Errorf("target %s normalized to %s, want %s", target, got, want)
		}
	}
	for _, rule := range []any{FilterRule{Protocol: "all", DstPort: "53"}, MangleRule{Protocol: "icmp", SrcPort: "53"}, NATRule{Protocol: "!udp", DstPort: "53"}} {
		if validateNewPortProtocol(rule) == nil {
			t.Fatalf("accepted unsafe port combination: %#v", rule)
		}
	}
	if validateNewPortProtocol(FilterRule{Protocol: "udp", DstPort: "53"}) != nil {
		t.Fatal("valid UDP port was rejected")
	}
}

func TestNeighboringChainRule(t *testing.T) {
	chains := []string{"srcnat", "input", "postrouting", "dstnat", "prerouting"}
	for _, tt := range []struct {
		index int
		dir   string
		nat   bool
		want  int
	}{{0, "down", true, 2}, {2, "up", true, 0}, {3, "down", true, 4}, {1, "up", true, -1}, {0, "down", false, -1}} {
		got, err := neighboringChainRule(chains, tt.index, tt.dir, tt.nat)
		if err != nil || got != tt.want {
			t.Errorf("%+v: got %d %v", tt, got, err)
		}
	}
	if _, err := neighboringChainRule(chains, 0, "sideways", true); err == nil {
		t.Fatal("invalid direction accepted")
	}
}