// Package firewall speaks the firewall found on an adopted node, in that firewall's own terms
// (novox/hq ADR 0100).
//
// **The found firewall stays in force.** On an adopted node the mesh loads nothing that drops by
// default or holds an accept; what it needs reachable it converges as openings through what it
// found, marks each rule as its own, and removes only what it marked. It never resets or flushes:
// the rules the machine already had are the operator's, and they are what keeps it serving.
package firewall

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	"github.com/novox/mesh-host/internal/declaration"
	"github.com/novox/mesh-host/internal/system"
)

// Runner executes a command.
type Runner = system.Runner

// Kind is what firewall a machine has, as far as the mesh is concerned.
type Kind string

const (
	// UFW is an active ufw — the one kind found on the machines measured, and the one spoken.
	UFW Kind = "ufw"
	// None is a machine where nothing refuses anything, which needs no openings.
	None Kind = "none"
	// Unsupported is a firewall no host speaks yet. A machine with one is refused adoption: the
	// mesh could neither open what it needs nor know what it would be closing.
	Unsupported Kind = "unsupported"
)

// deletion is the arguments that delete a rule as `ufw show added` printed it. A route rule is
// deleted with `route delete …`: ufw refuses `delete route …` as invalid syntax. And ufw answers
// success when asked to delete a rule it does not hold, so every deletion is read back.
func deletion(rule string) []string {
	w := words(rule)
	if len(w) > 0 && w[0] == "route" {
		return append([]string{"route", "delete"}, w[1:]...)
	}
	return append([]string{"delete"}, w...)
}

// MeshInterface is the private network's interface, the way an opening from the mesh is known.
// It must be the controller's overlay interface name.
const MeshInterface = "mesh0"

// Detect says which firewall this machine has. For Unsupported the string names it.
func Detect(ctx context.Context, run Runner) (Kind, string, error) {
	if out, err := run(ctx, "firewall-cmd", "--state"); err == nil && strings.TrimSpace(out) == "running" {
		return Unsupported, "firewalld", nil
	}
	ufwActive := false
	if out, err := run(ctx, "ufw", "status"); err == nil {
		ufwActive = statusActive(out)
	}

	out, err := run(ctx, "nft", "list", "ruleset")
	switch {
	case err == nil:
		if refusing := Refusing(out, ufwActive); len(refusing) > 0 {
			return Unsupported, "nftables rules that refuse traffic, in " + strings.Join(refusing, ", "), nil
		}
	case !missing(err):
		return "", "", fmt.Errorf("cannot read this machine's packet filter to know what it has: %w", err)
	}

	if !ufwActive {
		// iptables with the legacy backend is invisible to nft.
		for _, legacy := range []string{"iptables-legacy", "ip6tables-legacy"} {
			out, err := run(ctx, legacy, "-S")
			if err != nil {
				continue
			}
			if refusing := RefusingLegacy(out); len(refusing) > 0 {
				return Unsupported, legacy + " rules that refuse traffic, in " + strings.Join(refusing, ", "), nil
			}
		}
	}

	if ufwActive {
		return UFW, "ufw", nil
	}
	return None, "", nil
}

func missing(err error) bool {
	return errors.Is(err, exec.ErrNotFound)
}

func statusActive(out string) bool {
	for _, line := range strings.Split(out, "\n") {
		if strings.HasPrefix(strings.TrimSpace(line), "Status:") {
			return strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "Status:")) == "active"
		}
	}
	return false
}

