// Package trace follows packets through the whole ruleset, every table's
// chains, with nftables' own tracing: a table of its own (housecarl-trace,
// hooked before everything else) marks the packets a filter picks
// (meta nftrace set 1), and `nft monitor trace` reports each chain and
// rule they pass, and what decided their fate. Packets are counted per
// trace id: one packet through the hooks it crosses.
package trace

import (
	"fmt"
	"net/netip"
	"regexp"
	"strconv"
	"strings"
	"time"

	"housecarl/internal/api"
)

// Table is the tracing table's name (family inet).
const Table = "housecarl-trace"

// Script builds the tracing table for a filter: packets from or to the
// address, on the port (either way), of the protocol. A limit keeps a
// flood from swamping the trace.
func Script(in api.TraceInput) (string, error) {
	var addr netip.Prefix
	if in.Address != "" {
		p, err := netip.ParsePrefix(in.Address)
		if err != nil {
			a, err2 := netip.ParseAddr(in.Address)
			if err2 != nil {
				return "", fmt.Errorf("address %q: give an address or a prefix (192.168.1.50, 2001:db8::/64)", in.Address)
			}
			p = netip.PrefixFrom(a, a.BitLen())
		}
		addr = p.Masked()
	}
	if in.Port < 0 || in.Port > 65535 {
		return "", fmt.Errorf("port %d: 1 to 65535", in.Port)
	}
	proto := strings.ToLower(in.Protocol)
	switch proto {
	case "", "tcp", "udp", "sctp", "icmp", "icmpv6":
	default:
		return "", fmt.Errorf("protocol %q: tcp, udp, sctp, icmp, icmpv6, or nothing for any", in.Protocol)
	}
	if !addr.IsValid() && in.Port == 0 {
		return "", fmt.Errorf("say what to trace: an address, a port, or both")
	}
	if in.Port != 0 && (proto == "icmp" || proto == "icmpv6") {
		return "", fmt.Errorf("ICMP has no ports")
	}
	var l4 string
	switch {
	case proto != "":
		l4 = "meta l4proto " + proto
	case in.Port != 0:
		l4 = "meta l4proto { tcp, udp, sctp }"
	}
	addrs := []string{""}
	if addr.IsValid() {
		ip := "ip"
		if addr.Addr().Is6() {
			ip = "ip6"
		}
		addrs = []string{ip + " saddr " + addr.String(), ip + " daddr " + addr.String()}
	}
	ports := []string{""}
	if in.Port != 0 {
		ports = []string{fmt.Sprintf("th dport %d", in.Port), fmt.Sprintf("th sport %d", in.Port)}
	}
	var rules []string
	for _, a := range addrs {
		for _, p := range ports {
			parts := []string{}
			for _, x := range []string{a, l4, p, "limit rate 20/second", "meta nftrace set 1"} {
				if x != "" {
					parts = append(parts, x)
				}
			}
			rules = append(rules, "\t\t"+strings.Join(parts, " "))
		}
	}
	body := strings.Join(rules, "\n")
	return fmt.Sprintf(`table inet %[1]s {}
delete table inet %[1]s
table inet %[1]s {
	chain pre {
		type filter hook prerouting priority -301; policy accept;
%[2]s
	}
	chain out {
		type filter hook output priority -301; policy accept;
%[2]s
	}
}
`, Table, body), nil
}

// ---- reading `nft monitor trace`

// "trace id 38ff3662 inet housecarl input rule … (verdict accept)"
var lineRe = regexp.MustCompile(`^trace id (\S+) (\S+) (\S+) (\S+) (packet: (.*)|rule (.*) \(verdict ([^)]*)\)|policy (\S+)|(.*))$`)

var (
	fieldRe = regexp.MustCompile(`\b(iif|oif) "([^"]*)"`)
	addrRe  = regexp.MustCompile(`\b(ip6?) (saddr|daddr) (\S+)`)
	portRe  = regexp.MustCompile(`\b(tcp|udp|sctp) (sport|dport) (\d+)`)
	flagsRe = regexp.MustCompile(`\btcp flags == (\S+)`)
	icmpRe  = regexp.MustCompile(`\b(icmp|icmpv6) type (\S+)`)
	protoRe = regexp.MustCompile(`\bmeta l4proto (\S+)|\bip protocol (\S+)|\bip6 nexthdr (\S+)`)
	markRe  = regexp.MustCompile(`comment "hc:(\S+)`)
)

