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, a transport
// protocol their native config does not model (SCTP and the portless IP protocols
// GRE, ESP and AH), or an address-set reference (@set) — csf.allow/apf trust files
// take literal addresses only, so a `-m set --match-set` match lives in the hook
// beside the ipset commands that create the set.
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) ||
		isSetRef(r.Source) || isSetRef(r.Destination)
}

// ipv6Unavailable reports whether r resolves to IPv6 while the backend's own IPv6
// handling is switched off (csf.conf IPV6, conf.apf USE_IPV6). Neither backend can
// enforce such a rule: its native config silently drops or ignores every IPv6 form,
// and the raw-iptables hook is no escape hatch — with IPv6 off, neither firewall
// flushes ip6tables on (re)load (csf.pl guards the whole v6 flush behind IPV6; apf's
// ipt6() is a no-op unless USE_IPV6=1), so a hook-injected ip6tables line is
// re-appended on every reload and a line the library removes from the hook lives on
// in the kernel. Both backends therefore reject a concrete-IPv6 rule outright rather
// than write one they cannot manage. A FamilyAny rule is not concrete IPv6 and is
// unaffected: it is written for whichever family the backend enforces. Shared by CSF
// and APF.
func ipv6Unavailable(ipv6Enabled bool, r *Rule) bool {
	return !ipv6Enabled && r.impliedFamily() == IPv6
}

// 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
	}
	// A set reference (@set) is not a literal host: it matches through the hook's
	// `-m set` clause (ruleNeedsHook routes it there), never a plain trust-file line.
	if isSetRef(r.Source) || isSetRef(r.Destination) {
		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 a transport has no portless form: the plain
	// line is all-protocol and the advanced rule requires a port. TCPUDP counts —
	// it names transports, so it is not the all-protocol plain line either.
	if r.Source != "" || r.Destination != "" {
		return onProtocolAxis(r.Proto)
	}
	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.
// iptables has no both-transports match, so a TCPUDP rule fans out into a tcp line
// and a udp line; every other protocol writes one line. A portless ProtocolAny rule
// is a valid protocol-agnostic match (a bare `-j ACCEPT`) and is not fanned.
func hookRuleProtos(r *Rule) []Protocol {
	protos := make([]Protocol, 0, 2)
	for _, sub := range expandProtocols(r) {
		protos = append(protos, sub.Proto)
	}
	return protos
}

// 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)
}

// --- address sets (ipset commands in the hook) -----------------------------
//
// CSF and APF have no native address-set construct, so the library persists a
// set as `ipset` commands in the same hook that carries its raw iptables rules.
// The firewall sources the hook on every (re)start, so the ipset commands
// recreate the set before the `-m set --match-set` rule lines that follow can
// reference it — the set survives a reboot exactly as the hook's rules do. Every
// ipset line is kept ahead of every iptables/ip6tables line to preserve that
// ordering. Reading foreign, user-authored ipset lines is intended, as with
// rules: the library manages the actual hook state.

// ipsetLinesFor renders the hook lines that (re)create a set and load its
// entries: an idempotent create (-exist, so a reload does not fail on the
// existing set), a flush (so a reload drops entries removed since the last
// write, making the entry list declarative), then one add per entry.
func ipsetLinesFor(set *AddressSet) []string {
	fam := "inet"
	if set.Family == IPv6 {
		fam = "inet6"
	}
	lines := []string{
		fmt.Sprintf("ipset create %s %s family %s -exist", set.Name, set.Type.String(), fam),
		fmt.Sprintf("ipset flush %s", set.Name),
	}
	for _, e := range set.Entries {
		lines = append(lines, fmt.Sprintf("ipset add %s %s", set.Name, e))
	}
	return lines
}

// hookIPSetName returns the set a hook ipset line operates on, or "" when the
// line is not one of the library's ipset commands. Every such line names the set
// in its third field (`ipset <verb> <name> ...`).
func hookIPSetName(line string) string {
	f := strings.Fields(line)
	if len(f) >= 3 && f[0] == "ipset" {
		return f[2]
	}
	return ""
}

// isHookRuleLine reports whether a hook line is an iptables/ip6tables command, as
// opposed to an ipset command or user-authored shell.
func isHookRuleLine(line string) bool {
	t := strings.TrimSpace(line)
	return strings.HasPrefix(t, "iptables ") || strings.HasPrefix(t, "ip6tables ")
}