// Refusing names the tables of an `nft list ruleset` holding something that refuses traffic — a
// drop or reject, or a base chain whose policy drops — and that is neither the mesh's own nor the
// container runtime's. With ufw active, the tables iptables-nft manages are ufw's and the runtime's
// and are not counted.
//
// **A ban is not a firewall.** fail2ban refuses the sources it banned and passes everything else;
// captured on a lab machine with both of its backends (testdata/fail2ban-nftables.nft,
// testdata/fail2ban-iptables.nft). The mesh opens nothing through it and it closes nothing the
// mesh needs, so a refusal that names the sources it refuses, in a table or a chain that accepts
// nothing and is entered only from chains whose policy accepts, is not counted.
func Refusing(ruleset string, ufwActive bool) []string {
	type rule struct{ table, chain, line string }
	type chainOf struct {
		base, dropping, accepts bool
		policyLine            string
		jumpedFrom            []string
	}
	chains := map[string]*chainOf{} // by "table\x00chain"
	tableAccepts := map[string]bool{}
	var tables []string
	var refusals []rule
	managed := map[string]bool{}
	var table, chain string
	get := func(t, c string) *chainOf {
		k := t + "\x00" + c
		if chains[k] == nil {
			chains[k] = &chainOf{}
		}
		return chains[k]
	}
	for _, raw := range strings.Split(ruleset, "\n") {
		line := strings.TrimSpace(raw)
		switch {
		case strings.HasPrefix(line, "# Warning: table ") && strings.Contains(line, "managed by iptables-nft"):
			name := strings.TrimPrefix(line, "# Warning: table ")
			name, _, _ = strings.Cut(name, " is managed")
			managed[name] = true
			continue
		case strings.HasPrefix(line, "table "):
			table = strings.TrimSuffix(strings.TrimSpace(strings.TrimPrefix(line, "table ")), "{")
			table = strings.TrimSpace(table)
			tables = append(tables, table)
			chain = ""
			continue
		case strings.HasPrefix(line, "chain "):
			chain = strings.TrimSpace(strings.TrimSuffix(strings.TrimPrefix(line, "chain "), "{"))
			get(table, chain)
			continue
		case strings.HasPrefix(line, "set ") || strings.HasPrefix(line, "map ") ||
			strings.HasPrefix(line, "flowtable "):
			chain = ""
			continue
		case line == "" || line == "}" || strings.HasPrefix(line, "#") || chain == "":
			continue
		}
		c := get(table, chain)
		if strings.HasPrefix(line, "type ") {
			c.base = true
			c.policyLine = line
			c.dropping = strings.Contains(line, "policy drop")
			continue
		}
		for _, verb := range []string{"jump ", "goto "} {
			if i := strings.Index(line, verb); i >= 0 {
				target := strings.Fields(line[i+len(verb):])
				if len(target) > 0 {
					get(table, target[0]).jumpedFrom = append(get(table, target[0]).jumpedFrom, chain)
				}
			}
		}
		if accepts(line) {
			c.accepts = true
			tableAccepts[table] = true
		}
		if verdictRefuses(line) {
			refusals = append(refusals, rule{table, chain, line})
		}
	}

	skipped := func(table string) bool {
		if table == "inet mesh" || table == "inet mesh_guard" {
			return true
		}
		return (managed[table] || iptablesTable(table)) && ufwActive
	}
	// onlyBans is whether a refusal only refuses the sources it names: in a table that accepts
	// nothing and whose base chains all accept by default, or in a chain that accepts nothing and
	// is entered only from base chains that accept by default.
	onlyBans := func(r rule) bool {
		if !bansSources(r.line) {
			return false
		}
		allAccepting := true
		for k, c := range chains {
			if strings.HasPrefix(k, r.table+"\x00") && c.base && !strings.Contains(c.policyLine, "policy accept") {
				allAccepting = false
			}
		}
		if !tableAccepts[r.table] && allAccepting {
			return true
		}
		c := get(r.table, r.chain)
		if c.base || c.accepts || len(c.jumpedFrom) == 0 {
			return false
		}
		for _, from := range c.jumpedFrom {
			caller := get(r.table, from)
			if !caller.base || !strings.Contains(caller.policyLine, "policy accept") {
				return false
			}
		}
		return true
	}

	counted := map[string]bool{}
	for k, c := range chains {
		t, name, _ := strings.Cut(k, "\x00")
		if skipped(t) || !c.dropping {
			continue
		}
		if (managed[t] || iptablesTable(t)) && runtimes(t, name, c.policyLine) {
			continue
		}
		counted[t] = true
	}
	for _, r := range refusals {
		if skipped(r.table) || counted[r.table] {
			continue
		}
		if (managed[r.table] || iptablesTable(r.table)) && runtimes(r.table, r.chain, r.line) {
			continue
		}
		if onlyBans(r) {
			continue
		}
		counted[r.table] = true
	}
	var refusing []string
	for _, t := range tables {
		if counted[t] {
			counted[t] = false
			refusing = append(refusing, "table "+t)
		}
	}
	return refusing
}

