package conntrack

import (
	"encoding/binary"
	"net/netip"
	"testing"
)

func nested(typ uint16, parts ...[]byte) []byte {
	var b []byte
	for _, p := range parts {
		b = append(b, p...)
	}
	return attr(typ|nlaNested, b)
}

func tupleAttr(typ uint16, src, dst string, proto uint8, sport, dport uint16) []byte {
	s, d := netip.MustParseAddr(src), netip.MustParseAddr(dst)
	return nested(typ,
		nested(ctaTupleIP, attr(ctaIPv4Src, s.AsSlice()), attr(ctaIPv4Dst, d.AsSlice())),
		nested(ctaTupleProto, attr(ctaProtoNum, []byte{proto}), attr(ctaProtoSrcPort, be16(sport)), attr(ctaProtoDstPort, be16(dport))))
}

// A port-forwarded TCP connection from the internet, as ctnetlink
// describes it: the reply comes from the inside host.
func TestParse(t *testing.T) {
	status, timeout := make([]byte, 4), make([]byte, 4)
	binary.BigEndian.PutUint32(status, statusAssured|statusDstNAT)
	binary.BigEndian.PutUint32(timeout, 431999)
	msg := []byte{2, 0, 0, 0} // nfgenmsg: AF_INET
	msg = append(msg, tupleAttr(ctaTupleOrig, "198.51.100.7", "203.0.113.10", 6, 51432, 8443)...)
	msg = append(msg, tupleAttr(ctaTupleReply, "192.168.1.10", "198.51.100.7", 6, 443, 51432)...)
	msg = append(msg, attr(ctaStatus, status)...)
	msg = append(msg, attr(ctaTimeout, timeout)...)
	msg = append(msg, nested(ctaProtoinfo, nested(ctaProtoinfoTCP, attr(ctaProtoinfoTCPState, []byte{3})))...)
	c, ok := parse(msg)
	if !ok {
		t.Fatal("not parsed")
	}
	if c.Protocol != "tcp" || c.Source != "198.51.100.7" || c.SourcePort != 51432 || c.Destination != "203.0.113.10" || c.DestinationPort != 8443 ||
		c.DNATTo != "192.168.1.10" || c.DNATPort != 443 || c.SNATTo != "" || c.State != "established" || !c.Assured || c.Timeout != 431999 {
		t.Errorf("parsed %+v", c)
	}
	// Masqueraded on the way out: the reply comes back to the WAN address.
	msg = []byte{2, 0, 0, 0}
	msg = append(msg, tupleAttr(ctaTupleOrig, "192.168.1.20", "198.51.100.9", 17, 5353, 53)...)
	msg = append(msg, tupleAttr(ctaTupleReply, "198.51.100.9", "203.0.113.10", 17, 53, 5353)...)
	if c, _ := parse(msg); c.SNATTo != "203.0.113.10" || c.DNATTo != "" || c.Protocol != "udp" {
		t.Errorf("masqueraded: %+v", c)
	}
}