package kubenat

import (
	"context"
	"flag"
	"io/ioutil"
	"net"
	"os"
	"testing"

	"github.com/go-test/deep"
)

// testConntrack is the anonymized content of a production host.
// The first entry is an appservice-irc connection from a pod to an IRC server.
// The second connection is an UDP connection between two pods.
// The third to last entry is not a NAT entry, but an incoming external
// connection.
// The fourth connection has a mangled/incomplete entry.
const testConntrack = `
ipv4     2 tcp      6 86384 ESTABLISHED src=10.10.26.23 dst=192.0.2.180 sport=51336 dport=6697 src=192.0.2.180 dst=185.236.240.36 sport=6697 dport=28706 [ASSURED] mark=0 zone=0 use=2
ipv4     2 udp      17 35 src=10.10.24.162 dst=10.10.26.108 sport=49347 dport=53 src=10.10.26.108 dst=10.10.24.162 sport=53 dport=49347 [ASSURED] mark=0 zone=0 use=2
ipv4     2 tcp      6 2 SYN_SENT src=198.51.100.67 dst=185.236.240.56 sport=51053 dport=3359 [UNREPLIED] src=185.236.240.56 dst=198.51.100.67 sport=3359 dport=51053 mark=0 zone=0 use=2
ipv4     2 tcp      6 2
`

// TestConntrackParse exercises the conntrack parser for all entries in testConntrack.
func TestConntrackParse(t *testing.T) {
	// Last line is truncated and should be ignored.
	got, err := conntrackParse([]byte(testConntrack))
	if err != nil {
		t.Fatalf("conntrackParse: %v", err)
	}
	want := []conntrackEntry{
		{
			"ipv4", "tcp", 86384, "ESTABLISHED",
			map[string]string{
				"src": "10.10.26.23", "dst": "192.0.2.180", "sport": "57640", "dport": "6697",
				"mark": "0", "zone": "0", "use": "2",
			},
			map[string]string{
				"src": "192.0.2.180", "dst": "185.236.240.36", "sport": "6697", "dport": "28706",
			},
			map[string]bool{
				"ASSURED": true,
			},
		},
		{
			"ipv4", "udp", 35, "",
			map[string]string{
				"src": "10.10.24.162", "dst": "10.10.26.108", "sport": "49347", "dport": "53",
				"mark": "0", "zone": "0", "use": "2",
			},
			map[string]string{
				"src": "10.10.26.108", "dst": "10.10.24.162", "sport": "53", "dport": "49347",
			},
			map[string]bool{
				"ASSURED": true,
			},
		},
		{
			"ipv4", "tcp", 2, "SYN_SENT",
			map[string]string{
				"src": "198.51.100.67", "dst": "185.236.240.56", "sport": "51053", "dport": "3359",
				"mark": "0", "zone": "0", "use": "2",
			},
			map[string]string{
				"src": "185.236.240.56", "dst": "198.51.100.67", "sport": "3359", "dport": "51053",
			},
			map[string]bool{
				"UNREPLIED": true,
			},
		},
	}
	if diff := deep.Equal(want, got); diff != nil {
		t.Error(diff)
	}

	ix := buildIndex(got)
	if want, got := 0, len(ix.getByRequest("src", "1.2.3.4")); want != got {
		t.Errorf("by request, src, 1.2.3.4 should have returned %d result, wanted %d", want, got)
	}
	if want, got := 1, len(ix.getByRequest("src", "10.10.26.23")); want != got {
		t.Errorf("by request, src, 1.2.3.4 should have returned %d result, wanted %d", want, got)
	}
	if want, got := "10.10.26.23", ix.getByRequest("src", "10.10.26.23")[0].request["src"]; want != got {
		t.Errorf("by request, wanted src %q, got %q", want, got)
	}
	if want, got := 3, len(ix.getByRequest("mark", "0")); want != got {
		t.Errorf("by request, mark, 0 should have returned %d result, wanted %d", want, got)
	}
}

// TestTranslationWorker exercises a translation worker with a
// testConntrack-backed conntrack file.
func TestTranslationWorker(t *testing.T) {
	flag.Set("logtostderr", "true")
	tmpfile, err := ioutil.TempFile("", "conntack")
	if err != nil {
		t.Fatal(err)
	}
	defer os.Remove(tmpfile.Name())
	if _, err := tmpfile.Write([]byte(testConntrack)); err != nil {
		t.Fatal(err)
	}
	r := &Resolver{
		conntrackPath: tmpfile.Name(),
		translationC:  make(chan *translationReq),
	}
	ctx, ctxC := context.WithCancel(context.Background())
	defer ctxC()

	go r.runTranslationWorker(ctx)

	res, err := r.translate(ctx, &Tuple4{
		RemoteIP:   net.ParseIP("192.0.2.180"),
		RemotePort: 6697,
		LocalIP:    net.ParseIP("185.236.240.36"),
		LocalPort:  28706,
	})
	if err != nil {
		t.Fatalf("translate: %v", err)
	}
	if want, got := net.ParseIP("10.10.26.23"), res.localIP; !want.Equal(got) {
		t.Errorf("local ip: wanted %v, got %v", want, got)
	}
	if want, got := uint16(51336), res.localPort; want != got {
		t.Errorf("local port: wanted %d, got %d", want, got)
	}
}