// iptablesTable is whether a table is one iptables-nft writes. Named rather than read from the
// warning nft prints above it, because nft does not print that for every such table: a captured
// ruleset carried it on ip filter and not on ip raw, where the runtime keeps its drops.
func iptablesTable(table string) bool {
	family, name, _ := strings.Cut(table, " ")
	if family != "ip" && family != "ip6" {
		return false
	}
	switch name {
	case "filter", "nat", "raw", "mangle", "security":
		return true
	}
	return false
}

// runtimes is whether a refusal in an iptables-nft table is the container runtime's own: in its
// DOCKER chains, its forward policy, or its guard against reaching a container's address directly
// from outside its bridge, in the raw table.
func runtimes(table, chain, line string) bool {
	_, name, _ := strings.Cut(table, " ")
	switch {
	case strings.HasPrefix(chain, "DOCKER"):
		return true
	case name == "filter" && chain == "FORWARD" && strings.HasPrefix(line, "type "):
		return true
	case name == "raw" && chain == "PREROUTING":
		return strings.Contains(line, "daddr") && strings.Contains(line, "iifname !=")
	}
	return false
}

// iptables-nft prints a REJECT target it cannot translate as `xt target "REJECT"`, measured in
// testdata/fail2ban-iptables.nft; a refusal written that way is a refusal too.
var verdict = regexp.MustCompile(`(^|\s)(drop|reject)(\s|$)|xt target "(DROP|REJECT)"`)

func verdictRefuses(line string) bool {
	return verdict.MatchString(line)
}

var acceptVerdict = regexp.MustCompile(`(^|\s)accept(\s|;|$)|xt target "ACCEPT"`)

func accepts(line string) bool {
	return acceptVerdict.MatchString(line)
}

// bansSources is whether a refusal names the sources it refuses — a set or an address — rather
// than refusing everyone but some.
func bansSources(line string) bool {
	f := strings.Fields(line)
	for i, w := range f {
		if (w == "saddr" || w == "-s") && i+1 < len(f) && f[i+1] != "!=" && !strings.HasPrefix(f[i+1], "!") {
			return i == 0 || f[i-1] != "!"
		}
	}
	return false
}

// RefusingLegacy names the chains of an `iptables-legacy -S` that refuse traffic outside the
// container runtime's own. A ban — a refusal of the sources it names, in a chain that accepts
// nothing and is entered only from built-in chains whose policy accepts — is not counted, as in
// Refusing (testdata/fail2ban-iptables-S.txt).
func RefusingLegacy(rules string) []string {
	policy := map[string]string{}
	accepting := map[string]bool{}
	jumpedFrom := map[string][]string{}
	for _, line := range strings.Split(rules, "\n") {
		fields := strings.Fields(line)
		if len(fields) < 3 {
			continue
		}
		switch fields[0] {
		case "-P":
			policy[fields[1]] = fields[2]
		case "-A":
			for i, f := range fields {
				if (f == "-j" || f == "-g") && i+1 < len(fields) {
					switch fields[i+1] {
					case "ACCEPT":
						accepting[fields[1]] = true
					case "DROP", "REJECT", "RETURN", "LOG":
					default:
						jumpedFrom[fields[i+1]] = append(jumpedFrom[fields[i+1]], fields[1])
					}
				}
			}
		}
	}
	ban := func(chain, line string) bool {
		if !bansSources(line) || accepting[chain] || len(jumpedFrom[chain]) == 0 {
			return false
		}
		for _, from := range jumpedFrom[chain] {
			if policy[from] != "ACCEPT" {
				return false
			}
		}
		return true
	}
	var refusing []string
	seen := map[string]bool{}
	for _, line := range strings.Split(rules, "\n") {
		fields := strings.Fields(line)
		if len(fields) < 3 {
			continue
		}
		chain := fields[1]
		refuses := false
		switch fields[0] {
		case "-P":
			refuses = fields[2] == "DROP" && chain != "FORWARD"
		case "-A":
			for i, f := range fields {
				if f == "-j" && i+1 < len(fields) && (fields[i+1] == "DROP" || fields[i+1] == "REJECT") {
					refuses = !strings.HasPrefix(chain, "DOCKER") && !ban(chain, line)
				}
			}
		}
		if refuses && !seen[chain] {
			seen[chain] = true
			refusing = append(refusing, "chain "+chain)
		}
	}
	return refusing
}

