// Package firewall writes and applies the nftables rules that keep llmgw's
// internal ports to the cluster's own networks.
//
// Only llmgw's ports, in a table of its own. The machines llmgw is put on
// are often not llmgw's alone — a GPU node shared with a desktop's ComfyUI,
// a database host with other databases — so the policy is not "deny
// everything else": it is "these ports answer the cluster's networks and
// this machine, and nobody else". Whatever else the machine serves, and
// whatever firewall it already has, is left as it was. The client API
// (:8080) is not among them: it is meant to be reached.
package firewall

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net/netip"
	"os"
	"os/exec"
	"slices"
	"strconv"
	"strings"
)

// Runner runs a shell script on a machine, as onboarding's SSH and the
// local machine both do.
type Runner interface {
	Run(ctx context.Context, script string, stdin []byte) ([]byte, error)
}

// Port is one port, or a range, and what answers there.
type Port struct {
	From, To int
	What     string
}

// Single is one port.
func Single(port int, what string) Port { return Port{From: port, To: port, What: what} }

func (p Port) nft() string {
	if p.To > p.From {
		return strconv.Itoa(p.From) + "-" + strconv.Itoa(p.To)
	}
	return strconv.Itoa(p.From)
}

// The ports of each part of a cluster.
var (
	AgentPort    = Single(7443, "control plane: the agents' stream")
	NodeExporter = Single(9100, "node exporter: the machine's metrics, for Prometheus")
	MirrorFiles  = Single(8081, "mirror: weights, for the nodes")
	MirrorAPI    = Single(8082, "mirror: its management API, for the control planes")
	Postgres     = Single(5432, "store: PostgreSQL")
	PatroniAPI   = Single(8008, "store: Patroni's API")
	EtcdClient   = Single(2379, "store: etcd, for Patroni")
	EtcdPeer     = Single(2380, "store: etcd, between its members")
)

// Engines is the engines' port range on a node: their HTTP, for the
// router, and their rendezvous, for the other nodes of a replica.
func Engines(from, to int) Port {
	return Port{From: from, To: to, What: "engines: HTTP and rendezvous"}
}

// Spec is what a machine's ruleset holds: its llmgw ports, and the
// networks they answer besides the machine itself.
type Spec struct {
	Ports []Port
	Allow []netip.Prefix
}

// Table is the nftables table llmgw's rules live in, and nothing else does.
const Table = "llmgw"

// RulesFile is where the ruleset is kept, for the unit that loads it at boot.
const RulesFile = "/etc/llmgw/firewall.nft"

// Unit is the systemd unit that loads it.
const Unit = "llmgw-firewall.service"

// Script is the ruleset as nft reads it. It replaces llmgw's table whole,
// so applying it again, or a changed one, leaves no stale rule behind.
func Script(s Spec) string {
	ports := slices.Clone(s.Ports)
	slices.SortFunc(ports, func(a, b Port) int { return a.From - b.From })
	ports = slices.CompactFunc(ports, func(a, b Port) bool { return a.From == b.From && a.To == b.To })
	var v4, v6 []string
	for _, p := range s.Allow {
		p = p.Masked()
		if p.Addr().Is4() {
			v4 = append(v4, p.String())
		} else {
			v6 = append(v6, p.String())
		}
	}
	slices.Sort(v4)
	v4 = slices.Compact(v4)
	slices.Sort(v6)
	v6 = slices.Compact(v6)

	var list []string
	for _, p := range ports {
		list = append(list, p.nft())
	}
	set := "{ " + strings.Join(list, ", ") + " }"

	var b strings.Builder
	b.WriteString("# Written by llmgw. llmgw's internal ports answer this machine and the\n")
	b.WriteString("# cluster's networks only; nothing else on the machine is touched.\n")
	for _, p := range ports {
		fmt.Fprintf(&b, "#   %-11s %s\n", p.nft(), p.What)
	}
	fmt.Fprintf(&b, "table inet %s\ndelete table inet %s\n", Table, Table)
	fmt.Fprintf(&b, "table inet %s {\n", Table)
	b.WriteString("\tchain input {\n\t\ttype filter hook input priority filter - 1; policy accept;\n")
	b.WriteString("\t\tiif \"lo\" accept\n")
	if len(ports) > 0 {
		if len(v4) > 0 {
			fmt.Fprintf(&b, "\t\tip saddr { %s } tcp dport %s accept\n", strings.Join(v4, ", "), set)
		}
		if len(v6) > 0 {
			fmt.Fprintf(&b, "\t\tip6 saddr { %s } tcp dport %s accept\n", strings.Join(v6, ", "), set)
		}
		fmt.Fprintf(&b, "\t\ttcp dport %s drop\n", set)
	}
	b.WriteString("\t}\n}\n")
	return b.String()
}

