package firewall

import (
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestRuleNeedsHook(t *testing.T) {
	// Features with no native CSF/APF config path route through the hook.
	needs := []*Rule{
		{Proto: TCP, Port: 22, Action: Accept, State: StateNew},
		{Proto: TCP, Port: 80, Action: Accept, Log: true},
		{Proto: TCP, Port: 80, Action: Accept, InInterface: "eth0"},
		{Direction: DirOutput, Proto: TCP, Port: 80, Action: Accept, OutInterface: "eth0"},
		{Proto: TCP, Port: 25, Action: Drop, RateLimit: &RateLimit{Rate: 1, Unit: PerSecond}},
		{Proto: ICMPv6, ICMPType: Ptr[uint8](128), Action: Accept},
		// A forward rule has no native CSF/APF config path, so it routes through the
		// raw-iptables hook (which emits an -A FORWARD rule).
		{Direction: DirForward, Proto: TCP, Port: 8080, Action: Accept},
	}
	for _, r := range needs {
		require.True(t, ruleNeedsHook(r), "expected %+v to need the hook", *r)
	}
	// Natively expressible rules do not.
	native := []*Rule{
		{Proto: TCP, Port: 22, Action: Accept},
		{Proto: ICMP, ICMPType: Ptr[uint8](8), Action: Accept},
		{Proto: TCP, Port: 22, Action: Drop, ConnLimit: &ConnLimit{Count: 5, PerSource: true}},
	}
	for _, r := range native {
		require.False(t, ruleNeedsHook(r), "expected %+v to stay native", *r)
	}
}

// bareProtoNeedsHook routes a portless, addressless non-ICMP match to the hook —
// the shape CSF/APF cannot express natively but iptables applies directly — while
// leaving every rule that carries an address, a port, or an ICMP protocol on its
// own native/ICMP path.
func TestBareProtoNeedsHook(t *testing.T) {
	// Bare protocol matches with no address and no port go to the hook.
	needs := []*Rule{
		{Proto: TCP, Action: Accept},
		{Proto: UDP, Action: Drop},
		{Proto: ProtocolAny, Action: Accept},
		{Proto: TCP, Direction: DirOutput, Action: Accept},
	}
	for _, r := range needs {
		require.True(t, bareProtoNeedsHook(r), "expected %+v to route to the hook", *r)
	}
	// A port, an address, or an ICMP protocol keeps the rule off this path.
	native := []*Rule{
		{Proto: TCP, Port: 22, Action: Accept},
		{Proto: TCP, SourcePort: 1234, Action: Accept},
		{Proto: TCP, Source: "1.2.3.4", Action: Accept},
		{Proto: ProtocolAny, Destination: "1.2.3.4", Action: Accept},
		{Proto: ICMP, ICMPType: Ptr[uint8](8), Action: Accept},
		{Proto: ICMPv6, Action: Accept},
	}
	for _, r := range native {
		require.False(t, bareProtoNeedsHook(r), "expected %+v to stay off the hook route", *r)
	}
}

// Two rules that are Equal (port-set order is not part of rule identity) must
// inject the same command line, so a second add is a no-op and a remove using a
// reordered port set still finds the rule. The hook script matches on the exact
// marshalled line, so the marshaller must render Equal port sets identically.
func TestHookScriptPortOrderIdempotent(t *testing.T) {
	dir := t.TempDir()
	h := &hookScript{
		rulePrefix: "go_firewall",
		hookPath:   filepath.Join(dir, "csfpre.sh"),
		hookPerm:   0700,
	}

	// SCTP has no native CSF/APF config path, so a multi-port SCTP rule routes
	// through the hook. These two differ only in port order, so they are Equal.
	a := &Rule{Family: IPv4, Proto: SCTP, Ports: []PortRange{{Start: 80}, {Start: 443}}, Action: Accept}
	b := &Rule{Family: IPv4, Proto: SCTP, Ports: []PortRange{{Start: 443}, {Start: 80}}, Action: Accept}
	require.True(t, a.Equal(b, true), "the two rules must be Equal (order-independent)")

	changed, err := h.edit(a, false)
	require.NoError(t, err)
	require.True(t, changed)

	changed, err = h.edit(b, false)
	require.NoError(t, err)
	require.False(t, changed, "an Equal rule with reordered ports must not inject a duplicate")

	// Removing via the reordered form must still find and drop the rule.
	changed, err = h.edit(b, true)
	require.NoError(t, err)
	require.True(t, changed, "removing an Equal rule with reordered ports must drop it")

	got, err := h.getRules()
	require.NoError(t, err)
	require.Empty(t, got, "the rule must be gone after removal")
}

// A ProtocolAny rule that carries a port has no single iptables form — a --dport
// match requires a concrete -p tcp/udp — so the hook fans it out into a tcp line and
// a udp line, mirroring the tcp+udp fan-out csf/apf write in their native config.
// Both add and remove must fan out and never reject the rule for want of a concrete
// protocol. Regression: csf's GetRules merges a same-port TCP_IN/UDP_IN pair back
// into one ProtocolAny rule, so a Backup captured it and Restore's hook-copy clear
// then failed to marshal it, breaking the whole restore.
func TestHookScriptProtocolAnyPortFansOut(t *testing.T) {
	dir := t.TempDir()
	h := &hookScript{
		rulePrefix: "go_firewall",
		hookPath:   filepath.Join(dir, "csfpre.sh"),
		hookPerm:   0700,
	}

	// Adding a ProtocolAny port rule injects a tcp line and a udp line.
	any := &Rule{Family: IPv4, Proto: ProtocolAny, Port: 20, Action: Accept}
	changed, err := h.edit(any, false)
	require.NoError(t, err, "a ProtocolAny port rule must marshal, not be rejected")
	require.True(t, changed)

	got, err := h.getRules()
	require.NoError(t, err)
	require.Len(t, got, 2, "a ProtocolAny port rule fans out into a tcp and a udp hook line")
	protos := map[Protocol]bool{}
	for _, g := range got {
		protos[g.Proto] = true
	}
	require.True(t, protos[TCP] && protos[UDP], "the fan-out must cover both tcp and udp: %+v", got)

	// Removing the ProtocolAny form clears both concrete copies in one call, without
	// erroring on the port-without-concrete-protocol shape.
	changed, err = h.edit(any, true)
	require.NoError(t, err, "removing a ProtocolAny port rule must not fail to marshal")
	require.True(t, changed, "the ProtocolAny remove must clear the tcp and udp copies")

	got, err = h.getRules()
	require.NoError(t, err)
	require.Empty(t, got, "both fanned-out copies must be gone after the ProtocolAny remove")
}

// A deny whose action differs from the CSF/APF config's STOP action has no native
// form (deny_hosts/csf.deny encode no action of their own), so those backends
// inject it through the hook, whose iptables rule carries the exact action. The
// hook must marshal and read back the precise action, not coerce it — otherwise a
// Reject deny would read back as Drop and churn on every Sync.
func TestHookScriptCarriesExactDenyAction(t *testing.T) {
	dir := t.TempDir()
	h := &hookScript{
		rulePrefix: "go_firewall",
		hookPath:   filepath.Join(dir, "csfpre.sh"),
		hookPerm:   0700,
	}

	for _, deny := range []*Rule{
		{Family: IPv4, Proto: TCP, Port: 22, Source: "192.0.2.31/32", Action: Reject},
		{Family: IPv4, Proto: TCP, Port: 22, Source: "192.0.2.32/32", Action: Drop},
	} {
		changed, err := h.edit(deny, false)
		require.NoError(t, err)
		require.True(t, changed, "the deny must be injected: %+v", deny)

		got, err := h.getRules()
		require.NoError(t, err)
		var match *Rule
		for _, g := range got {
			if g.Equal(deny, true) {
				match = g
			}
		}
		require.NotNil(t, match, "the deny must read back from the hook: %+v", deny)
		require.Equal(t, deny.Action, match.Action,
			"the hook must carry the deny's exact action, not coerce it: %+v", deny)

		changed, err = h.edit(deny, true)
		require.NoError(t, err)
		require.True(t, changed, "the deny must be removable: %+v", deny)
	}
}

// Hook removal matches on the underlying rule, not the exact command line: a copy
// of a rule a customer added under a different comment (or a differently spelled
// address) must still be removed, since the comment is not part of rule identity.
func TestHookScriptRemoveIgnoresComment(t *testing.T) {
	dir := t.TempDir()
	h := &hookScript{
		rulePrefix: "go_firewall",
		hookPath:   filepath.Join(dir, "csfpre.sh"),
		hookPerm:   0700,
	}

	// Plant a rule the way a customer would: same underlying match, a foreign comment,
	// and an un-normalized address (no /32). A hookScript with a different prefix marks
	// it as not ours.
	foreign := &hookScript{rulePrefix: "acme", hookPath: h.hookPath, hookPerm: 0700}
	planted := &Rule{Family: IPv4, Proto: TCP, Port: 4567, Source: "192.0.2.60", Action: Accept, Comment: "ticket-42"}
	changed, err := foreign.edit(planted, false)
	require.NoError(t, err)
	require.True(t, changed)

	// Remove the same underlying rule with no comment and the normalized address.
	changed, err = h.edit(&Rule{Family: IPv4, Proto: TCP, Port: 4567, Source: "192.0.2.60/32", Action: Accept}, true)
	require.NoError(t, err)
	require.True(t, changed, "a rule with the same match but a different comment must still be removed")

	got, err := h.getRules()
	require.NoError(t, err)
	require.Empty(t, got, "the customer's differently-commented copy must be gone")
}

func TestHookScriptRoundTrip(t *testing.T) {
	dir := t.TempDir()
	h := &hookScript{
		rulePrefix: "go_firewall",
		hookPath:   filepath.Join(dir, "csfpre.sh"),
		hookPerm:   0700,
	}

	// A family-agnostic rule is injected for both v4 and v6.
	lines, err := h.rulesToLines(&Rule{Proto: TCP, Port: 8080, Action: Accept, State: StateNew})
	require.NoError(t, err)
	require.Len(t, lines, 2)
	require.True(t, strings.HasPrefix(lines[0], "iptables "), "want iptables line, got %q", lines[0])
	require.True(t, strings.HasPrefix(lines[1], "ip6tables "), "want ip6tables line, got %q", lines[1])

	// Family-pinned rules covering each non-native feature round-trip through the
	// hook.
	rules := []*Rule{
		{Family: IPv4, Proto: TCP, Port: 22, Action: Accept, State: StateNew | StateEstablished},
		{Family: IPv4, Proto: TCP, Port: 80, Action: Accept, Log: true, LogPrefix: "web"},
		{Family: IPv4, Proto: TCP, Port: 443, Action: Accept, InInterface: "eth0"},
		{Family: IPv6, Proto: ICMPv6, ICMPType: Ptr[uint8](128), Action: Accept},
		{Family: IPv4, Proto: TCP, Port: 25, Action: Drop, RateLimit: &RateLimit{Rate: 5, Unit: PerMinute, Burst: 3}},
	}
	for _, r := range rules {
		changed, err := h.edit(r, false)
		require.NoError(t, err, "add %+v", *r)
		require.True(t, changed, "expected add to change the script: %+v", *r)
	}

	// Adding again is idempotent.
	changed, err := h.edit(rules[0], false)
	require.NoError(t, err)
	require.False(t, changed, "expected a duplicate add to be a no-op")

	// The command lines live in the hook itself, under a single shebang.
	hookData, err := os.ReadFile(h.hookPath)
	require.NoError(t, err)
	require.Equal(t, 1, strings.Count(string(hookData), "#!/bin/sh"), "hook should carry one shebang")
	require.Contains(t, string(hookData), "iptables ")

	// Every rule reads back equal (family ignored, as the hook stores per-family).
	got, err := h.getRules()
	require.NoError(t, err)
	require.Len(t, got, len(rules))
	for _, want := range rules {
		found := false
		for _, g := range got {
			if g.EqualBase(want, true) {
				found = true
				break
			}
		}
		require.True(t, found, "rule not read back: %+v", *want)
	}

	// The logged rule round-trips with its prefix intact.
	for _, g := range got {
		if g.Port == 80 {
			require.True(t, g.Log, "expected the port 80 rule to be logged")
			require.Equal(t, "web", g.LogPrefix)
		}
	}

	// Removing one drops it (both its LOG and action lines) and leaves the rest.
	changed, err = h.edit(rules[1], true)
	require.NoError(t, err)
	require.True(t, changed)
	got, err = h.getRules()
	require.NoError(t, err)
	require.Len(t, got, len(rules)-1)
	for _, g := range got {
		require.False(t, g.EqualBase(rules[1], true), "removed rule still present")
	}

	// Removing an absent rule is a no-op.
	changed, err = h.edit(rules[1], true)
	require.NoError(t, err)
	require.False(t, changed, "expected removing an absent rule to be a no-op")
}

// Writing command lines into the existing hook must leave user-authored content
// untouched: arbitrary shell survives an add and a remove, and an iptables rule a
// user added by hand both survives edits and surfaces in getRules (the library
// reconciles the hook's actual state, not just the lines it wrote).
func TestHookPreservesUserContent(t *testing.T) {
	dir := t.TempDir()
	hookPath := filepath.Join(dir, "csfpre.sh")
	userContent := "#!/bin/sh\n" +
		"# operator's own pre-hook logic\n" +
		"logger firewall reloading\n" +
		"iptables -A INPUT -p tcp --dport 2222 -j ACCEPT\n"
	require.NoError(t, os.WriteFile(hookPath, []byte(userContent), 0700))

	h := &hookScript{rulePrefix: "go_firewall", hookPath: hookPath, hookPerm: 0700}

	// A hand-added iptables rule the library never wrote surfaces in getRules,
	// reported as foreign (no prefix tag).
	got, err := h.getRules()
	require.NoError(t, err)
	require.Len(t, got, 1)
	require.Equal(t, uint16(2222), got[0].Port)
	require.False(t, got[0].HasPrefix, "a user-authored rule must read back as foreign")

	// Adding our rule keeps every user line in place.
	added := &Rule{Family: IPv4, Proto: TCP, Port: 80, Action: Accept, State: StateNew}
	changed, err := h.edit(added, false)
	require.NoError(t, err)
	require.True(t, changed)
	data, err := os.ReadFile(hookPath)
	require.NoError(t, err)
	require.Contains(t, string(data), "logger firewall reloading")
	require.Contains(t, string(data), "iptables -A INPUT -p tcp --dport 2222 -j ACCEPT")
	require.Equal(t, 1, strings.Count(string(data), "#!/bin/sh"), "must not add a second shebang")

	// Removing our rule leaves the user's shell and rule behind.
	changed, err = h.edit(added, true)
	require.NoError(t, err)
	require.True(t, changed)
	data, err = os.ReadFile(hookPath)
	require.NoError(t, err)
	require.Contains(t, string(data), "logger firewall reloading")
	require.Contains(t, string(data), "iptables -A INPUT -p tcp --dport 2222 -j ACCEPT")

	// The user's rule still reads back after our churn.
	got, err = h.getRules()
	require.NoError(t, err)
	require.Len(t, got, 1)
	require.Equal(t, uint16(2222), got[0].Port)
}

// A hook line is sourced by /bin/sh, so a comment or log prefix containing $ or a
// backtick must be single-quoted (a literal), not left in strconv.Quote's double
// quotes where the shell would expand it. And it must still parse back intact.
func TestHookShellSafeLogPrefix(t *testing.T) {
	require.Equal(t, "-A", shellSafeToken("-A"))
	require.Equal(t, "INPUT", shellSafeToken("INPUT"))
	require.Equal(t, `'web $USER'`, shellSafeToken("web $USER"))
	require.Equal(t, `'a'\''b'`, shellSafeToken("a'b"))

	h := &hookScript{rulePrefix: "myapp"}
	lines, err := h.rulesToLines(&Rule{Family: IPv4, Proto: TCP, Port: 22, Action: Drop, Log: true, LogPrefix: "drop $x"})
	require.NoError(t, err)
	joined := strings.Join(lines, "\n")
	require.NotContains(t, joined, `"drop $x"`, "a $-bearing prefix must not stay double-quoted for the shell")
	require.Contains(t, joined, `'drop $x'`)

	found := false
	for _, l := range lines {
		if r, ok := h.parseLine(l); ok && r.Log {
			require.Equal(t, "drop $x", r.LogPrefix)
			found = true
		}
	}
	require.True(t, found, "the log line must parse back to the original prefix")
}

// A protocol CSF/APF cannot express natively (SCTP and the portless IP
// protocols) is routed through the raw-iptables hook and round-trips there.
func TestHookProtocolExtras(t *testing.T) {
	for _, p := range []Protocol{SCTP, GRE, ESP, AH} {
		require.True(t, hookOnlyProto(p), "%s should route through the hook", p)
		require.True(t, ruleNeedsHook(&Rule{Proto: p, Action: Accept}))
	}
	require.False(t, hookOnlyProto(TCP))
	require.False(t, ruleNeedsHook(&Rule{Proto: TCP, Port: 22, Action: Accept}))

	h := &hookScript{hookPath: "/tmp/unused", rulePrefix: "go_firewall"}
	cases := []*Rule{
		{Family: IPv4, Proto: GRE, Action: Accept},
		{Family: IPv4, Proto: SCTP, Port: 9000, Action: Accept},
	}
	for _, orig := range cases {
		lines, err := h.rulesToLines(orig)
		require.NoError(t, err, "%+v", orig)
		require.NotEmpty(t, lines)
		got, ok := h.parseLine(lines[len(lines)-1])
		require.True(t, ok, "line %q", lines[len(lines)-1])
		require.True(t, got.EqualBase(orig, true), "want %+v got %+v", orig, got)
	}
}

// bareHostOneWay classifies a one-way bare-address host rule — a single address,
// no ports, any-protocol, a concrete direction — which csf/apf must route to the
// hook because a plain line is bidirectional and an advanced rule needs a port. A
// DirAny bare host (the bidirectional plain line) and any ported or protocol-pinned
// rule are excluded.
func TestBareHostOneWay(t *testing.T) {
	yes := []*Rule{
		{Direction: DirInput, Source: "1.2.3.4", Action: Accept},
		{Direction: DirOutput, Destination: "1.2.3.4", Action: Accept},
		{Direction: DirInput, Source: "10.0.0.0/8", Action: Drop},
	}
	for _, r := range yes {
		require.Truef(t, bareHostOneWay(r), "expected one-way bare host: %+v", r)
	}
	no := []*Rule{
		{Direction: DirAny, Source: "1.2.3.4", Action: Accept},                           // bidirectional plain line
		{Direction: DirForward, Source: "1.2.3.4", Action: Accept},                       // forward is hooked separately
		{Direction: DirInput, Source: "1.2.3.4", Proto: TCP, Port: 22, Action: Accept},   // has a port (advanced)
		{Direction: DirInput, Source: "1.2.3.4", Proto: TCP, Action: Accept},             // pins a protocol
		{Direction: DirInput, Action: Accept},                                            // no address
		{Direction: DirInput, Source: "1.2.3.4", Destination: "5.6.7.8", Action: Accept}, // both addresses
	}
	for _, r := range no {
		require.Falsef(t, bareHostOneWay(r), "expected NOT one-way bare host: %+v", r)
	}
}

// hostNeedsHook selects the portless address rules a csf/apf trust file cannot
// express — a concrete tcp/udp host or a source+destination pair — so AddRule
// diverts them to the raw-iptables hook instead of rejecting them. An all-
// protocol single-address host (a native plain line or a bareHostOneWay hook
// rule), a port-bearing rule, an address-less rule, and ICMP stay off this path.
func TestHostNeedsHook(t *testing.T) {
	cases := []struct {
		name string
		rule *Rule
		want bool
	}{
		{"tcp host no port", &Rule{Proto: TCP, Source: "1.2.3.4"}, true},
		{"udp host no port outbound", &Rule{Proto: UDP, Destination: "1.2.3.4", Direction: DirOutput}, true},
		{"source and destination", &Rule{Source: "1.2.3.4", Destination: "5.6.7.8"}, true},
		{"source and destination with proto", &Rule{Proto: TCP, Source: "1.2.3.4", Destination: "5.6.7.8"}, true},
		{"all-protocol single host", &Rule{Source: "1.2.3.4"}, false},
		{"tcp host with port", &Rule{Proto: TCP, Port: 22, Source: "1.2.3.4"}, false},
		{"tcp host with source port", &Rule{Proto: TCP, SourcePort: 22, Source: "1.2.3.4"}, false},
		{"address-less port rule", &Rule{Proto: TCP, Port: 22}, false},
		{"icmp host", &Rule{Proto: ICMP, Source: "1.2.3.4"}, false},
		{"icmpv6 source and destination", &Rule{Proto: ICMPv6, Source: "2001:db8::1", Destination: "2001:db8::2"}, false},
	}
	for _, c := range cases {
		require.Equal(t, c.want, hostNeedsHook(c.rule), c.name)
	}
}