// setInUse reports whether any hook rule line references name through an
// `-m set --match-set <name>` match, so a set is not removed out from under a
// rule that still uses it (the kernel enforces the same on a live destroy).
func setInUse(lines []string, name string) bool {
	for _, l := range lines {
		if !isHookRuleLine(l) {
			continue
		}
		f := strings.Fields(l)
		for i := 0; i+1 < len(f); i++ {
			if f[i] == "--match-set" && f[i+1] == name {
				return true
			}
		}
	}
	return false
}

// getAddressSets parses the sets the hook carries, in the order their create
// lines appear. An ipset is pinned to a single family, so each create yields one
// set and its add lines supply the entries; flush lines carry no state and are
// ignored.
func (h *hookScript) getAddressSets() ([]*AddressSet, error) {
	lines, _, err := h.readHookLines()
	if err != nil {
		return nil, err
	}
	// ipsetParseType is an IPTables method that ignores its receiver; a zero value
	// reuses the same create-line parser the iptables backend uses.
	ipt := &IPTables{}
	sets := map[string]*AddressSet{}
	var order []string
	for _, line := range lines {
		f := strings.Fields(line)
		if len(f) >= 4 && f[0] == "ipset" && f[1] == "create" {
			// ipsetParseType scans a `create NAME <type> family <fam> ...` slice from
			// its third element, so drop the leading `ipset` word to line it up.
			fam, typ := ipt.ipsetParseType(f[1:])
			sets[f[2]] = &AddressSet{Name: f[2], Family: fam, Type: typ}
			order = append(order, f[2])
		}
	}
	for _, line := range lines {
		f := strings.Fields(line)
		if len(f) == 4 && f[0] == "ipset" && f[1] == "add" {
			if s, ok := sets[f[2]]; ok {
				s.Entries = append(s.Entries, f[3])
			}
		}
	}
	out := make([]*AddressSet, 0, len(order))
	for _, n := range order {
		out = append(out, sets[n])
	}
	return out, nil
}

// editAddressSet writes or removes a set's ipset lines in the hook. Adding drops
// any prior lines for the set and reinserts its block ahead of the first
// iptables/ip6tables line, so the set exists before any rule matches it; the
// write is idempotent. Removing drops the set's lines but refuses when a hook
// rule still references it. Every other hook line — user shell, rules, other
// sets — is preserved; it reports whether the hook changed.
func (h *hookScript) editAddressSet(set *AddressSet, remove bool) (bool, error) {
	lines, existed, err := h.readHookLines()
	if err != nil {
		return false, err
	}

	if remove && setInUse(lines, set.Name) {
		return false, fmt.Errorf("address set %q is in use by a rule", set.Name)
	}

	// Drop any existing lines for this set (idempotent re-add; also the removal path).
	kept := make([]string, 0, len(lines))
	dropped := false
	for _, l := range lines {
		if hookIPSetName(l) == set.Name {
			dropped = true
			continue
		}
		kept = append(kept, l)
	}

	if remove {
		if !dropped {
			return false, nil
		}
		return true, h.writeHook(kept, existed)
	}

	// Insert the set's block ahead of the first rule line (or at the end when the
	// hook has none yet), keeping every ipset line before every rule line.
	block := ipsetLinesFor(set)
	next := make([]string, 0, len(kept)+len(block))
	inserted := false
	for _, l := range kept {
		if !inserted && isHookRuleLine(l) {
			next = append(next, block...)
			inserted = true
		}
		next = append(next, l)
	}
	if !inserted {
		next = append(next, block...)
	}

	if equalLines(lines, next) {
		return false, nil
	}
	return true, h.writeHook(next, existed)
}

// editAddressSetEntry adds or removes a single entry in an existing set by
// rewriting the set's block. The set must already exist in the hook.
func (h *hookScript) editAddressSetEntry(name, entry string, remove bool) (bool, error) {
	sets, err := h.getAddressSets()
	if err != nil {
		return false, err
	}
	var target *AddressSet
	for _, s := range sets {
		if s.Name == name {
			target = s
			break
		}
	}
	if target == nil {
		return false, fmt.Errorf("address set %q not found", name)
	}
	if remove {
		next := target.Entries[:0]
		found := false
		for _, e := range target.Entries {
			if e == entry {
				found = true
				continue
			}
			next = append(next, e)
		}
		if !found {
			return false, nil
		}
		target.Entries = next
	} else {
		for _, e := range target.Entries {
			if e == entry {
				return false, nil
			}
		}
		target.Entries = append(target.Entries, entry)
	}
	return h.editAddressSet(target, false)
}

// equalLines reports whether two hook line slices are identical.
func equalLines(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}