// Package conntrack lists and closes the kernel's tracked connections,
// through ctnetlink (the netfilter netlink subsystem for connection
// tracking): what the "established" built-in lets through, and what a
// change of rules doesn't stop by itself.
//
// Closing a connection deletes its tracking entry. Its next packet is
// then checked by the rules again, as a new connection would be: one the
// rules now block is stopped, one they allow carries on (TCP connections
// are picked up again mid-stream).
package conntrack

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"net/netip"
	"sync/atomic"
	"syscall"
	"time"

	"golang.org/x/sys/unix"

	"housecarl/internal/api"
)

const (
	subsysCtnetlink = 1
	msgNew          = 0
	msgGet          = 1
	msgDelete       = 2

	ctaTupleOrig     = 1
	ctaTupleReply    = 2
	ctaStatus        = 3
	ctaProtoinfo     = 4
	ctaTimeout       = 7
	ctaMark          = 8
	ctaCountersOrig  = 9
	ctaCountersReply = 10
	ctaID            = 12

	ctaTupleIP    = 1
	ctaTupleProto = 2

	ctaIPv4Src = 1
	ctaIPv4Dst = 2
	ctaIPv6Src = 3
	ctaIPv6Dst = 4

	ctaProtoNum      = 1
	ctaProtoSrcPort  = 2
	ctaProtoDstPort  = 3
	ctaProtoICMPID   = 4
	ctaProtoICMPType = 5
	ctaProtoICMPCode = 6
	ctaProtoICMP6ID  = 7
	ctaProtoICMP6Typ = 8
	ctaProtoICMP6Cod = 9

	ctaProtoinfoTCP      = 1
	ctaProtoinfoTCPState = 1

	ctaCountersPackets = 1
	ctaCountersBytes   = 2

	statusAssured = 1 << 2
	statusSrcNAT  = 1 << 4
	statusDstNAT  = 1 << 5

	nlaNested   = 0x8000
	nlaTypeMask = 0x3fff
)

var tcpStates = []string{"none", "syn-sent", "syn-recv", "established", "fin-wait", "close-wait", "last-ack", "time-wait", "close", "syn-sent2"}

var protoNames = map[uint8]string{1: "icmp", 6: "tcp", 17: "udp", 58: "icmpv6", 132: "sctp", 136: "udplite", 33: "dccp", 47: "gre"}

// Client talks to ctnetlink. The zero value works.
type Client struct{}

var seq atomic.Uint32

// List returns every tracked connection, IPv4 and IPv6.
func (Client) List(ctx context.Context) ([]api.Connection, error) {
	msgs, err := request(ctx, msgGet, unix.NLM_F_REQUEST|unix.NLM_F_DUMP, unix.AF_UNSPEC, nil)
	if err != nil {
		return nil, err
	}
	out := []api.Connection{}
	for _, m := range msgs {
		if c, ok := parse(m); ok {
			out = append(out, c)
		}
	}
	return out, nil
}

// Close deletes a connection's tracking entry, found by its original
// direction (addresses, protocol, ports or ICMP id). Gone already is fine.
func (Client) Close(ctx context.Context, c api.Connection) error {
	src, err1 := netip.ParseAddr(c.Source)
	dst, err2 := netip.ParseAddr(c.Destination)
	if err1 != nil || err2 != nil || src.Is4() != dst.Is4() {
		return fmt.Errorf("connection %s -> %s: not a pair of addresses", c.Source, c.Destination)
	}
	num := protoNumber(c.Protocol)
	if num == 0 {
		return fmt.Errorf("connection %s -> %s: unknown protocol %q", c.Source, c.Destination, c.Protocol)
	}
	family, ipSrc, ipDst := uint8(unix.AF_INET), uint16(ctaIPv4Src), uint16(ctaIPv4Dst)
	if src.Is6() {
		family, ipSrc, ipDst = unix.AF_INET6, ctaIPv6Src, ctaIPv6Dst
	}
	ip := attr(ipSrc, src.AsSlice())
	ip = append(ip, attr(ipDst, dst.AsSlice())...)
	proto := attr(ctaProtoNum, []byte{num})
	switch num {
	case 1:
		proto = append(proto, attr(ctaProtoICMPID, be16(uint16(c.ICMPID)))...)
		proto = append(proto, attr(ctaProtoICMPType, []byte{uint8(c.ICMPType)})...)
		proto = append(proto, attr(ctaProtoICMPCode, []byte{uint8(c.ICMPCode)})...)
	case 58:
		proto = append(proto, attr(ctaProtoICMP6ID, be16(uint16(c.ICMPID)))...)
		proto = append(proto, attr(ctaProtoICMP6Typ, []byte{uint8(c.ICMPType)})...)
		proto = append(proto, attr(ctaProtoICMP6Cod, []byte{uint8(c.ICMPCode)})...)
	default:
		proto = append(proto, attr(ctaProtoSrcPort, be16(uint16(c.SourcePort)))...)
		proto = append(proto, attr(ctaProtoDstPort, be16(uint16(c.DestinationPort)))...)
	}
	tuple := attr(ctaTupleIP|nlaNested, ip)
	tuple = append(tuple, attr(ctaTupleProto|nlaNested, proto)...)
	_, err := request(ctx, msgDelete, unix.NLM_F_REQUEST|unix.NLM_F_ACK, family, attr(ctaTupleOrig|nlaNested, tuple))
	if errors.Is(err, unix.ENOENT) {
		return nil
	}
	return err
}

