package firewall

import (
	"fmt"
	"os"
	"strings"

	"github.com/anmitsu/go-shlex"
)

// hookScript injects raw iptables/ip6tables commands into a CSF or APF hook that
// the firewall sources at (re)load time. It lets those backends express filter
// rules their native config cannot — connection-state, per-rule interface,
// logging, rate limiting, ICMPv6, and the transport protocols SCTP, GRE, ESP and
// AH — by reusing the iptables rule marshaller/parser and writing the resulting
// commands directly into the firewall's documented hook.
//
// The command lines live in the hook file itself (hookPath) rather than in a
// separate library-owned script: edit touches only the exact command lines it
// manages, leaving any user-authored hook content in place, and getRules parses
// every iptables/ip6tables line the hook carries. Reading foreign, user-authored
// hook rules is intended — the library manages the actual firewall state, and
// HasPrefix (derived from the rule's comment tag) is the only signal of what it
// created.
type hookScript struct {
	// rulePrefix tags each injected rule with an iptables comment so it can be
	// told apart from other rules.
	rulePrefix string
	// hookPath is the firewall hook this library writes its command lines into. It
	// runs before the firewall adds its own rules, so injected rules sit at the top
	// of the INPUT/OUTPUT chains.
	hookPath string
	// hookPerm is the mode the hook file must carry to be executed (0700 for CSF,
	// 0750 for APF).
	hookPerm os.FileMode
}

// ruleNeedsHook reports whether a rule requires a feature that CSF/APF cannot
// express in their native config and so must be injected as a raw iptables rule
// through the hook: a forward-chain (routed) rule, connection-state matching,
// per-rule interface matching, logging, rate limiting, ICMPv6, or a transport
// protocol their native config does not model (SCTP and the portless IP protocols
// GRE, ESP and AH).
func ruleNeedsHook(r *Rule) bool {
	return r.IsForward() || r.State != 0 || r.InInterface != "" || r.OutInterface != "" ||
		r.Log || r.RateLimit != nil || r.Proto == ICMPv6 || hookOnlyProto(r.Proto)
}

// bareHostShape reports whether a rule has the shape a plain csf.allow/apf
// allow_hosts line expresses: exactly one source or destination address, no ports,
// and the any-protocol match. Its direction is not considered — a DirAny bare host
// is the single bidirectional plain line, while a concrete-direction one is one-way
// (see bareHostOneWay). Shared by CSF and APF.
func bareHostShape(r *Rule) bool {
	if r.HasPorts() || r.HasSourcePorts() || r.Proto != ProtocolAny {
		return false
	}
	return (r.Source != "") != (r.Destination != "")
}

// bareHostOneWay reports whether a rule is a ONE-WAY bare-address host allow/deny:
// the bare host shape with a concrete input or output direction. A plain line matches
// a host in BOTH directions, and neither backend's advanced-rule format can carry an
// address without a port, so a one-way bare host rule is expressed through the
// raw-iptables hook instead.
func bareHostOneWay(r *Rule) bool {
	return bareHostShape(r) && (r.Direction == DirInput || r.Direction == DirOutput)
}

// dirAnyPlainLine reports whether a DirAny rule maps to a single bidirectional plain
// csf.allow/apf line: a bare host carrying no feature that would force the
// raw-iptables hook (connection state, interface, logging, etc.). Every other DirAny
// rule fans out into a concrete input rule plus its role-swapped output rule on
// add/remove, since csf/apf have no single native both-directions construct for it.
func dirAnyPlainLine(r *Rule) bool {
	return r.Direction == DirAny && bareHostShape(r) && !ruleNeedsHook(r)
}

// hostNeedsHook reports whether an address-bearing rule with no port must be
// injected through the raw-iptables hook because a csf/apf trust file has no form
// for it: a source+destination pair (both store a single address) or a host
// pinned to a concrete tcp/udp protocol (the plain line is all-protocol and the
// advanced rule requires a port). Both are expressed natively as an iptables
// rule. A single-address all-protocol host is not covered here — a one-way one is
// routed by bareHostOneWay, a bidirectional one is a native plain line — and ICMP
// keeps its own handling (checkICMP/apfCheckICMP), so it is excluded. Shared by
// CSF and APF.
func hostNeedsHook(r *Rule) bool {
	if r.Proto.IsICMP() || r.HasPorts() || r.HasSourcePorts() {
		return false
	}
	// A source+destination pair has no single-address advanced/plain form.
	if r.Source != "" && r.Destination != "" {
		return true
	}
	// A single-address host pinned to tcp/udp has no portless form.
	if r.Source != "" || r.Destination != "" {
		return r.Proto == TCP || r.Proto == UDP
	}
	return false
}

