#!/usr/bin/env bash
#
# harden-firewall.sh
#
# Deny-by-default host firewall for Alpine, Debian, and Alma Linux.
#
# Backend per family (override with FW_BACKEND=iptables|firewalld):
#   alpine/debian  iptables -- hardens the INPUT chain to loopback,
#                  established/related, ICMP, SSH (configured port) + registered
#                  ports; everything else dropped. Persisted natively, no boot hook.
#   rhel (Alma)    firewalld -- the native RHEL firewall. The default zone is
#                  deny-by-default; we open SSH + registered ports and let
#                  sshguard block via the sshguard-firewalld backend.
# OUTPUT/egress stays open; FORWARD is left untouched so Docker networking is
# unaffected. The allow/deny/list/disable subcommands work on both backends.
#
# 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

# Firewall backend: firewalld on RHEL/Alma (native, and the sshguard-firewalld
# backend blocks there), iptables everywhere else. Override with
# FW_BACKEND=iptables|firewalld.
: "${FW_BACKEND:=}"
if [[ -z "$FW_BACKEND" ]]; then
    [[ "$OS_FAMILY" == rhel ]] && FW_BACKEND=firewalld || FW_BACKEND=iptables
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"
}

# ============================================================================
# firewalld backend (RHEL/Alma) -- same CLI as the iptables engine, driven via
# firewall-cmd. firewalld zones are deny-by-default for unsolicited inbound and
# persist natively (--permanent), and sshguard blocks via the sshguard-firewalld
# backend -- so there is no INPUT->sshguard jump, no native save, no boot hook.
# The default zone is the managed surface.
# ============================================================================
_fw_zone() { firewall-cmd --get-default-zone 2>/dev/null || echo public; }

# Add/remove one "PORT/PROTO" spec on the zone (permanent). A range a:b becomes
# a-b; a source CIDR turns the rule into a rich rule.
_fw_spec() {  # <add|remove> <zone> <port/proto> [cidr]
    local op="$1" z="$2" spec="$3" cidr="${4:-}" port proto flag
    port="${spec%%/*}"; proto="${spec##*/}"; [[ "$proto" == "$spec" ]] && proto=tcp
    port="${port//:/-}"
    [[ "$op" == add ]] && flag=--add || flag=--remove
    if [[ -n "$cidr" ]]; then
        firewall-cmd -q --permanent --zone="$z" \
            "${flag}-rich-rule=rule family=\"ipv4\" source address=\"$cidr\" port port=\"$port\" protocol=\"$proto\" accept" 2>/dev/null || true
    else
        firewall-cmd -q --permanent --zone="$z" "${flag}-port=$port/$proto" 2>/dev/null || true
    fi
}

# Apply each "PORT/PROTO [CIDR]" spec line from stdin to the zone.
_fw_apply_specs() {  # <add|remove> <zone>   (spec lines on stdin)
    local op="$1" z="$2" line spec cidr
    while IFS= read -r line; do
        [[ -n "$line" ]] || continue
        spec="$(awk '{print $1}' <<< "$line")"; cidr="$(awk '{print $2}' <<< "$line")"
        _fw_spec "$op" "$z" "$spec" "$cidr"
    done
}