// Tracer collects packets from trace lines.
type Tracer struct {
	// Max is how many packets it keeps (0: 200).
	Max     int
	Now     func() time.Time
	packets []*api.TracePacket
	byID    map[string]*api.TracePacket
	// seen are the chains a packet entered: the kernel reuses trace ids
	// (they come from the packet's buffer), so entering one again means
	// another packet.
	seen map[*api.TracePacket]map[string]bool
}

// Packets returns what was traced so far, oldest first.
func (t *Tracer) Packets() []api.TracePacket {
	out := make([]api.TracePacket, 0, len(t.packets))
	for _, p := range t.packets {
		out = append(out, *p)
	}
	return out
}

// Full reports whether it keeps no more packets.
func (t *Tracer) Full() bool {
	max := t.Max
	if max == 0 {
		max = 200
	}
	return len(t.packets) >= max
}

// Line reads one line of `nft monitor trace`. The tracing table's own
// steps are left out.
func (t *Tracer) Line(line string) {
	m := lineRe.FindStringSubmatch(strings.TrimSpace(line))
	if m == nil {
		return
	}
	id, table, chain := m[1], m[2]+" "+m[3], m[4]
	if m[3] == Table {
		return
	}
	if t.byID == nil {
		t.byID, t.seen = map[string]*api.TracePacket{}, map[*api.TracePacket]map[string]bool{}
	}
	p := t.byID[id]
	isPacket := m[6] != "" || strings.HasPrefix(m[5], "packet:")
	if p != nil && isPacket && t.seen[p][table+" "+chain] {
		p = nil // the same id, another packet
	}
	if p == nil {
		if t.Full() {
			return
		}
		now := time.Now
		if t.Now != nil {
			now = t.Now
		}
		p = &api.TracePacket{ID: id, Time: now(), Verdict: "accepted"}
		t.byID[id] = p
		t.seen[p] = map[string]bool{}
		t.packets = append(t.packets, p)
	}
	switch {
	case isPacket:
		t.seen[p][table+" "+chain] = true
		readPacket(p, m[6])
	case m[7] != "":
		step := api.TraceStep{Table: table, Chain: chain, Kind: "rule", Rule: m[7], Verdict: m[8]}
		if mk := markRe.FindStringSubmatch(m[7]); mk != nil {
			step.Mark = strings.TrimSuffix(mk[1], `"`)
		}
		p.Steps = append(p.Steps, step)
		decide(p, step)
	case m[9] != "":
		step := api.TraceStep{Table: table, Chain: chain, Kind: "policy", Verdict: m[9]}
		p.Steps = append(p.Steps, step)
		decide(p, step)
	}
}

// decide notes a packet's fate: the first drop (or reject) ends it.
func decide(p *api.TracePacket, s api.TraceStep) {
	if p.Verdict == "dropped" {
		return
	}
	switch strings.Fields(s.Verdict + " ")[0] {
	case "drop", "reject":
		p.Verdict, p.DecidedBy = "dropped", s
	case "accept":
		// What let it through is a filter chain's accept: Housecarl's
		// input, forward or output (NAT chains accept everything they
		// see). Another table's accept counts when Housecarl's says
		// nothing.
		filter := s.Table == "inet housecarl" && (s.Chain == "input" || s.Chain == "forward" || s.Chain == "output")
		if filter || p.DecidedBy.Verdict == "" && s.Table != "inet housecarl" {
			p.DecidedBy = s
		}
	}
}

// readPacket fills in a packet's headers from its "packet:" line (once:
// the first hook's view, before any NAT).
func readPacket(p *api.TracePacket, f string) {
	for _, m := range fieldRe.FindAllStringSubmatch(f, -1) {
		if m[1] == "iif" && p.In == "" {
			p.In = m[2]
		}
		if m[1] == "oif" && p.Out == "" {
			p.Out = m[2]
		}
	}
	if p.Source != "" {
		return
	}
	for _, m := range addrRe.FindAllStringSubmatch(f, -1) {
		if m[2] == "saddr" {
			p.Source = m[3]
		} else {
			p.Destination = m[3]
		}
	}
	for _, m := range portRe.FindAllStringSubmatch(f, -1) {
		p.Protocol = m[1]
		n, _ := strconv.Atoi(m[3])
		if m[2] == "sport" {
			p.SourcePort = n
		} else {
			p.DestinationPort = n
		}
	}
	if m := flagsRe.FindStringSubmatch(f); m != nil {
		p.TCPFlags = m[1]
	}
	if m := icmpRe.FindStringSubmatch(f); m != nil {
		p.Protocol, p.ICMPType = m[1], m[2]
	}
	if p.Protocol == "" {
		if m := protoRe.FindStringSubmatch(f); m != nil {
			p.Protocol = strings.TrimSpace(m[1] + m[2] + m[3])
		}
	}
}