func protoNumber(name string) uint8 {
	for n, s := range protoNames {
		if s == name {
			return n
		}
	}
	return 0
}

// ---- netlink

func request(ctx context.Context, msg uint16, flags uint16, family uint8, payload []byte) ([][]byte, error) {
	fd, err := unix.Socket(unix.AF_NETLINK, unix.SOCK_RAW|unix.SOCK_CLOEXEC, unix.NETLINK_NETFILTER)
	if err != nil {
		return nil, err
	}
	defer unix.Close(fd)
	if err := unix.Bind(fd, &unix.SockaddrNetlink{Family: unix.AF_NETLINK}); err != nil {
		return nil, err
	}
	deadline := time.Now().Add(10 * time.Second)
	if d, ok := ctx.Deadline(); ok && d.Before(deadline) {
		deadline = d
	}
	tv := unix.NsecToTimeval(time.Until(deadline).Nanoseconds())
	unix.SetsockoptTimeval(fd, unix.SOL_SOCKET, unix.SO_RCVTIMEO, &tv)
	unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_RCVBUF, 1<<22)

	s := seq.Add(1)
	b := make([]byte, unix.SizeofNlMsghdr+4, unix.SizeofNlMsghdr+4+len(payload))
	binary.NativeEndian.PutUint16(b[4:], subsysCtnetlink<<8|msg)
	binary.NativeEndian.PutUint16(b[6:], flags)
	binary.NativeEndian.PutUint32(b[8:], s)
	b[16] = family // nfgenmsg: family, version 0, res_id 0
	b = append(b, payload...)
	binary.NativeEndian.PutUint32(b[0:], uint32(len(b)))
	if err := unix.Sendto(fd, b, 0, &unix.SockaddrNetlink{Family: unix.AF_NETLINK}); err != nil {
		return nil, err
	}
	var out [][]byte
	buf := make([]byte, 1<<16)
	for {
		n, _, err := unix.Recvfrom(fd, buf, 0)
		if err != nil {
			if errors.Is(err, unix.EAGAIN) {
				return nil, errors.New("no answer from connection tracking (is nf_conntrack loaded?)")
			}
			return nil, err
		}
		for off := 0; off+unix.SizeofNlMsghdr <= n; {
			l := int(binary.NativeEndian.Uint32(buf[off:]))
			typ := binary.NativeEndian.Uint16(buf[off+4:])
			if l < unix.SizeofNlMsghdr || off+l > n {
				return nil, errors.New("connection tracking: a malformed answer")
			}
			body := buf[off+unix.SizeofNlMsghdr : off+l]
			switch typ {
			case unix.NLMSG_DONE:
				return out, nil
			case unix.NLMSG_ERROR:
				if len(body) >= 4 {
					if code := int32(binary.NativeEndian.Uint32(body)); code != 0 {
						return nil, syscall.Errno(-code)
					}
				}
				return out, nil // the acknowledgement
			default:
				out = append(out, append([]byte(nil), body...))
			}
			off += (l + 3) &^ 3
		}
		if flags&unix.NLM_F_DUMP == 0 && flags&unix.NLM_F_ACK == 0 {
			return out, nil
		}
	}
}

func attr(typ uint16, data []byte) []byte {
	l := 4 + len(data)
	b := make([]byte, (l+3)&^3)
	binary.NativeEndian.PutUint16(b[0:], uint16(l))
	binary.NativeEndian.PutUint16(b[2:], typ)
	copy(b[4:], data)
	return b
}

func be16(v uint16) []byte {
	b := make([]byte, 2)
	binary.BigEndian.PutUint16(b, v)
	return b
}

// attrs splits netlink attributes: type (without the nested flag) → data.
func attrs(b []byte) map[uint16][]byte {
	m := map[uint16][]byte{}
	for len(b) >= 4 {
		l := int(binary.NativeEndian.Uint16(b[0:]))
		typ := binary.NativeEndian.Uint16(b[2:]) & nlaTypeMask
		if l < 4 || l > len(b) {
			break
		}
		m[typ] = b[4:l]
		b = b[min((l+3)&^3, len(b)):]
	}
	return m
}