// unitFile loads the ruleset at boot, and again whenever nftables.service
// is (re)started — whose own configuration commonly begins by flushing every
// table, ours included.
const unitFile = `[Unit]
Description=llmgw firewall: llmgw's internal ports for the cluster's networks only
After=nftables.service
PartOf=nftables.service
Before=network-online.target llmgw-server.service llmgw-agent.service

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/sbin/nft -f ` + RulesFile + `
ExecReload=/usr/sbin/nft -f ` + RulesFile + `
ExecStop=/usr/sbin/nft delete table inet ` + Table + `

[Install]
WantedBy=multi-user.target nftables.service
`

// Install writes the ruleset and its unit on a machine and applies them,
// installing nftables first where it is missing. It reports whether
// anything changed. The ruleset is checked by nft before it is put in
// place, and a broken one changes nothing.
//
// What the machine's ruleset already holds is kept: a machine can be more
// than one thing — a store member and a node — and whichever of them is set
// up second must not take the first one's ports out of the rules.
func Install(ctx context.Context, r Runner, s Spec) (bool, error) {
	if old, err := r.Run(ctx, "sudo cat "+RulesFile+" 2>/dev/null || true", nil); err == nil {
		prev := parse(string(old))
		s.Ports = append(prev.Ports, s.Ports...)
		s.Allow = append(prev.Allow, s.Allow...)
	}
	if len(s.Allow) == 0 {
		// Ports for this machine alone would cut every node and control
		// plane off. Not a firewall anybody asked for.
		return false, errors.New("no network to let in; the ports were left open rather than closed to every other machine")
	}
	rules := Script(s)
	body := rules + "\x00" + unitFile
	out, err := r.Run(ctx, `set -e
if ! command -v nft >/dev/null 2>&1 && [ ! -x /usr/sbin/nft ]; then
  if command -v apt-get >/dev/null 2>&1; then
    sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq nftables >/dev/null
  elif command -v dnf >/dev/null 2>&1; then
    sudo dnf install -y -q nftables >/dev/null
  else
    echo "no nft, and no apt-get or dnf to install it with" >&2; exit 1
  fi
fi
t=$(mktemp -d); trap 'rm -rf "$t"' EXIT
cat > "$t/all"
python3 - "$t" <<'PY'
import sys
d = sys.argv[1]
rules, unit = open(d + "/all", "rb").read().split(b"\x00")
open(d + "/rules", "wb").write(rules)
open(d + "/unit", "wb").write(unit)
PY
sudo /usr/sbin/nft -c -f "$t/rules"
changed=no
if ! sudo cmp -s "$t/rules" `+RulesFile+`; then
  sudo install -d -m 755 /etc/llmgw
  sudo install -m 644 "$t/rules" `+RulesFile+`
  changed=yes
fi
if ! sudo cmp -s "$t/unit" /etc/systemd/system/`+Unit+`; then
  sudo install -m 644 "$t/unit" /etc/systemd/system/`+Unit+`
  sudo systemctl daemon-reload
  changed=yes
fi
sudo systemctl enable `+Unit+` >/dev/null 2>&1
if [ "$changed" = yes ] || ! sudo /usr/sbin/nft list table inet `+Table+` >/dev/null 2>&1; then
  sudo /usr/sbin/nft -f `+RulesFile+`
  sudo systemctl start `+Unit+` >/dev/null 2>&1 || true
  echo changed
else
  echo same
fi`, []byte(body))
	if err != nil {
		return false, fmt.Errorf("%v: %s", err, strings.TrimSpace(string(out)))
	}
	return strings.HasSuffix(strings.TrimSpace(string(out)), "changed"), nil
}

