#!/usr/bin/env bash
#
# harden-firewall.sh
#
# Deny-by-default host firewall (iptables) for Alpine, Debian, and Alma Linux.
# Hardens the INPUT chain to: loopback, established/related, ICMP, SSH (on the
# configured port), and any explicitly-registered ports -- everything else is
# dropped. OUTPUT stays open (egress allow); FORWARD is left untouched so Docker
# container networking is unaffected.
#
# All distro-specific operations go through scripts/oslib.sh. The OS-specific
# surface here is exactly three things, all in oslib: which packages provide
# iptables + native persistence (install_iptables), how to persist the live
# ruleset (fw_save_cmd), and which service restores it at boot (fw_enable_restore).
#
# Persistence is NATIVE -- there is no boot hook. The installer builds the rules
# live, saves them via the distro's own iptables/iptables-persistent/
# iptables-services package, and enables that package's restore service. The
# saved ruleset carries the INPUT->sshguard jump (and the empty sshguard chain),
# so brute-force protection survives reboot without any custom hook.
#
# A small on-host CLI, /usr/local/sbin/firewall-apply, rebuilds INPUT from the
# declarative drop-ins under /etc/firewall and then runs the native save. It is
# generated here (with the save command baked in) so deployments can register a
# port and re-apply without this repo present:
#
#     printf '80/tcp\n443/tcp\n' > /etc/firewall/ports.d/mystack.rule
#     /usr/local/sbin/firewall-apply
#
# Note on Docker: containers published with `-p` reach the host through the nat
# PREROUTING + FORWARD chains, which BYPASS INPUT. This firewall therefore can
# neither block nor needs to open Docker-published ports (e.g. Caddy's 80/443) --
# the INPUT rules here protect host-bound services only.
#
# Usage:
#   bash harden-firewall.sh                          # install + apply (deny-by-default)
#   SSH_PORT=2222 bash harden-firewall.sh            # bastion port
#   OPEN_PORTS="80/tcp 443/tcp" bash harden-firewall.sh   # open extra ports at install
#   FW_SSH_SOURCE=10.0.0.0/8 bash harden-firewall.sh # restrict SSH to a source CIDR
#   FW_ALLOW_PING=0 bash harden-firewall.sh          # drop ICMP echo (ping)
#
#   bash harden-firewall.sh allow 443/tcp 51820/udp  # register + apply
#   bash harden-firewall.sh allow web                # preset: 80/tcp + 443/tcp
#   bash harden-firewall.sh allow 2222/tcp@10.0.0.0/8 # port limited to a source
#   bash harden-firewall.sh deny 51820/udp           # unregister + apply
#   bash harden-firewall.sh list                     # show drop-ins + live INPUT
#   bash harden-firewall.sh disable                  # recovery: flush + INPUT ACCEPT

set -euo pipefail

# Load the OS abstraction layer (sits next to this script).
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=scripts/oslib.sh
. "$SCRIPT_DIR/oslib.sh"

# ============================================================================
# CONFIG
# ============================================================================
: "${FW_ALLOW_PING:=1}"     # 1 = allow ICMP echo-request (ping); 0 = drop it
: "${FW_SSH_SOURCE:=}"      # optional CIDR: restrict SSH to this source (default any)
: "${OPEN_PORTS:=}"         # space/comma list of specs to open at install
: "${FORCE:=0}"

FW_DIR=/etc/firewall
PORTS_DIR="$FW_DIR/ports.d"
CONF="$FW_DIR/firewall.conf"
APPLY=/usr/local/sbin/firewall-apply

# log()/warn()/die() come from oslib (_log/_warn/_die); alias for readability.
log()  { _log  "$@"; }
warn() { _warn "$@"; }
die()  { _die  "$@"; }

[[ $EUID -eq 0 ]] || die "Run as root."
os_detect

# SSH port: explicit env wins; otherwise read it live from sshd_config; else 22.
# (The generated engine re-derives this on every run so a later port change is
# picked up; SSH_PORT here is the install-time fallback baked into the conf.)
if [[ -z "${SSH_PORT:-}" ]]; then
    SSH_PORT="$(awk '/^[Pp]ort / {print $2; exit}' /etc/ssh/sshd_config 2>/dev/null || true)"
    [[ -n "$SSH_PORT" ]] || SSH_PORT=22
fi

