package firewall

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"os/exec"
	"strings"
	"testing"
	"time"
)

func TestDNSForwardingInNamespace(t *testing.T) {
	if os.Getenv("DEBTIK_NAT_INTEGRATION") != "1" {
		t.Skip("isolated network namespace required")
	}
	run := func(args ...string) {
		t.Helper()
		out, err := exec.Command("ip", args...).CombinedOutput()
		if err != nil {
			t.Fatalf("ip %v: %s: %v", args, out, err)
		}
	}
	run("link", "set", "lo", "up")
	run("addr", "add", "192.0.2.53/32", "dev", "lo")
	run("link", "add", "fwserver", "type", "veth", "peer", "name", "fwclient")
	run("addr", "add", "192.0.2.1/24", "dev", "fwserver")
	run("link", "set", "fwserver", "up")
	udp, err := net.ListenPacket("udp4", "192.0.2.53:1053")
	if err != nil {
		t.Fatal(err)
	}
	defer udp.Close()
	tcp, err := net.Listen("tcp4", "192.0.2.53:1053")
	if err != nil {
		t.Fatal(err)
	}
	defer tcp.Close()
	go func() {
		b := make([]byte, 16)
		n, a, e := udp.ReadFrom(b)
		if e == nil {
			udp.WriteTo(b[:n], a)
		}
	}()
	go func() {
		c, e := tcp.Accept()
		if e == nil {
			defer c.Close()
			b := make([]byte, 16)
			n, e := c.Read(b)
			if e == nil {
				c.Write(b[:n])
			}
		}
	}()
	rules := []NATRule{{ID: "nat-dns-udp", Chain: "dstnat", Action: "dst-nat", Protocol: "udp", DstPort: "53", ToAddresses: "192.0.2.53", ToPorts: "1053"}, {ID: "nat-dns-tcp", Chain: "dstnat", Action: "dst-nat", Protocol: "tcp", DstPort: "53", ToAddresses: "192.0.2.53", ToPorts: "1053"}}
	if err := ApplyDualStack(nil, rules, nil); err != nil {
		t.Fatal(err)
	}
	script := `import socket,subprocess,sys
print('READY',flush=True)
sys.stdin.readline()
for args in [['link','set','lo','up'],['addr','add','192.0.2.2/24','dev','fwclient'],['link','set','fwclient','up'],['route','add','default','via','192.0.2.1']]: subprocess.run(['ip']+args,check=True)
for typ in [socket.SOCK_DGRAM,socket.SOCK_STREAM]:
 s=socket.socket(socket.AF_INET,typ);s.settimeout(3);s.connect(('198.51.100.53',53));s.sendall(b'dns');assert s.recv(3)==b'dns';s.close()
print('UDP and TCP forwarded successfully',flush=True)`
	cmd := exec.Command("unshare", "--net", "python3", "-u", "-c", script)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		t.Fatal(err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		t.Fatal(err)
	}
	cmd.Stderr = os.Stderr
	if err = cmd.Start(); err != nil {
		t.Fatal(err)
	}
	defer func() {
		if cmd.ProcessState == nil {
			cmd.Process.Kill()
			cmd.Wait()
		}
	}()
	reader := bufio.NewReader(stdout)
	ready, err := reader.ReadString('\n')
	if err != nil || strings.TrimSpace(ready) != "READY" {
		t.Fatalf("client not ready: %q %v", ready, err)
	}
	run("link", "set", "fwclient", "netns", fmt.Sprint(cmd.Process.Pid))
	fmt.Fprintln(stdin, "go")
	stdin.Close()
	done := make(chan error, 1)
	go func() { done <- cmd.Wait() }()
	select {
	case err := <-done:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(10 * time.Second):
		cmd.Process.Kill()
		t.Fatal("client timed out")
	}
}

func TestLocalOutputDNATInNamespace(t *testing.T) {
	if os.Getenv("DEBTIK_NAT_INTEGRATION") != "1" {
		t.Skip("isolated network namespace required")
	}
	if out, err := exec.Command("ip", "link", "set", "lo", "up").CombinedOutput(); err != nil {
		t.Fatalf("loopback: %s: %v", out, err)
	}
	server, err := net.ListenPacket("udp4", "127.0.0.2:1053")
	if err != nil {
		t.Fatal(err)
	}
	defer server.Close()
	go func() {
		b := make([]byte, 16)
		n, a, e := server.ReadFrom(b)
		if e == nil {
			server.WriteTo(b[:n], a)
		}
	}()
	rule := NATRule{ID: "nat-local", Chain: "output", Action: "dst-nat", Protocol: "udp", DstAddress: "127.0.0.1", DstPort: "53", ToAddresses: "127.0.0.2", ToPorts: "1053"}
	if err := ValidateRule("nat", rule); err != nil {
		t.Fatal(err)
	}
	if err := ApplyDualStack(nil, []NATRule{rule}, nil); err != nil {
		t.Fatal(err)
	}
	client, err := net.DialTimeout("udp4", "127.0.0.1:53", 2*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	client.SetDeadline(time.Now().Add(2 * time.Second))
	if _, err := client.Write([]byte("dns")); err != nil {
		t.Fatal(err)
	}
	b := make([]byte, 3)
	if n, err := client.Read(b); err != nil || string(b[:n]) != "dns" {
		t.Fatalf("local DNAT response %q: %v", b[:n], err)
	}
}