// Networks are the networks a machine reaches others on: for each address,
// the prefix of the interface its route leaves by — or the address alone,
// when the route goes through a gateway and the network behind it is not
// the machine's to know.
func Networks(ctx context.Context, r Runner, toward []string) ([]netip.Prefix, error) {
	var out []netip.Prefix
	for _, t := range toward {
		a, err := netip.ParseAddr(t)
		if err != nil {
			continue
		}
		b, err := r.Run(ctx, "ip -o route get "+a.String()+" && ip -o addr show", nil)
		if err != nil {
			return nil, fmt.Errorf("finding the network toward %s: %v: %s", a, err, strings.TrimSpace(string(b)))
		}
		p, ok := networkToward(string(b), a)
		if !ok {
			return nil, fmt.Errorf("cannot tell the network toward %s from: %s", a, strings.TrimSpace(string(b)))
		}
		out = append(out, p)
	}
	return out, nil
}

// networkToward reads `ip -o route get A` followed by `ip -o addr show`.
func networkToward(text string, a netip.Addr) (netip.Prefix, bool) {
	lines := strings.Split(strings.TrimSpace(text), "\n")
	if len(lines) == 0 {
		return netip.Prefix{}, false
	}
	route := strings.Fields(lines[0])
	field := func(name string) string {
		for i := 0; i+1 < len(route); i++ {
			if route[i] == name {
				return route[i+1]
			}
		}
		return ""
	}
	if field("via") != "" {
		// Routed: the address alone.
		return netip.PrefixFrom(a, a.BitLen()), true
	}
	dev, src := field("dev"), field("src")
	if route[0] == "local" {
		// This machine's own address: the network of the interface it is on.
		dev, src = "", a.String()
	}
	for _, l := range lines[1:] {
		f := strings.Fields(l)
		// "2: eth0    inet 10.0.0.147/24 brd ..."
		if len(f) < 4 || (dev != "" && strings.TrimSuffix(f[1], ":") != dev) || (f[2] != "inet" && f[2] != "inet6") {
			continue
		}
		p, err := netip.ParsePrefix(f[3])
		if err != nil || (src != "" && p.Addr().String() != src) {
			continue
		}
		if p.Contains(a) {
			return p.Masked(), true
		}
	}
	return netip.PrefixFrom(a, a.BitLen()), true
}

// parse reads back a ruleset Script wrote: its ports from the comments
// naming them, its networks from the rules.
func parse(text string) Spec {
	var s Spec
	for _, l := range strings.Split(text, "\n") {
		if rest, ok := strings.CutPrefix(l, "#   "); ok {
			f := strings.Fields(rest)
			if len(f) == 0 {
				continue
			}
			from, to, isRange := strings.Cut(f[0], "-")
			a, err1 := strconv.Atoi(from)
			b := a
			var err2 error
			if isRange {
				b, err2 = strconv.Atoi(to)
			}
			if err1 == nil && err2 == nil {
				s.Ports = append(s.Ports, Port{From: a, To: b, What: strings.Join(f[1:], " ")})
			}
			continue
		}
		l = strings.TrimSpace(l)
		if !strings.HasPrefix(l, "ip saddr {") && !strings.HasPrefix(l, "ip6 saddr {") {
			continue
		}
		inner := l[strings.Index(l, "{")+1 : strings.Index(l, "}")]
		for _, n := range strings.Split(inner, ",") {
			if p, err := netip.ParsePrefix(strings.TrimSpace(n)); err == nil {
				s.Allow = append(s.Allow, p)
			}
		}
	}
	return s
}