# ============================================================================
# Spec parsing: PORT[/PROTO][@CIDR] and named presets -> "PORT/PROTO[ CIDR]" lines
# ============================================================================
_preset_lines() {  # echo drop-in lines for a preset name, else return 1
    case "$1" in
        web)   printf '80/tcp\n443/tcp\n' ;;
        http)  printf '80/tcp\n' ;;
        https) printf '443/tcp\n' ;;
        *) return 1 ;;
    esac
}

_normalize_spec() {  # PORT[/PROTO][@CIDR] -> a single "PORT/PROTO[ CIDR]" line
    local spec="$1" cidr="" port proto
    case "$spec" in *@*) cidr="${spec##*@}"; spec="${spec%@*}" ;; esac
    port="${spec%%/*}"
    proto=tcp
    case "$spec" in */*) proto="${spec##*/}" ;; esac
    case "$proto" in tcp|udp) : ;; *) die "Invalid protocol '$proto' in '$1' (tcp|udp)." ;; esac
    case "$port" in ''|*[!0-9:]*) die "Invalid port '$port' in '$1' (digits, or a:b range)." ;; esac
    if [[ -n "$cidr" ]]; then printf '%s/%s %s\n' "$port" "$proto" "$cidr"
    else printf '%s/%s\n' "$port" "$proto"; fi
}

_expand_args() {  # mix of presets + specs -> newline-separated drop-in lines
    local a
    for a in "$@"; do
        _preset_lines "$a" 2>/dev/null || _normalize_spec "$a"
    done
}

# ============================================================================
# firewall.conf -- the declarative inputs the engine reads on every run
# ============================================================================
write_conf() {
    install -d -m 0755 "$FW_DIR" "$PORTS_DIR"
    cat > "$CONF" <<EOF
# Generated by harden-firewall.sh -- $(date -u +%FT%TZ)
# Read by ${APPLY} on every run. After editing, re-apply with: ${APPLY##*/}
FW_SSH_PORT="${SSH_PORT}"
FW_ALLOW_PING="${FW_ALLOW_PING}"
FW_SSH_SOURCE="${FW_SSH_SOURCE}"
EOF
    chmod 0644 "$CONF"
}

