#!/bin/bash
# nftables ruleset, its verification and sshd hardening.
#
# Sourced by setup_wireguard_ssh.sh; split out to keep wg_keys.sh under the
# 250-line cap. Sourced rather than run, so it inherits the caller's
# strict mode and the variables defined above the source line.

# Rendering is separate from writing so that 'verify' can regenerate the
# ruleset into a scratch file and diff it against what is actually on disk,
# without backing anything up or touching the live configuration.
render_nftables_ruleset() {
	local target="$1"
	detect_lan_subnet
	local web_rule=""
	if [[ ${ALLOW_WEB:-false} == "true" ]]; then
		web_rule=$'\n\t\ttcp dport { 80, 443 } accept'
	fi
	# DNS blocker (setup_dns_blocker.sh): serve DNS (53) -- and DHCP (67) when
	# this PC is the LAN DHCP server -- to LAN clients only. Restricted to the
	# LAN subnet so nothing is ever exposed to the internet.
	local dns_rule="" dhcp_rule=""
	if [[ ${ALLOW_DNS:-false} == "true" ]]; then
		dns_rule=$'\n\t\tip saddr '"${LAN_SUBNET}"$' udp dport 53 accept'
		dns_rule+=$'\n\t\tip saddr '"${LAN_SUBNET}"$' tcp dport 53 accept'
		# DHCP clients have no IP yet (saddr 0.0.0.0 -> 255.255.255.255) and the
		# broadcast is classified INVALID by conntrack -- so this rule MUST sit
		# BEFORE 'ct state invalid drop' and match the client source port, not a
		# saddr. Link-local only: :67 broadcasts never route in from the internet.
		dhcp_rule=$'\n\t\tudp sport 68 udp dport 67 accept'
	fi
	# Home Assistant (setup_homeassistant.sh): its UI/API on 8123 for the phone
	# app, from the LAN and over WireGuard only -- never the internet.
	local ha_rule=""
	if [[ ${ALLOW_HOMEASSISTANT:-false} == "true" ]]; then
		ha_rule=$'\n\t\tip saddr '"${LAN_SUBNET}"$' tcp dport 8123 accept'
		ha_rule+=$'\n\t\tiifname "'"${WG_IFACE}"$'" tcp dport 8123 accept'
	fi
	# Workout poke (screen-locker workout-poke.service): the phone POSTs a
	# finished workout to tcp/8774 the moment it ends. LAN only, deliberately
	# NO wg0 rule: off the home LAN the poke fails fast and the normal sync
	# credits the workout instead.
	local poke_rule=""
	if [[ ${ALLOW_WORKOUT_POKE:-false} == "true" ]]; then
		poke_rule=$'\n\t\tip saddr '"${LAN_SUBNET}"$' tcp dport 8774 accept'
	fi
	# DNS-over-TLS (setup_dot_resolver.sh, bound to the WireGuard address):
	# rendered here so a re-apply keeps it -- a runtime 'nft add rule' did not.
	local dot_rule=""
	if [[ ${ALLOW_DOT:-false} == "true" ]]; then
		dot_rule=$'\n\t\tiifname "'"${WG_IFACE}"$'" tcp dport 853 accept'
	fi
	# LAN gateway (setup_dns_blocker.sh gateway), rendered by wg_gateway.sh.
	local gw_set gw_forward gw_table
	gw_set="$(lan_gateway_set)"
	gw_forward="$(lan_gateway_forward)"
	gw_table="$(lan_gateway_table)"
	cat >"$target" <<EOF
#!/usr/sbin/nft -f
# Replaces ONLY this table, atomically. 'flush ruleset' also deleted every
# other table (Docker's under iptables-nft, Tailscale's), which is what made
# each apply restart docker. Declaring the table first makes the delete safe
# on a boot where it does not exist yet.
table inet filter
delete table inet filter

table inet filter {${gw_set}
	chain input {
		type filter hook input priority 0; policy drop;

		iif "lo" accept
		ct state established,related accept${dhcp_rule}
		ct state invalid drop

		icmp type { destination-unreachable, time-exceeded, parameter-problem, echo-request } accept
		icmpv6 type { destination-unreachable, packet-too-big, time-exceeded, parameter-problem, nd-neighbor-solicit, nd-neighbor-advert, echo-request } accept

		udp dport ${WG_PORT} accept

		iifname "${WG_IFACE}" tcp dport 22 accept
		ip saddr ${LAN_SUBNET} tcp dport 22 accept${web_rule}${dns_rule}${ha_rule}${poke_rule}${dot_rule}
	}
	chain forward {
		type filter hook forward priority 0; policy drop;${gw_forward}

		# Docker keeps its own FORWARD rules (DOCKER-USER, DOCKER-FORWARD,
		# the isolation chains) in the ip/filter table. This chain runs at
		# the same hook, so a bare 'policy drop' here VETOES them and every
		# bridged container loses all outbound networking -- while the host
		# itself stays perfectly online, so nothing looks broken. That is
		# what killed signal-cli's connection to Signal on 2026-08-29: the
		# bot stayed up, healthy and deaf for three hours.
		ct state established,related accept
		iifname "docker0" accept
		oifname "docker0" accept
		iifname "br-*" accept
		oifname "br-*" accept
	}
	chain output {
		type filter hook output priority 0; policy accept;
	}
}

${gw_table}
EOF
}