// ---- parsing

type tuple struct {
	src, dst         netip.Addr
	proto            uint8
	sport, dport     uint16
	icmpID           uint16
	icmpType, icmpCd uint8
}

func parseTuple(b []byte) (tuple, bool) {
	var t tuple
	a := attrs(b)
	ip := attrs(a[ctaTupleIP])
	if v, ok := ip[ctaIPv4Src]; ok {
		t.src, _ = netip.AddrFromSlice(v)
		t.dst, _ = netip.AddrFromSlice(ip[ctaIPv4Dst])
	} else if v, ok := ip[ctaIPv6Src]; ok {
		t.src, _ = netip.AddrFromSlice(v)
		t.dst, _ = netip.AddrFromSlice(ip[ctaIPv6Dst])
	}
	p := attrs(a[ctaTupleProto])
	if v := p[ctaProtoNum]; len(v) == 1 {
		t.proto = v[0]
	}
	if v := p[ctaProtoSrcPort]; len(v) == 2 {
		t.sport = binary.BigEndian.Uint16(v)
	}
	if v := p[ctaProtoDstPort]; len(v) == 2 {
		t.dport = binary.BigEndian.Uint16(v)
	}
	for _, k := range [][3]uint16{{ctaProtoICMPID, ctaProtoICMPType, ctaProtoICMPCode}, {ctaProtoICMP6ID, ctaProtoICMP6Typ, ctaProtoICMP6Cod}} {
		if v := p[k[0]]; len(v) == 2 {
			t.icmpID = binary.BigEndian.Uint16(v)
		}
		if v := p[k[1]]; len(v) == 1 {
			t.icmpType = v[0]
		}
		if v := p[k[2]]; len(v) == 1 {
			t.icmpCd = v[0]
		}
	}
	return t, t.src.IsValid() && t.dst.IsValid()
}

// parse reads one connection (an nfgenmsg and its attributes).
func parse(m []byte) (api.Connection, bool) {
	if len(m) < 4 {
		return api.Connection{}, false
	}
	a := attrs(m[4:])
	orig, ok1 := parseTuple(a[ctaTupleOrig])
	reply, ok2 := parseTuple(a[ctaTupleReply])
	if !ok1 || !ok2 {
		return api.Connection{}, false
	}
	c := api.Connection{
		Family: "ipv4", Protocol: protoNames[orig.proto],
		Source: orig.src.String(), SourcePort: int(orig.sport), Destination: orig.dst.String(), DestinationPort: int(orig.dport),
		ICMPID: int(orig.icmpID), ICMPType: int(orig.icmpType), ICMPCode: int(orig.icmpCd),
	}
	if orig.src.Is6() {
		c.Family = "ipv6"
	}
	if c.Protocol == "" {
		c.Protocol = fmt.Sprint(orig.proto)
	}
	// NAT shows in the reply direction: it comes from where the packets
	// really went, to where they really came from.
	if reply.src != orig.dst || reply.sport != orig.dport {
		c.DNATTo = reply.src.String()
		if reply.sport != orig.dport {
			c.DNATPort = int(reply.sport)
		}
	}
	if reply.dst != orig.src || reply.dport != orig.sport {
		c.SNATTo = reply.dst.String()
	}
	if v := a[ctaStatus]; len(v) == 4 {
		st := binary.BigEndian.Uint32(v)
		c.Assured = st&statusAssured != 0
	}
	if v := a[ctaTimeout]; len(v) == 4 {
		c.Timeout = int(binary.BigEndian.Uint32(v))
	}
	if v := a[ctaMark]; len(v) == 4 {
		c.Mark = binary.BigEndian.Uint32(v)
	}
	if v := a[ctaID]; len(v) == 4 {
		c.ID = binary.BigEndian.Uint32(v)
	}
	if v := attrs(attrs(a[ctaProtoinfo])[ctaProtoinfoTCP])[ctaProtoinfoTCPState]; len(v) == 1 && int(v[0]) < len(tcpStates) {
		c.State = tcpStates[v[0]]
	}
	for _, k := range []uint16{ctaCountersOrig, ctaCountersReply} {
		ct := attrs(a[k])
		if v := ct[ctaCountersPackets]; len(v) == 8 {
			c.Packets += binary.BigEndian.Uint64(v)
		}
		if v := ct[ctaCountersBytes]; len(v) == 8 {
			c.Bytes += binary.BigEndian.Uint64(v)
		}
	}
	return c, true
}