# OPEN_PORTS env -> a declarative seed drop-in (rewritten each install).
seed_open_ports() {
    [[ -n "$OPEN_PORTS" ]] || return 0
    install -d -m 0755 "$PORTS_DIR"
    # shellcheck disable=SC2086  # word-split the space/comma list intentionally
    _expand_args ${OPEN_PORTS//,/ } | awk 'NF' | sort -u > "$PORTS_DIR/seed.rule"
    log "Seeded ports.d/seed.rule from OPEN_PORTS: ${OPEN_PORTS}"
}

# ============================================================================
# Generate /usr/local/sbin/firewall-apply (the on-demand rebuild+save engine).
# The native save command is resolved here and baked in, so the helper is fully
# self-contained on the host (no repo / no oslib needed to re-apply).
# ============================================================================
write_apply() {
    local save_cmd; save_cmd="$(fw_save_cmd)"
    cat > "$APPLY" <<'ENGINE'
#!/bin/sh
# Managed by harden-firewall.sh -- do not edit by hand (regenerated on install).
#
# Rebuilds the INPUT chain from /etc/firewall (declarative) and persists it via
# the distro's native iptables save. This is NOT a boot service: boot-time
# restore is handled by the native iptables/iptables-persistent/iptables-services
# unit that harden-firewall.sh enabled. Run this after changing /etc/firewall.
set -eu

FW_DIR=/etc/firewall
PORTS_DIR="$FW_DIR/ports.d"
CONF="$FW_DIR/firewall.conf"

# Declarative inputs (overridden by firewall.conf when present).
FW_SSH_PORT=22
FW_ALLOW_PING=1
FW_SSH_SOURCE=
# shellcheck disable=SC1090
[ -r "$CONF" ] && . "$CONF"

# SSH port: prefer the live sshd_config value so a port change is picked up;
# fall back to the conf value, then 22.
SSH_PORT="$(awk '/^[Pp]ort / {print $2; exit}' /etc/ssh/sshd_config 2>/dev/null || true)"
[ -n "${SSH_PORT:-}" ] || SSH_PORT="${FW_SSH_PORT:-22}"

ssh_src=""
[ -n "${FW_SSH_SOURCE:-}" ] && ssh_src="-s $FW_SSH_SOURCE"

# Rebuild one address family. $1 = binary, $2 = its ICMP protocol keyword.
# Order matters: open the policy, flush, add every ACCEPT, THEN drop -- so there
# is never a window where INPUT is DROP without the SSH/established rules, and
# the ESTABLISHED accept keeps the current SSH session alive across the rebuild.
apply_family() {
    ipt="$1"; icmp="$2"
    command -v "$ipt" >/dev/null 2>&1 || return 0

    # Ensure the sshguard chain exists; NEVER flush it -- the sshguard daemon
    # owns its contents (the live block list).
    "$ipt" -N sshguard 2>/dev/null || true

    "$ipt" -P INPUT ACCEPT
    "$ipt" -F INPUT
    "$ipt" -A INPUT -i lo -j ACCEPT
    "$ipt" -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    "$ipt" -A INPUT -m conntrack --ctstate INVALID -j DROP

    # ICMP. IPv6 needs NDP/MLD (and PMTUD) or the link breaks, so accept all
    # ipv6-icmp. For IPv4 accept the essential diagnostics and, optionally, ping.
    if [ "$icmp" = "ipv6-icmp" ]; then
        "$ipt" -A INPUT -p ipv6-icmp -j ACCEPT
    else
        "$ipt" -A INPUT -p icmp --icmp-type destination-unreachable -j ACCEPT
        "$ipt" -A INPUT -p icmp --icmp-type time-exceeded -j ACCEPT
        "$ipt" -A INPUT -p icmp --icmp-type parameter-problem -j ACCEPT
        if [ "${FW_ALLOW_PING:-1}" = "1" ]; then
            "$ipt" -A INPUT -p icmp --icmp-type echo-request \
                -m limit --limit 5/second --limit-burst 10 -j ACCEPT
        fi
    fi

    # SSH: hand NEW connections to sshguard first (it drops blocked sources),
    # then accept. Both honour the optional source restriction.
    # shellcheck disable=SC2086
    "$ipt" -A INPUT -p tcp --dport "$SSH_PORT" $ssh_src -j sshguard
    # shellcheck disable=SC2086
    "$ipt" -A INPUT -p tcp --dport "$SSH_PORT" $ssh_src -m conntrack --ctstate NEW -j ACCEPT

    # Registered drop-in ports: "<port>/<proto> [source-cidr]" per line.
    if [ -d "$PORTS_DIR" ]; then
        for f in "$PORTS_DIR"/*.rule; do
            [ -e "$f" ] || continue
            while IFS= read -r line || [ -n "$line" ]; do
                case "$line" in ''|\#*) continue ;; esac
                spec=$(printf '%s' "$line" | awk '{print $1}')
                src=$(printf '%s'  "$line" | awk '{print $2}')
                port=${spec%%/*}
                proto=${spec##*/}
                [ "$proto" = "$spec" ] && proto=tcp
                case "$proto" in tcp|udp) : ;; *) continue ;; esac
                case "$port" in ''|*[!0-9:]*) continue ;; esac
                srcarg=""
                [ -n "$src" ] && srcarg="-s $src"
                # shellcheck disable=SC2086
                "$ipt" -A INPUT -p "$proto" --dport "$port" $srcarg -j ACCEPT
            done < "$f"
        done
    fi

    "$ipt" -P INPUT DROP
}

apply_family iptables  icmp
apply_family ip6tables ipv6-icmp
ENGINE

    # Append the native persist step (resolved for this distro, runs last).
    {
        echo ''
        echo '# Persist the live ruleset so it survives reboot (native, per-distro).'
        echo "$save_cmd"
    } >> "$APPLY"
    chmod 0755 "$APPLY"
}

run_apply() {
    [[ -x "$APPLY" ]] || die "Firewall engine missing ($APPLY). Run: bash $0 apply"
    "$APPLY"
}

ensure_installed() {
    [[ -x "$APPLY" ]] || die "Firewall not installed. Run: bash $0 apply"
    install -d -m 0755 "$PORTS_DIR"
}

