#!/usr/bin/env bash
#
# Dezky production host firewall (nftables) — k3s-safe.
#
# Why this design:
#   - k3s/kube-proxy/flannel manage their OWN nftables tables (ip/ip6: filter,
#     nat, mangle). We must never `flush ruleset` or use ufw/firewalld, or we
#     wipe/clobber cluster networking. Instead we own a single dedicated table,
#     `inet dezky_fw`, with only an INPUT chain. Separate tables coexist; a
#     packet is dropped if ANY base chain drops it, so our default-drop INPUT
#     is the gate for host-bound traffic while k3s keeps owning FORWARD/NAT.
#   - We explicitly accept the pod/service CIDRs and CNI interfaces so
#     cluster<->host traffic (API server, kubelet, CoreDNS) is never dropped.
#
# Idempotent: re-running replaces only our table (`destroy table` first).
#
# Usage (as root, on the server):
#   ./firewall.sh            # render from ../config.env, install unit, apply
#   ./firewall.sh --dry-run  # print the ruleset, apply nothing

set -euo pipefail

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'
info()  { echo -e "${BLUE}[INFO]${NC}  $*"; }
ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
HOST_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
CONFIG_FILE="${CONFIG_FILE:-$HOST_DIR/config.env}"
NFT_OUT="/etc/nftables.d/dezky-fw.nft"
UNIT_SRC="$SCRIPT_DIR/dezky-firewall.service"
UNIT_DST="/etc/systemd/system/dezky-firewall.service"

DRY_RUN=0
[[ "${1:-}" == "--dry-run" ]] && DRY_RUN=1

# ── Load config ───────────────────────────────────────────────────────────
if [[ ! -f "$CONFIG_FILE" ]]; then
    error "Config not found: $CONFIG_FILE"
    error "Copy config.env.example → config.env and fill it in."
    exit 1
fi
# shellcheck disable=SC1090
source "$CONFIG_FILE"

: "${MGMT_ALLOW_V4:?MGMT_ALLOW_V4 is required in config.env}"
: "${SSH_PORT:=22}"
: "${K3S_POD_CIDR:=10.42.0.0/16}"
: "${K3S_SERVICE_CIDR:=10.43.0.0/16}"

# ── Build the management v6 block only if a v6 address is configured ───────
V6_SET=""
V6_RULE=""
if [[ -n "${MGMT_ALLOW_V6:-}" ]]; then
    V6_SET=$(cat <<EOF

    set mgmt_v6 {
        type ipv6_addr
        flags interval
        elements = { ${MGMT_ALLOW_V6} }
    }
EOF
)
    V6_RULE="        ip6 saddr @mgmt_v6 tcp dport { ${SSH_PORT}, 6443 } accept"
fi

# ── Render the ruleset ─────────────────────────────────────────────────────
RULESET=$(cat <<EOF
#!/usr/sbin/nft -f
#
# Managed by Dezky firewall.sh — DO NOT edit by hand.
# Owns only 'inet dezky_fw'. k3s manages its own ip/ip6 tables separately.
# NEVER add 'flush ruleset' here: it would wipe k3s networking.

destroy table inet dezky_fw

table inet dezky_fw {
    # Management source allowlist (SSH + k3s API). Intervals allow CIDRs.
    set mgmt_v4 {
        type ipv4_addr
        flags interval
        elements = { ${MGMT_ALLOW_V4} }
    }${V6_SET}

    chain input {
        type filter hook input priority filter; policy drop;

        # Stateful fast-path
        ct state established,related accept
        ct state invalid drop

        # Loopback
        iif "lo" accept

        # ICMP — keep ping working and (critically) IPv6 NDP/RA + PMTUD
        ip  protocol icmp accept
        ip6 nexthdr icmpv6 accept

        # ── k3s internal: never block cluster <-> host traffic ──────────────
        iifname "cni0"      accept
        iifname "flannel.1" accept
        ip saddr ${K3S_POD_CIDR}     accept
        ip saddr ${K3S_SERVICE_CIDR} accept

        # ── Public services (world-reachable) ──────────────────────────────
        # Web + ACME HTTP-01 challenge
        tcp dport { 80, 443 } accept
        # Mail: smtp, submissions, submission, imap, imaps, managesieve
        tcp dport { 25, 465, 587, 143, 993, 4190 } accept

        # ── Management surfaces: home IP only ──────────────────────────────
        ip  saddr @mgmt_v4 tcp dport { ${SSH_PORT}, 6443 } accept
${V6_RULE}

        # Rate-limited drop logging for debugging (then policy drop applies)
        limit rate 5/minute burst 5 packets log prefix "dezky-fw drop: " level info
    }
}
EOF
)

if [[ $DRY_RUN -eq 1 ]]; then
    echo "$RULESET"
    info "Dry run — nothing applied."
    exit 0
fi

if [[ $EUID -ne 0 ]]; then
    error "Must run as root to apply the firewall."
    exit 1
fi

# ── Write, validate, install, apply ────────────────────────────────────────
mkdir -p /etc/nftables.d
echo "$RULESET" > "$NFT_OUT"
chmod 0644 "$NFT_OUT"
info "Wrote ruleset → $NFT_OUT"

# Validate syntax before touching the live ruleset
if ! nft -c -f "$NFT_OUT"; then
    error "nft syntax check FAILED — not applying. Live firewall unchanged."
    exit 1
fi
ok "Ruleset syntax valid."

# Install the systemd unit so the rules survive reboot (and never flush global)
if [[ -f "$UNIT_SRC" ]]; then
    install -m 0644 "$UNIT_SRC" "$UNIT_DST"
    systemctl daemon-reload
    systemctl enable dezky-firewall.service >/dev/null 2>&1 || true
    ok "Installed + enabled dezky-firewall.service"
fi

# Apply now
nft -f "$NFT_OUT"
ok "Firewall applied. Management restricted to: ${MGMT_ALLOW_V4} ${MGMT_ALLOW_V6:-}"
warn "Open a SECOND SSH session NOW and confirm you still have access before"
warn "closing this one. Hetzner KVM/LARA is your out-of-band fallback."