package main

import (
	"net"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/google/nftables"
	"github.com/vishvananda/netlink"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)

func TestNetworkLifecycle(t *testing.T) {
	if os.Getenv("RELAY_NETWORK_TEST") != "1" {
		t.Skip("RELAY_NETWORK_TEST is not enabled")
	}
	privateKey, err := wgtypes.GeneratePrivateKey()
	if err != nil {
		t.Fatal(err)
	}
	peerKey, err := wgtypes.GeneratePrivateKey()
	if err != nil {
		t.Fatal(err)
	}
	privateKeyPath := filepath.Join(t.TempDir(), "private_key")
	if err := os.WriteFile(privateKeyPath, []byte(privateKey.String()), 0400); err != nil {
		t.Fatal(err)
	}
	tunnelAddress, err := parseCIDR("198.18.0.2/30")
	if err != nil {
		t.Fatal(err)
	}
	localReturnCIDR, err := parseCIDR("203.0.113.0/24")
	if err != nil {
		t.Fatal(err)
	}
	targets := []relayTarget{
		{address: net.ParseIP("198.18.0.10").To4(), ports: []uint16{80, 443}},
		{address: net.ParseIP("198.18.0.11").To4(), ports: []uint16{25, 465, 993}},
	}
	loopback, err := netlink.LinkByName("lo")
	if err != nil {
		t.Fatal(err)
	}
	if err := netlink.LinkSetUp(loopback); err != nil {
		t.Fatal(err)
	}
	for _, target := range targets {
		if err := netlink.AddrAdd(loopback, &netlink.Addr{
			IPNet: &net.IPNet{IP: target.address, Mask: net.CIDRMask(32, 32)},
		}); err != nil {
			t.Fatal(err)
		}
	}
	if err := netlink.RouteAdd(&netlink.Route{
		LinkIndex: loopback.Attrs().Index,
		Dst:       localReturnCIDR,
		Scope:     netlink.SCOPE_LINK,
	}); err != nil {
		t.Fatal(err)
	}
	instance := &relay{config: config{
		privateKeyPath:   privateKeyPath,
		peerPublicKey:    peerKey.PublicKey().String(),
		peerEndpoint:     "192.0.2.1:51820",
		listenPort:       51820,
		tunnelAddress:    tunnelAddress,
		targets:          targets,
		endpointAddress:  net.ParseIP("192.0.2.1").To4(),
		localReturnCIDRs: []*net.IPNet{localReturnCIDR},
	}}
	if err := instance.configure(); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		if err := instance.cleanup(); err != nil {
			t.Error(err)
		}
	})
	if err := instance.refreshFence(); err != nil {
		t.Fatal(err)
	}
	if err := verifyFirewall(instance.config.endpointAddress, instance.config.targets); err != nil {
		t.Fatal(err)
	}
	if _, err := netlink.LinkByName(interfaceName); err != nil {
		t.Fatal(err)
	}
	rules, err := netlink.RuleList(netlink.FAMILY_V4)
	if err != nil {
		t.Fatal(err)
	}
	for _, target := range instance.config.targets {
		foundRule := false
		for _, rule := range rules {
			if isRelayPolicyRule(rule, target.address) {
				foundRule = true
			}
		}
		if !foundRule {
			t.Fatalf("relay policy rule for %s is absent", target.address)
		}
		externalRoutes, err := netlink.RouteGetWithOptions(
			net.ParseIP("192.0.2.2").To4(),
			&netlink.RouteGetOptions{SrcAddr: target.address},
		)
		if err != nil {
			t.Fatal(err)
		}
		if len(externalRoutes) != 1 || externalRoutes[0].LinkIndex != linkIndex(t, interfaceName) {
			t.Fatalf("external traffic from %s does not select %s: %+v", target.address, interfaceName, externalRoutes)
		}
		localRoutes, err := netlink.RouteGetWithOptions(
			net.ParseIP("203.0.113.1").To4(),
			&netlink.RouteGetOptions{SrcAddr: target.address},
		)
		if err != nil {
			t.Fatal(err)
		}
		if len(localRoutes) != 1 || localRoutes[0].LinkIndex != loopback.Attrs().Index {
			t.Fatalf("local traffic from %s does not select the main table: %+v", target.address, localRoutes)
		}
	}
	foundTable := false
	tables, err := (&nftables.Conn{}).ListTables()
	if err != nil {
		t.Fatal(err)
	}
	for _, table := range tables {
		if table.Family == nftables.TableFamilyINet && table.Name == firewallTable {
			foundTable = true
		}
	}
	if !foundTable {
		t.Fatal("relay firewall table is absent")
	}
	time.Sleep(16 * time.Second)
	if err := verifyFirewall(instance.config.endpointAddress, instance.config.targets); err == nil {
		t.Fatal("expected the unrefreshed ownership fence to expire")
	}
	if err := instance.cleanup(); err != nil {
		t.Fatal(err)
	}
	if _, err := netlink.LinkByName(interfaceName); err == nil {
		t.Fatalf("%s remains after cleanup", interfaceName)
	}
	rules, err = netlink.RuleList(netlink.FAMILY_V4)
	if err != nil {
		t.Fatal(err)
	}
	for _, target := range instance.config.targets {
		for _, rule := range rules {
			if isRelayPolicyRule(rule, target.address) ||
				isLocalReturnRule(rule, target.address, localReturnCIDR, localPriority) {
				t.Fatalf("policy rule for %s remains after cleanup: %+v", target.address, rule)
			}
		}
	}
}

func linkIndex(t *testing.T, name string) int {
	t.Helper()
	link, err := netlink.LinkByName(name)
	if err != nil {
		t.Fatal(err)
	}
	return link.Attrs().Index
}