# ============================================================================
# Subcommands
# ============================================================================
cmd_apply() {
    log "Detected OS: ${OS_ID} (family ${OS_FAMILY}, init ${INIT_SYSTEM})"

    # RHEL/Alma ships firewalld; it and iptables-services fight over nftables.
    # Decoupled from FORCE on purpose: harden-ssh passes FORCE=1 to skip its SSH
    # prompt, and we must NOT let that bulldoze an active firewalld. Bail clearly
    # instead (harden-ssh treats this as a warning and leaves firewalld in place).
    if [[ "$OS_FAMILY" == rhel ]] && command -v firewall-cmd >/dev/null 2>&1 \
        && firewall-cmd --state >/dev/null 2>&1; then
        warn "firewalld is active. It conflicts with this iptables firewall."
        warn "Disable it first:  systemctl disable --now firewalld"
        [[ "${FW_IGNORE_FIREWALLD:-0}" == "1" ]] || die "Refusing to proceed while firewalld is active (set FW_IGNORE_FIREWALLD=1 to override)."
    fi

    log "Installing iptables + native persistence package..."
    install_iptables

    write_conf
    seed_open_ports
    write_apply
    fw_enable_restore

    log "Building and applying firewall rules (deny-by-default INPUT)..."
    run_apply

    log "Firewall active. Allowed: loopback, established, ICMP$([[ $FW_ALLOW_PING == 1 ]] || echo ' (no ping)'), SSH/${SSH_PORT}${FW_SSH_SOURCE:+ from $FW_SSH_SOURCE}, + registered ports."
    log "Add ports:  ${APPLY##*/}-managed -> 'bash $0 allow <port/proto>'   Recover:  bash $0 disable"
}

cmd_allow() {
    [[ "$#" -ge 1 ]] || die "usage: allow <port[/proto][@cidr]|web|http|https>..."
    ensure_installed
    local lines file="$PORTS_DIR/manual.rule"
    lines="$(_expand_args "$@")"
    touch "$file"
    local l
    while IFS= read -r l; do
        [[ -n "$l" ]] || continue
        grep -qxF "$l" "$file" 2>/dev/null || printf '%s\n' "$l" >> "$file"
    done <<< "$lines"
    log "Registered: $(printf '%s' "$lines" | tr '\n' ' ')"
    run_apply
    log "Applied."
}

cmd_deny() {
    [[ "$#" -ge 1 ]] || die "usage: deny <port[/proto][@cidr]|web|http|https>..."
    ensure_installed
    local lines l file tmp
    lines="$(_expand_args "$@")"
    while IFS= read -r l; do
        [[ -n "$l" ]] || continue
        for file in "$PORTS_DIR"/*.rule; do
            [[ -e "$file" ]] || continue
            tmp="$(mktemp)"
            grep -vxF "$l" "$file" > "$tmp" 2>/dev/null || true
            mv "$tmp" "$file"
        done
    done <<< "$lines"
    # Drop any drop-in file we emptied.
    for file in "$PORTS_DIR"/*.rule; do
        [[ -e "$file" ]] || continue
        [[ -s "$file" ]] || rm -f "$file"
    done
    log "Unregistered: $(printf '%s' "$lines" | tr '\n' ' ')"
    run_apply
    log "Applied."
}

cmd_list() {
    echo "Registered ports ($PORTS_DIR):"
    if ls "$PORTS_DIR"/*.rule >/dev/null 2>&1; then
        for f in "$PORTS_DIR"/*.rule; do
            printf '  [%s]\n' "$(basename "$f")"
            sed 's/^/    /' "$f"
        done
    else
        echo "  (none)"
    fi
    echo
    echo "Live INPUT chain (iptables):"
    iptables -S INPUT 2>/dev/null | sed 's/^/  /' || echo "  (unavailable)"
}

# Recovery escape hatch: open the policy and flush our rules, then persist that
# open state so a reboot does not re-DROP. The sshguard chain is left intact.
cmd_disable() {
    local ipt
    for ipt in iptables ip6tables; do
        command -v "$ipt" >/dev/null 2>&1 || continue
        "$ipt" -P INPUT ACCEPT || true
        "$ipt" -F INPUT || true
    done
    eval "$(fw_save_cmd)"
    warn "Firewall disabled: INPUT policy ACCEPT, rules flushed and saved."
    warn "Re-enable with: bash $0 apply"
}

usage() {
    sed -n '2,50p' "$0" | sed 's/^#\{0,1\} \{0,1\}//'
}

# ============================================================================
# Dispatch
# ============================================================================
cmd="${1:-apply}"
case "$cmd" in
    apply)        cmd_apply ;;
    allow)        shift; cmd_allow "$@" ;;
    deny|remove)  shift; cmd_deny "$@" ;;
    list|status)  cmd_list ;;
    disable)      cmd_disable ;;
    -h|--help|help) usage ;;
    *) die "Unknown command '$cmd'. Use: apply | allow | deny | list | disable." ;;
esac