package firewall

import (
	"context"
	"net/netip"
	"strings"
	"testing"
)

func TestTheRulesetGuardsOnlyLlmgwsPorts(t *testing.T) {
	s := Script(Spec{
		Ports: []Port{NodeExporter, Engines(20000, 20999), NodeExporter},
		Allow: []netip.Prefix{netip.MustParsePrefix("10.0.0.147/24"), netip.MustParsePrefix("fd00::1/64"),
			netip.MustParsePrefix("10.0.0.0/24")},
	})
	for _, want := range []string{
		"table inet llmgw\ndelete table inet llmgw\ntable inet llmgw {",
		"policy accept;",
		`iif "lo" accept`,
		"ip saddr { 10.0.0.0/24 } tcp dport { 9100, 20000-20999 } accept",
		"ip6 saddr { fd00::/64 } tcp dport { 9100, 20000-20999 } accept",
		"tcp dport { 9100, 20000-20999 } drop",
	} {
		if !strings.Contains(s, want) {
			t.Errorf("no %q in:\n%s", want, s)
		}
	}
	if strings.Contains(s, "policy drop") || strings.Contains(s, "flush ruleset") {
		t.Errorf("the ruleset reaches beyond llmgw's ports:\n%s", s)
	}
}

func TestTheNetworkIsTheInterfacesOrTheAddressAlone(t *testing.T) {
	addrs := `1: lo    inet 127.0.0.1/8 scope host lo\       valid_lft forever preferred_lft forever
2: eth0    inet 10.0.0.147/24 brd 10.0.0.255 scope global eth0\       valid_lft forever preferred_lft forever
3: eth1    inet 192.168.5.7/16 brd 192.168.255.255 scope global eth1\       valid_lft forever preferred_lft forever`
	cases := []struct {
		route, addr, want string
	}{
		{"10.0.0.140 dev eth0 src 10.0.0.147 uid 1000 \\    cache", "10.0.0.140", "10.0.0.0/24"},
		{"172.16.3.4 via 10.0.0.1 dev eth0 src 10.0.0.147 uid 1000 \\    cache", "172.16.3.4", "172.16.3.4/32"},
		{"192.168.9.9 dev eth1 src 192.168.5.7 uid 1000 \\    cache", "192.168.9.9", "192.168.0.0/16"},
		// Its own address: the network it is on.
		{"local 10.0.0.147 dev lo src 10.0.0.147 uid 1000 \\    cache <local>", "10.0.0.147", "10.0.0.0/24"},
	}
	for _, c := range cases {
		got, ok := networkToward(c.route+"\n"+addrs, netip.MustParseAddr(c.addr))
		if !ok || got.String() != c.want {
			t.Errorf("toward %s: %v %v, expected %s", c.addr, got, ok, c.want)
		}
	}
}

func TestARulesetReadsBackAsItsSpec(t *testing.T) {
	in := Spec{Ports: []Port{Postgres, Engines(20000, 20999)},
		Allow: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/24"), netip.MustParsePrefix("fd00::/64")}}
	got := parse(Script(in))
	if Script(got) != Script(in) {
		t.Errorf("read back as\n%s\nfrom\n%s", Script(got), Script(in))
	}
}

// fakeMachine answers the scripts Admit runs: the ruleset file, the route,
// and an install, which it records.
type fakeMachine struct {
	rules     string
	route     string
	installed []byte
}

func (f *fakeMachine) Run(_ context.Context, script string, stdin []byte) ([]byte, error) {
	switch {
	case strings.HasPrefix(script, "sudo cat "):
		return []byte(f.rules), nil
	case strings.HasPrefix(script, "ip -o route get"):
		return []byte(f.route), nil
	default:
		f.installed = stdin
		return []byte("changed\n"), nil
	}
}

func TestANodeOnAnotherNetworkIsLetIn(t *testing.T) {
	rules := Script(Spec{Ports: []Port{AgentPort, MirrorFiles},
		Allow: []netip.Prefix{netip.MustParsePrefix("10.104.130.128/25")}})
	m := &fakeMachine{rules: rules,
		route: "10.20.0.5 via 10.104.130.129 dev eth0 src 10.104.130.140 uid 0\n" +
			"2: eth0    inet 10.104.130.140/25 brd 10.104.130.255 scope global eth0\n"}

	added, err := Admit(context.Background(), m, []string{"10.104.130.200"})
	if err != nil || len(added) != 0 || m.installed != nil {
		t.Fatalf("an address the rules hold: added %v, err %v, installed %v", added, err, m.installed != nil)
	}

	added, err = Admit(context.Background(), m, []string{"10.20.0.5", "gpu01"})
	if err != nil {
		t.Fatal(err)
	}
	if len(added) != 1 || added[0].String() != "10.20.0.5/32" {
		t.Errorf("added %v, want the routed node's address alone", added)
	}
	got := parse(strings.Split(string(m.installed), "\x00")[0])
	if len(got.Ports) != 2 || len(got.Allow) != 2 {
		t.Errorf("installed %+v, want both ports kept and both networks", got)
	}

	// No llmgw rules on the machine: none are made.
	bare := &fakeMachine{}
	if added, err := Admit(context.Background(), bare, []string{"10.20.0.5"}); err != nil || added != nil || bare.installed != nil {
		t.Errorf("a machine without rules: added %v, err %v", added, err)
	}
}