apply_firewalld() {
    log "Detected OS: ${OS_ID} (family ${OS_FAMILY}) -- firewall backend: firewalld"
    install_firewalld
    local z; z="$(_fw_zone)"
    log "Hardening firewalld zone '$z' (deny-by-default; SSH/${SSH_PORT}${FW_SSH_SOURCE:+ from $FW_SSH_SOURCE}$([[ $FW_ALLOW_PING == 1 ]] || echo '; no ping'))..."

    # We own SSH via an explicit port/rich-rule, so drop the stock services that
    # would otherwise leave 22 (and cockpit's 9090) open on the zone.
    local s
    for s in ssh cockpit; do firewall-cmd -q --permanent --zone="$z" --remove-service="$s" 2>/dev/null || true; done

    # SSH (source-restricted -> rich rule; otherwise a plain port).
    _fw_spec add "$z" "$SSH_PORT/tcp" "$FW_SSH_SOURCE"

    # Extra ports from OPEN_PORTS (same spec grammar as the iptables path).
    if [[ -n "$OPEN_PORTS" ]]; then
        # shellcheck disable=SC2086
        _fw_apply_specs add "$z" <<< "$(_expand_args ${OPEN_PORTS//,/ })"
    fi

    # Ping: firewalld permits echo-request by default; block it only when asked.
    if [[ "$FW_ALLOW_PING" == "1" ]]; then
        firewall-cmd -q --permanent --zone="$z" --remove-icmp-block=echo-request 2>/dev/null || true
    else
        firewall-cmd -q --permanent --zone="$z" --add-icmp-block=echo-request 2>/dev/null || true
    fi

    firewall-cmd -q --reload
    log "firewalld active on zone '$z'. Add ports: 'bash $0 allow <port/proto>'   Recover: 'bash $0 disable'"
}

allow_firewalld() {
    local z lines; z="$(_fw_zone)"; lines="$(_expand_args "$@")"
    _fw_apply_specs add "$z" <<< "$lines"
    firewall-cmd -q --reload
    log "Allowed in firewalld zone '$z': $(printf '%s' "$lines" | tr '\n' ' ')"
}

deny_firewalld() {
    local z lines; z="$(_fw_zone)"; lines="$(_expand_args "$@")"
    _fw_apply_specs remove "$z" <<< "$lines"
    firewall-cmd -q --reload
    log "Removed from firewalld zone '$z': $(printf '%s' "$lines" | tr '\n' ' ')"
}

list_firewalld() {
    local z; z="$(_fw_zone)"
    echo "firewalld default zone: $z"
    firewall-cmd --list-all --zone="$z" 2>/dev/null | sed 's/^/  /' || echo "  (firewalld not running)"
}

# Recovery: re-open SSH so a misconfig can't lock you out. Leaves firewalld
# running (sshguard-firewalld keeps working). To turn the firewall off entirely:
# systemctl disable --now firewalld.
disable_firewalld() {
    local z; z="$(_fw_zone)"
    firewall-cmd -q --permanent --zone="$z" --add-service=ssh 2>/dev/null || true
    firewall-cmd -q --permanent --zone="$z" --add-port="${SSH_PORT}/tcp" 2>/dev/null || true
    firewall-cmd -q --reload || true
    warn "Re-opened SSH (service ssh + ${SSH_PORT}/tcp) on firewalld zone '$z'; firewalld left running."
    warn "Re-harden: 'bash $0 apply'   |   turn firewalld off: 'systemctl disable --now firewalld'"
}

# ============================================================================
# Subcommands
# ============================================================================
cmd_apply() {
    if [[ "$FW_BACKEND" == firewalld ]]; then apply_firewalld; return; fi

    log "Detected OS: ${OS_ID} (family ${OS_FAMILY}, init ${INIT_SYSTEM}) -- firewall backend: iptables"

    # FW_BACKEND=iptables was forced on a RHEL/Alma host where firewalld is the
    # native firewall; the two fight over nftables. Bail clearly unless overridden.
    if [[ "$OS_FAMILY" == rhel ]] && command -v firewall-cmd >/dev/null 2>&1 \
        && firewall-cmd --state >/dev/null 2>&1; then
        warn "firewalld is active but FW_BACKEND=iptables was forced; they conflict."
        warn "Prefer the firewalld backend (unset FW_BACKEND), or: systemctl disable --now firewalld"
        [[ "${FW_IGNORE_FIREWALLD:-0}" == "1" ]] || die "Refusing iptables 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>..."
    [[ "$FW_BACKEND" == firewalld ]] && { allow_firewalld "$@"; return; }
    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>..."
    [[ "$FW_BACKEND" == firewalld ]] && { deny_firewalld "$@"; return; }
    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() {
    [[ "$FW_BACKEND" == firewalld ]] && { list_firewalld; return; }
    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() {
    [[ "$FW_BACKEND" == firewalld ]] && { disable_firewalld; return; }
    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