// bareProtoNeedsHook reports whether a rule is a bare protocol match — a non-ICMP
// transport with no address and no port — that CSF/APF cannot express in their
// native config (the trust files key on an address and the conf lists key on a
// port or icmp type) but iptables applies directly (`-p tcp -j ACCEPT`, or a bare
// `-j ACCEPT`). Such a rule is injected through the raw-iptables hook rather than
// rejected. ICMP/ICMPv6 keep their own native/hook handling (checkICMP/apfCheckICMP,
// nativeICMPv6, ruleNeedsHook), so they are excluded here. Shared by CSF and APF.
func bareProtoNeedsHook(r *Rule) bool {
	return r.Source == "" && r.Destination == "" && !r.HasPorts() && !r.HasSourcePorts() && !r.Proto.IsICMP()
}

// hookOnlyProto reports whether a protocol has no representation in CSF's or
// APF's native config and so can only be applied through the raw-iptables hook.
func hookOnlyProto(p Protocol) bool {
	switch p {
	case SCTP, GRE, ESP, AH:
		return true
	}
	return false
}

// hookRuleProtos lists the transport protocols a rule is written for in the hook.
// A ProtocolAny rule that carries a port has no single iptables form — a --dport/
// --sport match requires a concrete -p tcp/udp — so it fans out into a tcp line and
// a udp line, mirroring the tcp+udp fan-out csf/apf write for a ProtocolAny port
// rule in their native config. A portless ProtocolAny rule is a valid protocol-
// agnostic match (a bare `-j ACCEPT`), so it keeps its own protocol and is not fanned.
func hookRuleProtos(r *Rule) []Protocol {
	if r.Proto == ProtocolAny && (r.HasPorts() || r.HasSourcePorts()) {
		return []Protocol{TCP, UDP}
	}
	return []Protocol{r.Proto}
}

// hookRuleFamilies lists the address families a rule is written for: a rule
// pinned to a family (by address or an ICMP protocol) touches only that command,
// a family-agnostic rule (e.g. a bare state match) is written for both v4 and v6.
func hookRuleFamilies(r *Rule) []Family {
	switch r.impliedFamily() {
	case IPv4:
		return []Family{IPv4}
	case IPv6:
		return []Family{IPv6}
	default:
		return []Family{IPv4, IPv6}
	}
}

// hookCommand returns the iptables command name for a family.
func hookCommand(fam Family) string {
	if fam == IPv6 {
		return "ip6tables"
	}
	return "iptables"
}

// shellSafeToken quotes a token so /bin/sh passes it through verbatim. The
// iptables marshaller quotes free-text fields (a comment, a log prefix) with
// strconv.Quote — Go double-quoting, which is right for an iptables-restore file
// but NOT for the hook script, which /bin/sh sources: inside double quotes the
// shell still expands $var, $(...) and backticks. A token made of ordinary
// argument characters is returned bare for readability; anything else is wrapped
// in single quotes (with any embedded single quote escaped), which the shell
// treats as a literal. shlex.Split reverses either form on read-back.
func shellSafeToken(tok string) string {
	safe := tok != ""
	for _, r := range tok {
		if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' ||
			strings.ContainsRune("_./:=,+-@%", r) {
			continue
		}
		safe = false
		break
	}
	if safe {
		return tok
	}
	return "'" + strings.ReplaceAll(tok, "'", `'\''`) + "'"
}

// rulesToLines encodes a rule as the raw command line(s) to inject: one iptables
// (or ip6tables) command per underlying iptables line and per family. A logged
// rule yields a LOG line followed by its action line, as with the iptables
// backend. Each marshalled line is re-tokenized and re-quoted shell-safely,
// because the hook script is sourced by /bin/sh rather than exec'd argv-style.
func (h *hookScript) rulesToLines(r *Rule) ([]string, error) {
	var out []string
	for _, proto := range hookRuleProtos(r) {
		for _, fam := range hookRuleFamilies(r) {
			rc := *r
			rc.Proto = proto
			rc.Family = fam
			ipt := &IPTables{rulePrefix: h.rulePrefix}
			base, err := ipt.marshalRuleLines(&rc)
			if err != nil {
				return nil, err
			}
			cmd := hookCommand(fam)
			for _, line := range base {
				tokens, terr := shlex.Split(line, true)
				if terr != nil {
					return nil, terr
				}
				for i, t := range tokens {
					tokens[i] = shellSafeToken(t)
				}
				out = append(out, cmd+" "+strings.Join(tokens, " "))
			}
		}
	}
	return out, nil
}

// ruleMatchesAny reports whether e — a rule parsed from a hook line — is the same
// underlying rule as any of targets, honoring direction (the hook emits explicit
// -A INPUT/-A OUTPUT lines) but ignoring the comment, which is not part of rule
// identity. It backs comment-agnostic hook removal.
func ruleMatchesAny(e *Rule, targets []*Rule) bool {
	for _, t := range targets {
		if e.Equal(t, true) {
			return true
		}
	}
	return false
}