// --- ufw ---------------------------------------------------------------------------------------

// Mark is the comment every rule the mesh adds carries: whose it is, which opening, and a digest
// of the rule itself, so a rule the opening no longer describes is recognised as stale without the
// host having to know how ufw prints a rule back.
func Mark(o *declaration.Opening) string {
	sum := sha256.Sum256([]byte(strings.Join(Rule(o), " ")))
	return marker(o.ID) + " " + hex.EncodeToString(sum[:])[:8]
}

func marker(id string) string { return "mesh-host " + id }

// markedFor is whether a comment is the mesh's, for this opening.
func markedFor(comment, id string) bool {
	return comment == marker(id) || strings.HasPrefix(comment, marker(id)+" ")
}

// Rule is the ufw rule an opening becomes, without its comment.
//
//	incoming from everywhere  allow proto tcp to any port P
//	incoming from the mesh    allow in on mesh0 proto tcp to any port P
//	forwarded                 route allow [in on mesh0] proto tcp to any port <container port>
//
// A forwarded opening names the container's port because ufw's route rules are matched after the
// runtime's destination translation.
func Rule(o *declaration.Opening) []string {
	var rule []string
	port := o.Port
	if o.Path == declaration.PathForwarded {
		rule = append(rule, "route")
		port = o.To
	}
	rule = append(rule, "allow")
	if o.From == declaration.FromMesh {
		rule = append(rule, "in", "on", MeshInterface)
	}
	return append(rule, "proto", o.Protocol, "to", "any", "port", strconv.Itoa(port))
}

var commentOf = regexp.MustCompile(`comment '([^']*)'`)

// added is every rule `ufw show added` lists, each without its leading "ufw".
func added(ctx context.Context, run Runner) ([]string, error) {
	out, err := run(ctx, "ufw", "show", "added")
	if err != nil {
		return nil, fmt.Errorf("reading ufw's rules: %w", err)
	}
	var rules []string
	for _, line := range strings.Split(out, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "ufw ") {
			rules = append(rules, strings.TrimPrefix(line, "ufw "))
		}
	}
	return rules, nil
}

func comment(rule string) string {
	m := commentOf.FindStringSubmatch(rule)
	if m == nil {
		return ""
	}
	return m[1]
}

// words splits a rule as ufw printed it into arguments, keeping a quoted comment whole.
func words(rule string) []string {
	var out []string
	var cur strings.Builder
	quoted, any := false, false
	for _, r := range rule {
		switch {
		case r == '\'':
			quoted = !quoted
			any = true
		case r == ' ' && !quoted:
			if any {
				out = append(out, cur.String())
				cur.Reset()
				any = false
			}
		default:
			cur.WriteRune(r)
			any = true
		}
	}
	if any {
		out = append(out, cur.String())
	}
	return out
}

// Converge makes one opening true in ufw: its marked rule present, and any rule marked for it that
// no longer describes it deleted. Nothing unmarked is touched. The outcome is created, updated or
// unchanged, read back from ufw rather than assumed.
func Converge(ctx context.Context, run Runner, o *declaration.Opening) (string, error) {
	rules, err := added(ctx, run)
	if err != nil {
		return "", err
	}
	mark := Mark(o)
	present := false
	var stale []string
	for _, rule := range rules {
		c := comment(rule)
		switch {
		case c == mark:
			present = true
		case markedFor(c, o.ID):
			stale = append(stale, rule)
		}
	}
	if present && len(stale) == 0 {
		return "unchanged", nil
	}
	for _, rule := range stale {
		if _, err := run(ctx, "ufw", deletion(rule)...); err != nil {
			return "", fmt.Errorf("deleting the mesh's stale ufw rule %q: %w", rule, err)
		}
	}
	if !present {
		args := append(Rule(o), "comment", mark)
		if _, err := run(ctx, "ufw", args...); err != nil {
			return "", fmt.Errorf("adding the ufw rule for %s: %w", o.Target(), err)
		}
	}
	after, err := added(ctx, run)
	if err != nil {
		return "", err
	}
	found, leftover := false, 0
	for _, rule := range after {
		c := comment(rule)
		if c == mark {
			found = true
		} else if markedFor(c, o.ID) {
			leftover++
		}
	}
	if !found {
		return "", fmt.Errorf("ufw was asked for %s and does not list it afterwards", o.Target())
	}
	if leftover > 0 {
		return "", fmt.Errorf("ufw still lists %d stale rule(s) marked for %s after deleting them", leftover, o.ID)
	}
	if len(stale) > 0 {
		return "updated", nil
	}
	return "created", nil
}