// Present reports whether a machine has llmgw's table loaded.
func Present(ctx context.Context, r Runner) bool {
	_, err := r.Run(ctx, "sudo -n /usr/sbin/nft list table inet "+Table+" >/dev/null 2>&1", nil)
	return err == nil
}

// Admit lets addresses through a machine's ruleset when it has one: each
// one's network, as Networks finds it, unless a network already there holds
// it. A machine without llmgw's rules is left without them. What was added
// is returned, nothing when every address was already let in.
func Admit(ctx context.Context, r Runner, addrs []string) ([]netip.Prefix, error) {
	old, err := r.Run(ctx, "sudo cat "+RulesFile+" 2>/dev/null || true", nil)
	if err != nil {
		return nil, err
	}
	prev := parse(string(old))
	if len(prev.Ports) == 0 {
		return nil, nil
	}
	var missing []string
	for _, s := range addrs {
		a, err := netip.ParseAddr(s)
		if err != nil {
			continue
		}
		if !slices.ContainsFunc(prev.Allow, func(p netip.Prefix) bool { return p.Contains(a) }) {
			missing = append(missing, s)
		}
	}
	if len(missing) == 0 {
		return nil, nil
	}
	nets, err := Networks(ctx, r, missing)
	if err != nil {
		return nil, err
	}
	if _, err := Install(ctx, r, Spec{Allow: nets}); err != nil {
		return nil, err
	}
	return nets, nil
}

// Disable takes llmgw's table out of a machine's ruleset and stops it being
// loaded at boot, keeping the rules file: Install, or enabling the unit,
// puts it back as it was. Reports whether there was anything to take out.
func Disable(ctx context.Context, r Runner) (bool, error) {
	had := Present(ctx, r)
	out, err := r.Run(ctx, `sudo systemctl disable --now `+Unit+` >/dev/null 2>&1 || true
if sudo /usr/sbin/nft list table inet `+Table+` >/dev/null 2>&1; then
  sudo /usr/sbin/nft delete table inet `+Table+`
fi`, nil)
	if err != nil {
		return false, fmt.Errorf("%v: %s", err, strings.TrimSpace(string(out)))
	}
	return had, nil
}

// Enable loads the rules file a machine already has and has it loaded at
// boot again, after Disable. Reports false when there is no rules file.
func Enable(ctx context.Context, r Runner) (bool, error) {
	out, err := r.Run(ctx, `[ -f `+RulesFile+` ] || { echo none; exit 0; }
sudo systemctl enable `+Unit+` >/dev/null 2>&1
sudo /usr/sbin/nft -f `+RulesFile+`
sudo systemctl start `+Unit+` >/dev/null 2>&1 || true
echo loaded`, nil)
	if err != nil {
		return false, fmt.Errorf("%v: %s", err, strings.TrimSpace(string(out)))
	}
	return strings.TrimSpace(string(out)) == "loaded", nil
}

// Local runs scripts on this machine. As root, sudo in them runs its command
// directly, for a machine that has no sudo.
type Local struct{}

// Run runs script with sh.
func (Local) Run(ctx context.Context, script string, stdin []byte) ([]byte, error) {
	if os.Geteuid() == 0 {
		script = "sudo() { \"$@\"; }\n" + script
	}
	cmd := exec.CommandContext(ctx, "sh", "-c", script)
	if stdin != nil {
		cmd.Stdin = bytes.NewReader(stdin)
	}
	var out bytes.Buffer
	cmd.Stdout, cmd.Stderr = &out, &out
	err := cmd.Run()
	return out.Bytes(), err
}