#!/usr/bin/env bash
# Own one nftables table. Never flush Docker, Tailscale or the host ruleset.
set -euo pipefail
export LC_ALL=C
fail() { printf 'ERROR: %s\n' "$*" >&2; return 1; }
validate_policy() {
  jq -e '
    def ip:
      type == "string" and test("^(0|[1-9][0-9]{0,2})(\\.(0|[1-9][0-9]{0,2})){3}$") and
      (split(".") | all(.[]; tonumber <= 255));
    def tailip: ip and (split(".") | map(tonumber) | .[0] == 100 and .[1] >= 64 and .[1] <= 127);
    def ports: type == "array" and all(.[]; type == "number" and . == floor and . >= 1 and . <= 65535);
    ((keys - ["trusted_sources"]) == ["hostname","management","public_tcp","public_udp","tailnet","version"]) and .version == 1 and
    (.hostname | type == "string" and length <= 63 and test("^(dev|test|stage|prod)-[a-z][a-z0-9]{0,15}-[a-z][a-z0-9]{0,15}(-[a-z][a-z0-9]{0,15})?-(0[1-9]|[1-9][0-9])$")) and
    (.public_tcp | ports) and (.public_udp | ports) and
    ((.trusted_sources // []) | type == "array" and all(.[]; ip and (split(".")[0] | tonumber | . > 0 and . < 224))) and
    (.management | keys == ["port","sources"]) and ([.management.port] | ports) and
    (.management.sources | type == "array" and all(.[]; ip and (tailip | not))) and
    (.tailnet | type == "array" and all(.[];
      keys == ["sources","tcp","udp"] and (.sources | type == "array" and length > 0 and all(.[]; tailip)) and
      (.tcp | ports) and (.udp | ports) and ((.tcp + .udp) | length > 0))) and
    ((.management.sources | length > 0) or
      ( .management.port as $ssh | any(.tailnet[]; .tcp | index($ssh) != null))) and
    (.management.port as $ssh | (.public_tcp | index($ssh) == null))
  ' "$1" >/dev/null || { fail 'Invalid firewall policy: exact IPv4 sources, ports, hostname and an explicit SSH path required'; return 1; }
}
render_policy() {
  local policy=$1 ports row sources protocol
  validate_policy "$policy" || return 1
  cat <<'NFT'
add table inet platform_host
delete table inet platform_host
table inet platform_host {
  chain input {
    type filter hook input priority 10; policy drop;
    iifname "lo" accept
    iifname != "tailscale0" ip saddr 100.64.0.0/10 drop
    ct state invalid drop
    ct state established,related accept
    iifname "tailscale0" ip saddr 100.64.0.0/10 accept
    iifname != "tailscale0" ip saddr { 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16 } accept
    meta l4proto icmp accept
    meta l4proto ipv6-icmp accept
    udp sport 67 udp dport 68 accept
    ip6 saddr fe80::/10 udp sport 547 udp dport 546 accept
NFT
  # Public control hosts are explicit per-environment inputs, never compiled IPs.
  sources=$(jq -r '(.trusted_sources // []) | unique | join(", ")' "$policy")
  if [[ -n "$sources" ]]; then
    printf '    ip saddr { %s } accept\n' "$sources"
  fi
  sources=$(jq -r '.management.sources | join(", ")' "$policy")
  if [[ -n "$sources" ]]; then
    printf '    iifname != "tailscale0" ip saddr { %s } tcp dport %s accept\n' "$sources" "$(jq -r .management.port "$policy")"
  fi
  for row in tcp udp; do
    ports=$(jq -r --arg field "public_$row" '.[$field] | unique | map(tostring) | join(", ")' "$policy")
    if [[ -n "$ports" ]]; then printf '    iifname != "tailscale0" %s dport { %s } accept\n' "$row" "$ports"; fi
  done
  # Match the decrypted interface AND Headscale-assigned source; no LAN spoofing.
  while IFS= read -r row; do
    sources=$(jq -r '.sources | unique | join(", ")' <<< "$row")
    for protocol in tcp udp; do
      ports=$(jq -r --arg protocol "$protocol" '.[$protocol] | unique | map(tostring) | join(", ")' <<< "$row")
      if [[ -n "$ports" ]]; then printf '    iifname "tailscale0" ip saddr { %s } %s dport { %s } accept\n' "$sources" "$protocol" "$ports"; fi
    done
  done < <(jq -c '.tailnet[]' "$policy")
  cat <<'NFT'
  }
  chain forward {
    type filter hook forward priority 10; policy accept;
    iifname "tailscale0" drop
    oifname "tailscale0" drop
  }
}
NFT
}
trusted_path() {
  local path=$1
  [[ "$path" == /* && "$path" != *'/../'* && "$path" != *'/./'* ]] || { fail 'Absolute trusted path required'; return 1; }
  while [[ "$path" != / ]]; do
    [[ ! -L "$path" ]] || { fail 'Symlink in firewall path'; return 1; }
    if [[ -e "$path" ]]; then
      [[ $(stat -c %u "$path") == 0 ]] && (( (8#$(stat -c %a "$path") & 0022) == 0 )) || { fail 'Firewall path must be root owned without group/other write permission'; return 1; }
    fi
    path=$(dirname -- "$path")
  done
}
save_boot_gate() {
  local state=$1 file
  install -d -m 0700 "$state/boot"
  for file in platform-firewall.service docker.service.d/20-platform-firewall.conf tailscaled.service.d/20-platform-firewall.conf; do
    trusted_path "/etc/systemd/system/$file"
    install -d -m 0700 "$state/boot/$(dirname -- "$file")"
    if [[ -f "/etc/systemd/system/$file" ]]; then
      cp -p -- "/etc/systemd/system/$file" "$state/boot/$file"
    else rm -f -- "$state/boot/$file"; fi
  done
  rm -f -- "$state/boot/changed" "$state/boot/enabled" "$state/boot/active"
  if systemctl is-enabled --quiet platform-firewall.service; then touch "$state/boot/enabled"; fi
  if systemctl is-active --quiet platform-firewall.service; then touch "$state/boot/active"; fi
}
install_boot_gate() {
  local directory
  for directory in /etc/platform-firewall /etc/systemd/system/docker.service.d /etc/systemd/system/tailscaled.service.d; do
    trusted_path "$directory"; install -d -m 0755 "$directory"
  done
  trusted_path /etc/systemd/system/platform-firewall.service
  cat > /etc/systemd/system/platform-firewall.service <<'UNIT'
[Unit]
Description=Platform host service firewall
After=nftables.service
Before=docker.service tailscaled.service
[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/sbin/nft -f /etc/platform-firewall/active.nft
ExecReload=/usr/sbin/nft -f /etc/platform-firewall/active.nft
[Install]
WantedBy=multi-user.target
UNIT
  for directory in docker tailscaled; do
    trusted_path "/etc/systemd/system/$directory.service.d/20-platform-firewall.conf"
    cat > "/etc/systemd/system/$directory.service.d/20-platform-firewall.conf" <<'UNIT'
[Unit]
Requires=platform-firewall.service
After=platform-firewall.service
UNIT
  done
  systemctl daemon-reload
  systemctl enable platform-firewall.service >/dev/null
  # Reload an active gate: restarting a Required unit stops Docker/Tailscale.
  systemctl reload-or-restart platform-firewall.service
}
check_ssh_transaction() {
  local transaction=$1 digest=$2 ssh_state=${3:-/var/lib/platform-bootstrap/ssh}
  if [[ -e "$ssh_state/pending" ]]; then
    trusted_path "$ssh_state/pending" || return 1
    trusted_path "$ssh_state/expected-firewall-sha" || return 1
    [[ "$transaction" =~ ^[a-f0-9]{64}$ && "$transaction" == "$(cat "$ssh_state/pending")" &&
       "$digest" == "$(cat "$ssh_state/expected-firewall-sha")" ]] || {
      fail 'SSH maintenance owns the firewall; exact transaction and target policy required'; return 1;
    }
  else
    [[ -z "$transaction" ]] || { fail 'SSH transaction expired or already completed'; return 1; }
  fi
}
main() {
  local action='' policy='' expected='' ssh_transaction='' state=/run/platform-firewall task digest previous connection
  while (($#)); do
    case "$1" in
      --help|-h) echo 'Usage: bash firewall.sh --render|--check|--apply --policy ROOT_0600_JSON | --confirm --sha256 SHA [--ssh-transaction SHA during SSH migration]'; return 0;;
      --render|--check|--apply|--confirm) [[ -z "$action" ]] || return 1; action=$1; shift;;
      --policy|--sha256|--ssh-transaction) (($# >= 2)) || return 1; case "$1" in --policy) policy=$2;; --sha256) expected=$2;; --ssh-transaction) ssh_transaction=$2;; esac; shift 2;;
      *) fail "Unknown argument: $1"; return 1;;
    esac
  done
  if [[ "$action" == --render ]]; then render_policy "$policy"; return; fi
  umask 077
  [[ $EUID == 0 && -d /run/systemd/system ]] || { fail 'Native root systemd host required'; return 1; }
  for task in nft jq sha256sum flock systemctl systemd-run; do command -v "$task" >/dev/null || { fail "Missing command: $task"; return 1; }; done
  trusted_path "$state"; install -d -m 0700 "$state"
  trusted_path "$state/lock"
  exec 9>>"$state/lock"; flock -n 9 || { fail 'Firewall operation in progress'; return 1; }
  if [[ "$action" == --confirm ]]; then
    [[ "$expected" =~ ^[a-f0-9]{64}$ && -f "$state/pending.nft" ]] || { fail 'Exact pending policy digest required'; return 1; }
    digest=$(sha256sum "$state/pending.nft"); digest=${digest%% *}
    [[ "$expected" == "$digest" ]] || { fail 'Pending policy differs'; return 1; }
    check_ssh_transaction "$ssh_transaction" "$digest"
    connection=${SSH_CONNECTION:-}
    [[ -n "$connection" && "$connection" != "$(cat "$state/original-ssh")" ]] || { fail 'Confirm from a new SSH transport after checking allowed and denied ports'; return 1; }
    nft -s list table inet platform_host > "$state/current.nft"
    cmp -s "$state/current.nft" "$state/applied.nft" || { fail 'Active policy changed after apply'; return 1; }
    trusted_path /etc/platform-firewall/active.nft
    trusted_path /etc/platform-firewall/active.nft.pending
    install -d -m 0755 /etc/platform-firewall
    install -m 0600 "$state/pending.nft" /etc/platform-firewall/active.nft.pending
    mv -fT /etc/platform-firewall/active.nft.pending /etc/platform-firewall/active.nft
    touch "$state/boot/changed"
    install_boot_gate
    rm -f -- "$state/pending.nft"
    systemctl stop platform-firewall-rollback.timer
    echo 'Firewall confirmed and persisted; Docker/Tailscale startup now requires it.'
    return
  fi
  [[ "$action" == --check || "$action" == --apply ]] || { fail 'Choose render, check, apply or confirm'; return 1; }
  trusted_path "$policy"
  [[ -f "$policy" && $(stat -c %a "$policy") == 600 ]] || { fail 'Use a root-owned 0600 JSON policy file'; return 1; }
  validate_policy "$policy"
  [[ $(jq -r .hostname "$policy") == "$(hostname)" ]] || { fail 'Policy belongs to another host'; return 1; }
  task=$(mktemp -d "$state/check-XXXXXXXX")
  local cleanup
  printf -v cleanup 'rm -rf -- %q' "$task"; trap "$cleanup" EXIT
  render_policy "$policy" > "$task/rules.nft"
  nft --check --file "$task/rules.nft"
  digest=$(sha256sum "$task/rules.nft"); digest=${digest%% *}
  if [[ "$action" == --check ]]; then echo "Firewall policy checked: $digest"; return; fi
  check_ssh_transaction "$ssh_transaction" "$digest"
  [[ ! -f "$state/pending.nft" ]] || { fail 'A policy is pending confirmation; inspect it or wait for rollback'; return 1; }
  # Rollback only our table, preserving concurrent native Tailscale/Docker rules.
  printf 'add table inet platform_host\ndelete table inet platform_host\n' > "$task/rollback.nft"
  if nft list table inet platform_host >/dev/null 2>&1; then nft -s list table inet platform_host >> "$task/rollback.nft"; fi
  install -m 0600 "$task/rollback.nft" "$state/rollback.nft"
  # Save the durable version too, in case confirmation fails midway.
  trusted_path /etc/platform-firewall/active.nft
  if [[ -f /etc/platform-firewall/active.nft ]]; then
    install -m 0600 /etc/platform-firewall/active.nft "$state/previous.nft"
  else rm -f -- "$state/previous.nft"; fi
  save_boot_gate "$state"
  cat > "$state/rollback.sh" <<'ROLLBACK'
#!/usr/bin/env bash
set -euo pipefail
exec 9>>/run/platform-firewall/lock
flock 9
test -f /run/platform-firewall/pending.nft || exit 0
if test -f /run/platform-firewall/previous.nft; then
  install -m 0600 /run/platform-firewall/previous.nft /etc/platform-firewall/active.nft
else
  rm -f /etc/platform-firewall/active.nft
fi
if test -f /run/platform-firewall/boot/changed; then
  if ! test -f /run/platform-firewall/boot/active; then systemctl stop platform-firewall.service; fi
  if ! test -f /run/platform-firewall/boot/enabled; then systemctl disable platform-firewall.service; fi
  for file in platform-firewall.service docker.service.d/20-platform-firewall.conf tailscaled.service.d/20-platform-firewall.conf; do
    if test -f "/run/platform-firewall/boot/$file"; then
      cp -p -- "/run/platform-firewall/boot/$file" "/etc/systemd/system/$file"
    else rm -f -- "/etc/systemd/system/$file"; fi
  done
  systemctl daemon-reload
  if test -f /run/platform-firewall/boot/enabled; then systemctl enable platform-firewall.service; fi
  if test -f /run/platform-firewall/boot/active; then systemctl start platform-firewall.service; fi
fi
/usr/sbin/nft -f /run/platform-firewall/rollback.nft
rm -f /run/platform-firewall/pending.nft
ROLLBACK
  chmod 0700 "$state/rollback.sh"
  install -m 0600 "$task/rules.nft" "$state/pending.nft"
  printf '%s\n' "${SSH_CONNECTION:-}" > "$state/original-ssh"
  systemctl stop platform-firewall-rollback.timer platform-firewall-rollback.service 2>/dev/null || true
  systemctl reset-failed platform-firewall-rollback.service 2>/dev/null || true
  if ! systemd-run --quiet --collect --unit=platform-firewall-rollback --on-active=180s /bin/bash "$state/rollback.sh"; then
    rm -f -- "$state/pending.nft"; fail 'Rollback timer could not be armed; no rules changed'; return 1
  fi
  nft --file "$state/pending.nft"
  nft -s list table inet platform_host > "$state/applied.nft"
  printf 'Policy applied temporarily. Confirm from a NEW SSH connection within 180 seconds: --confirm --sha256 %s\n' "$digest"
}
if [[ ${BASH_SOURCE[0]:-} == "$0" || -z ${BASH_SOURCE[0]:-} ]]; then main "$@"; fi