write_nftables_ruleset() {
	if [[ -f $NFT_CONF ]]; then
		cp "$NFT_CONF" "${NFT_CONF}.bak.$(date +%s)"
		log_warn "Backed up existing ${NFT_CONF} before overwriting."
	fi
	render_nftables_ruleset "${NFT_CONF}.new"
}

# The one invariant that has now broken twice: the forward chain must let
# Docker's own chains decide. Both times the repo carried the rule and
# /etc/nftables.conf did not, so a reboot reloaded the stale file and cut
# every bridged container off the internet while the host stayed online.
# Nothing compared the two, so nothing noticed.
readonly DOCKER_FORWARD_RULE='iifname "docker0" accept'

verify_nft() {
	local rendered status=0
	rendered="$(mktemp)"
	render_nftables_ruleset "$rendered"

	if diff -q "$rendered" "$NFT_CONF" >/dev/null 2>&1; then
		log_ok "${NFT_CONF} matches what this script generates."
	else
		log_error "${NFT_CONF} has DRIFTED from this script. A reboot would apply the stale file."
		diff -u "$NFT_CONF" "$rendered" | head -40 || true
		status=1
	fi
	rm -f "$rendered"

	if nft list chain inet filter forward 2>/dev/null | grep -qF "$DOCKER_FORWARD_RULE"; then
		log_ok "live forward chain lets Docker through."
	else
		log_error "live forward chain does NOT accept docker0 -- bridged containers have no network."
		status=1
	fi

	if ((status != 0)); then
		log_error "Re-apply with: sudo $0 allow-web"
	fi
	return "$status"
}

# Docker keeps its NAT/forward rules in iptables (legacy on this host, so not
# even visible to nft). Replacing only 'inet filter' leaves them alone, so a
# restart -- which bounces every container -- is reserved for when they are
# really gone (e.g. a hand-run 'nft flush ruleset' under iptables-nft).
restore_docker_rules() {
	if ! is_service_active docker; then
		return 0
	fi
	if iptables -t nat -S DOCKER >/dev/null 2>&1; then
		log_ok "docker's own firewall rules are intact; containers untouched."
		return 0
	fi
	log_warn "docker's nat rules are missing; restarting docker to rebuild them."
	systemctl restart docker
	log_ok "docker restarted; its firewall rules are back."
}

verify_nftables_then_apply() {
	nft -c -f "${NFT_CONF}.new" || die "nftables ruleset failed syntax check -- not applying."
	mv "${NFT_CONF}.new" "$NFT_CONF"
	log_warn "Applying a default-drop firewall now."
	nft -f "$NFT_CONF"
	sleep 2
	if ! is_service_active sshd; then
		nft delete table inet filter
		die "sshd died after applying nftables -- rolled back. Investigate before retrying."
	fi
	log_ok "nftables applied; sshd is still active."
	restore_docker_rules
	log_warn "Before closing this terminal, open a SECOND ssh session now and confirm it connects."
	enable_service nftables
}

allow_dot() {
	export ALLOW_DOT=true
	save_config
	write_nftables_ruleset
	verify_nftables_then_apply
	log_ok "Opened tcp/853 on ${WG_IFACE} (persisted)."
}

allow_workout_poke() {
	export ALLOW_WORKOUT_POKE=true
	save_config
	write_nftables_ruleset
	verify_nftables_then_apply
	log_ok "Opened tcp/8774 from ${LAN_SUBNET:-the LAN} only, not ${WG_IFACE} (persisted)."
}