// Remove deletes the rules marked for one opening, and nothing else.
func Remove(ctx context.Context, run Runner, id string) (int, error) {
	rules, err := added(ctx, run)
	if err != nil {
		return 0, err
	}
	removed := 0
	for _, rule := range rules {
		if !markedFor(comment(rule), id) {
			continue
		}
		if _, err := run(ctx, "ufw", deletion(rule)...); err != nil {
			return removed, fmt.Errorf("deleting the mesh's ufw rule %q: %w", rule, err)
		}
		removed++
	}
	after, err := added(ctx, run)
	if err != nil {
		return removed, err
	}
	for _, rule := range after {
		if markedFor(comment(rule), id) {
			return removed, fmt.Errorf("ufw still lists a rule marked for %s after deleting it", id)
		}
	}
	return removed, nil
}

// Enable turns ufw back on, as found, and reads back that it is.
func Enable(ctx context.Context, run Runner) error {
	if _, err := run(ctx, "ufw", "--force", "enable"); err != nil {
		return fmt.Errorf("enabling ufw again: %w", err)
	}
	return expectActive(ctx, run, true)
}

// Disable retires ufw without flushing it: its configuration stays on disk, and the container
// runtime's rules are not its to remove.
//
// **Nor is the forward policy ufw's to open.** Measured on a lab machine running the container
// runtime with a published port (testdata/ufw-disable-iptables-before.txt and -after.txt):
// `ufw disable` sets every built-in chain's policy to accept, the forward chain's among them. The
// runtime had set that one to drop when it turned forwarding on, and it does not set it again while
// forwarding stays on — not even on a restart. Left so, a retired ufw turns the machine into a
// router for anyone who can reach it. So each family's forward policy is read before, and one that
// was drop is put back and read back.
func Disable(ctx context.Context, run Runner) error {
	type family struct{ tool, policy string }
	var before []family
	for _, tool := range []string{"iptables", "ip6tables"} {
		if policy, ok := forwardPolicy(ctx, run, tool); ok {
			before = append(before, family{tool, policy})
		}
	}
	if _, err := run(ctx, "ufw", "disable"); err != nil {
		return fmt.Errorf("disabling ufw: %w", err)
	}
	if err := expectActive(ctx, run, false); err != nil {
		return err
	}
	for _, f := range before {
		if f.policy != "DROP" {
			continue
		}
		if now, ok := forwardPolicy(ctx, run, f.tool); ok && now == "DROP" {
			continue
		}
		if _, err := run(ctx, f.tool, "-P", "FORWARD", "DROP"); err != nil {
			return fmt.Errorf("ufw is disabled, and %s's forward policy, which was drop, could not be put back: %w",
				f.tool, err)
		}
		if now, ok := forwardPolicy(ctx, run, f.tool); !ok || now != "DROP" {
			return fmt.Errorf("ufw is disabled, and %s's forward policy was put back to drop and reads %q",
				f.tool, now)
		}
	}
	return nil
}

// forwardPolicy reads the forward chain's policy the way iptables prints it: "-P FORWARD DROP".
// Not ok when the tool is absent or says nothing readable.
func forwardPolicy(ctx context.Context, run Runner, tool string) (string, bool) {
	out, err := run(ctx, tool, "-S", "FORWARD")
	if err != nil {
		return "", false
	}
	for _, line := range strings.Split(out, "\n") {
		f := strings.Fields(line)
		if len(f) == 3 && f[0] == "-P" && f[1] == "FORWARD" {
			return f[2], true
		}
	}
	return "", false
}

func expectActive(ctx context.Context, run Runner, want bool) error {
	out, err := run(ctx, "ufw", "status")
	if err != nil {
		return fmt.Errorf("reading ufw's status back: %w", err)
	}
	if statusActive(out) != want {
		state := "inactive"
		if want {
			state = "active"
		}
		return fmt.Errorf("ufw was asked to be %s and says: %s", state, strings.TrimSpace(out))
	}
	return nil
}