// parseLine decodes an injected command line back into the rule it represents
// (one line, so a LOG line yields a rule with Log set and no action), reporting
// whether the line is one this backend recognizes.
func (h *hookScript) parseLine(line string) (*Rule, bool) {
	line = strings.TrimSpace(line)
	var fam Family
	var rest string
	switch {
	case strings.HasPrefix(line, "iptables "):
		fam, rest = IPv4, strings.TrimPrefix(line, "iptables ")
	case strings.HasPrefix(line, "ip6tables "):
		fam, rest = IPv6, strings.TrimPrefix(line, "ip6tables ")
	default:
		return nil, false
	}
	r, err := unmarshalIPTablesRule(rest, fam)
	if err != nil {
		return nil, false
	}
	// The iptables parser captures the prefixed comment; strip the prefix (a tag)
	// so only the user-facing comment surfaces, and record whether the injected
	// rule carried the prefix so HasPrefix reflects it just like the iptables
	// backend.
	text, hasPrefix := prefixedComment(h.rulePrefix, r.Comment)
	r.Comment = text
	r.HasPrefix = hasPrefix
	return r, true
}

// commandLines returns the iptables/ip6tables command lines currently in the
// hook (a missing hook contributes none). Every such line is returned, including
// any a user authored by hand, so the library reconciles the hook's real state.
func (h *hookScript) commandLines() ([]string, error) {
	lines, _, err := h.readHookLines()
	if err != nil {
		return nil, err
	}
	var cmds []string
	for _, line := range lines {
		t := strings.TrimSpace(line)
		if strings.HasPrefix(t, "iptables ") || strings.HasPrefix(t, "ip6tables ") {
			cmds = append(cmds, t)
		}
	}
	return cmds, nil
}

// getRules parses the hook into logical rules, coalescing each LOG line with the
// action line that follows it. Family merging is left to the caller, which
// unions these with the backend's native rules.
func (h *hookScript) getRules() ([]*Rule, error) {
	cmds, err := h.commandLines()
	if err != nil {
		return nil, err
	}
	var rules []*Rule
	for _, line := range cmds {
		if r, ok := h.parseLine(line); ok {
			rules = append(rules, r)
		}
	}
	return coalesceLoggedRules(rules), nil
}

// readHookLines returns the hook's lines and whether the hook file exists. A
// single trailing newline is trimmed so a rewrite does not accrue a blank line;
// a missing hook yields no lines.
func (h *hookScript) readHookLines() ([]string, bool, error) {
	data, err := os.ReadFile(h.hookPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, false, nil
		}
		return nil, false, err
	}
	content := strings.TrimSuffix(string(data), "\n")
	if content == "" {
		return nil, true, nil
	}
	return strings.Split(content, "\n"), true, nil
}

// writeHook atomically replaces the hook with lines, giving a freshly created
// hook a shebang, and keeps it executable so the firewall can source it.
func (h *hookScript) writeHook(lines []string, existed bool) error {
	var b strings.Builder
	if !existed {
		b.WriteString("#!/bin/sh\n")
	}
	for _, l := range lines {
		b.WriteString(l)
		b.WriteByte('\n')
	}
	// Preserve an existing hook's mode and ownership; a freshly created hook gets
	// the executable hookPerm so the firewall can source it.
	if err := writeConfigFile(h.hookPath, []byte(b.String()), h.hookPerm); err != nil {
		return fmt.Errorf("failed to move updated hook into place: %s", err)
	}
	return nil
}

// edit adds or removes a rule's command line(s) directly in the hook. Adding
// matches on the exact deterministic line the marshaller emits, so the library's
// own tagged line is written once and re-adds are idempotent. Removal instead
// matches on the underlying rule: a line is dropped when the rule it encodes is
// the same as one r resolves to, ignoring the comment tag, so a copy of the rule a
// customer added under a different comment (or none) is cleared too — the comment
// is not part of rule identity. A logged rule's LOG and action lines are matched
// independently, exactly as they were written. It preserves every other hook
// line — user-authored shell and rules alike — and reports whether the hook
// changed. Adding to an absent hook creates it; removing from one is a no-op.
func (h *hookScript) edit(r *Rule, remove bool) (bool, error) {
	desired, err := h.rulesToLines(r)
	if err != nil {
		return false, err
	}

	lines, existed, err := h.readHookLines()
	if err != nil {
		return false, err
	}

	changed := false
	if remove {
		// Parse each desired line back into the rule it encodes; a hook line whose own
		// rule matches one of these (comment ignored, see Rule.Equal) is dropped. A line
		// that is not a rule the hook recognizes (foreign shell, a comment) never parses
		// and is preserved.
		var targets []*Rule
		for _, l := range desired {
			if tr, ok := h.parseLine(l); ok {
				targets = append(targets, tr)
			}
		}
		next := lines[:0:0]
		for _, l := range lines {
			if er, ok := h.parseLine(l); ok && ruleMatchesAny(er, targets) {
				changed = true
				continue
			}
			next = append(next, l)
		}
		if !changed {
			return false, nil
		}
		lines = next
	} else {
		present := make(map[string]bool, len(lines))
		for _, l := range lines {
			present[strings.TrimSpace(l)] = true
		}
		for _, l := range desired {
			if !present[l] {
				lines = append(lines, l)
				present[l] = true
				changed = true
			}
		}
		if !changed {
			return false, nil
		}
	}

	return true, h.writeHook(lines, existed)
}