#!/usr/bin/env bash
#
# harden-ssh.sh
#
# SSH hardening for Alpine, Debian, and Alma Linux. Run on a fresh box (and,
# for the simplex relay, BEFORE deploy-simplex.sh).
#
# All distro-specific operations go through scripts/oslib.sh. The OS-specific
# surface for this script is: package install, the sshd service name, the
# external sftp-server path, host-key keygen suppression (Alpine only), the
# sshguard log source + firewall backend, and the boot hook that installs the
# iptables jump. Each is clearly marked.
#
# What this does:
#   1. Generates fresh Ed25519 host keys; removes RSA/ECDSA/DSA host keys
#   2. Generates an Ed25519 root keypair, installs the public key into
#      /root/.ssh/authorized_keys, and PRINTS the private key to stdout once.
#   3. Forces post-quantum hybrid KEX only (mlkem768x25519, sntrup761x25519).
#   4. Modern ciphers and MACs only.
#   5. Disables everything but an interactive terminal + SFTP (no forwarding,
#      tunneling, X11, agent, password auth).
#   6. Optional non-default port (SSH_PORT).
#   7. Installs sshguard for brute-force protection.
#   8. Validates with `sshd -t` and prompts before reloading (so you don't
#      lock yourself out).
#
# A note on "quantum-safe": stock OpenSSH gives PQ KEY EXCHANGE (protects the
# session key against store-now-decrypt-later) but classical Ed25519 AUTH
# keys -- the strongest practical posture available without breaking client
# compatibility.
#
# Usage:
#   bash harden-ssh.sh                      # port stays 22
#   SSH_PORT=2222 bash harden-ssh.sh        # change port
#   ALLOWED_IP=1.2.3.4 bash harden-ssh.sh   # whitelist your client IP
#   FORCE=1 bash harden-ssh.sh              # skip the confirm prompt

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
# ============================================================================
: "${SSH_PORT:=22}"
: "${ALLOWED_IP:=}"        # optional: your client IP, sshguard-whitelisted
: "${KEY_COMMENT:=root@$(hostname)-$(date +%Y%m%d)}"
: "${FORCE:=0}"

# 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
log "Detected OS: ${OS_ID} (family ${OS_FAMILY}, init ${INIT_SYSTEM})"

# ----------------------------------------------------------------------------
# 1. Pre-flight: ensure an ssh client exists before probing its version
# ----------------------------------------------------------------------------
if ! command -v ssh >/dev/null 2>&1; then
    log "ssh not found; installing openssh..."
    install_openssh || die "Could not install OpenSSH; cannot harden. Fix the package error above, then re-run."
fi

log "Checking OpenSSH version supports PQ KEX..."
SSH_VER=$(ssh -V 2>&1 | grep -oE 'OpenSSH_[0-9]+\.[0-9]+' | head -1 | sed 's/OpenSSH_//')
SSH_MAJOR=${SSH_VER%%.*}
SSH_MINOR=${SSH_VER##*.}

# OpenSSH 9.0+ has sntrup761x25519-sha512; 9.9+ adds mlkem768x25519-sha256.
HAS_MLKEM=0
HAS_SNTRUP=0
[[ $SSH_MAJOR -gt 9 || ( $SSH_MAJOR -eq 9 && $SSH_MINOR -ge 0 ) ]] && HAS_SNTRUP=1
[[ $SSH_MAJOR -gt 9 || ( $SSH_MAJOR -eq 9 && $SSH_MINOR -ge 9 ) ]] && HAS_MLKEM=1
[[ $HAS_SNTRUP -eq 1 || $HAS_MLKEM -eq 1 ]] \
    || die "OpenSSH ${SSH_VER} has no PQ KEX. Need >= 9.0. Upgrade the base OS first."
log "OpenSSH ${SSH_VER}: ML-KEM=${HAS_MLKEM} sntrup761=${HAS_SNTRUP}"

KEX_LIST=""
[[ $HAS_MLKEM  -eq 1 ]] && KEX_LIST="mlkem768x25519-sha256"
[[ $HAS_SNTRUP -eq 1 ]] && KEX_LIST="${KEX_LIST:+$KEX_LIST,}sntrup761x25519-sha512"

# ----------------------------------------------------------------------------
# 2. Install packages (OS-gated inside oslib)
# ----------------------------------------------------------------------------
log "Installing OpenSSH server + sshguard + iptables..."
install_openssh || die "OpenSSH packages failed to install; cannot harden SSH. Fix the package error above, then re-run."
# sshguard is best-effort: a host where it can't install right now (e.g. EPEL
# momentarily unreachable) must still get the sshd_config hardening AND the login
# notifier -- not a silently half-configured box. Warn and press on.
install_bruteforce_protection \
    || warn "sshguard not installed; brute-force protection is OFF. Add it later with: dnf install -y epel-release sshguard. Continuing with the rest of the hardening."

# The external SFTP subsystem binary path differs per distro.
SFTP_PATH="$(sftp_server_path)"
[[ -x "$SFTP_PATH" ]] || warn "sftp-server not found at expected path ($SFTP_PATH); SFTP may not work until installed."

# ----------------------------------------------------------------------------
# 3. Host keys -- regenerate with Ed25519 only
# ----------------------------------------------------------------------------
log "Regenerating host keys (Ed25519 only)..."
rm -f /etc/ssh/ssh_host_rsa_key*   \
      /etc/ssh/ssh_host_ecdsa_key* \
      /etc/ssh/ssh_host_dsa_key*

if [[ ! -f /etc/ssh/ssh_host_ed25519_key ]]; then
    ssh-keygen -q -t ed25519 -f /etc/ssh/ssh_host_ed25519_key -N "" \
        -C "host@$(hostname)-$(date +%Y%m%d)"
fi
chmod 600 /etc/ssh/ssh_host_ed25519_key
chmod 644 /etc/ssh/ssh_host_ed25519_key.pub

log "Host key fingerprint (verify on first connect):"
ssh-keygen -l -f /etc/ssh/ssh_host_ed25519_key.pub | sed 's/^/    /'

# Alpine's OpenRC sshd init regenerates RSA/ECDSA keys on each start; pin off.
# No-op on systemd distros.
sshd_disable_keygen

# ----------------------------------------------------------------------------
# 4. Root user keypair
# ----------------------------------------------------------------------------
log "Generating Ed25519 keypair for root..."
mkdir -p /root/.ssh
chmod 700 /root/.ssh
touch /root/.ssh/authorized_keys
chmod 600 /root/.ssh/authorized_keys

TMP_KEY=$(mktemp -u /tmp/root_ed25519.XXXXXX)
ssh-keygen -q -t ed25519 -f "$TMP_KEY" -N "" -C "$KEY_COMMENT"
ROOT_PUB=$(cat "${TMP_KEY}.pub")
ROOT_PRIV=$(cat "$TMP_KEY")

# Idempotency: don't add the same pubkey twice.
grep -qxF "$ROOT_PUB" /root/.ssh/authorized_keys || echo "$ROOT_PUB" >> /root/.ssh/authorized_keys

# Seed root's authorized_keys with the shared admin keys from globals/ so the
# box has a known, secure default login (SSH_KEYS_URL first, else
# globals/authorized_keys). Best-effort: needs the repo's lib.sh present.
if [[ "${SEED_KEYS:-1}" == "1" && -f "$SCRIPT_DIR/lib.sh" ]]; then
    # shellcheck source=scripts/lib.sh
    . "$SCRIPT_DIR/lib.sh"
    load_globals
    if declare -f resolve_ssh_keys >/dev/null 2>&1; then
        SEEDED=0
        while IFS= read -r k; do
            [[ -n "$k" ]] || continue
            grep -qxF "$k" /root/.ssh/authorized_keys || { echo "$k" >> /root/.ssh/authorized_keys; SEEDED=$((SEEDED+1)); }
        done <<< "$(resolve_ssh_keys 2>/dev/null || true)"
        [[ "$SEEDED" -gt 0 ]] && log "Seeded ${SEEDED} admin key(s) into /root/.ssh/authorized_keys from globals."
    fi
fi

# ----------------------------------------------------------------------------
# 5. sshd_config
# ----------------------------------------------------------------------------
log "Writing /etc/ssh/sshd_config..."
[[ -f /etc/ssh/sshd_config.orig ]] || cp /etc/ssh/sshd_config /etc/ssh/sshd_config.orig

cat > /etc/ssh/sshd_config <<EOF
# Generated by harden-ssh.sh on ${OS_ID} -- $(date -u +%FT%TZ)
# Original config preserved at /etc/ssh/sshd_config.orig

Port ${SSH_PORT}
AddressFamily any
ListenAddress 0.0.0.0
ListenAddress ::
PidFile /run/sshd.pid

# --- Host key: Ed25519 only ---
HostKey /etc/ssh/ssh_host_ed25519_key

# --- Post-quantum hybrid KEX only ---
# Anything not in this list (every classical-only KEX) is rejected, which is
# what protects the session key against "store now, decrypt later".
KexAlgorithms ${KEX_LIST}

# --- Modern ciphers and MACs ---
Ciphers chacha20-poly1305@openssh.com,aes256-gcm@openssh.com,aes128-gcm@openssh.com
MACs    hmac-sha2-512-etm@openssh.com,hmac-sha2-256-etm@openssh.com,umac-128-etm@openssh.com

# --- Host key signature algorithms ---
HostKeyAlgorithms ssh-ed25519,ssh-ed25519-cert-v01@openssh.com

# --- Public key algorithms accepted from clients ---
# Ed25519 + RSA-4096 (RSA kept for older YubiKey PIV firmware pre-5.7).
PubkeyAcceptedAlgorithms ssh-ed25519,ssh-ed25519-cert-v01@openssh.com,sk-ssh-ed25519@openssh.com,sk-ssh-ed25519-cert-v01@openssh.com,rsa-sha2-512,rsa-sha2-512-cert-v01@openssh.com,rsa-sha2-256,rsa-sha2-256-cert-v01@openssh.com
RequiredRSASize 4096

# --- Authentication ---
PermitRootLogin prohibit-password
PubkeyAuthentication yes
PasswordAuthentication no
PermitEmptyPasswords no
ChallengeResponseAuthentication no
KbdInteractiveAuthentication no
UsePAM yes
AuthenticationMethods publickey
MaxAuthTries 3
MaxSessions 4
LoginGraceTime 30s
# Expose the authenticated key to the session (file at \$SSH_USER_AUTH) so the
# pam_exec login notifier can report which key was used.
ExposeAuthInfo yes

# --- Session restrictions: terminal + SFTP only, no forwarding/tunneling ---
X11Forwarding no
X11UseLocalhost yes
AllowAgentForwarding no
AllowTcpForwarding no
AllowStreamLocalForwarding no
DisableForwarding yes         # belt-and-braces: kills *all* forwarding types
GatewayPorts no
PermitTunnel no
PermitUserRC no
PermitListen none
PermitOpen none
PrintMotd no
TCPKeepAlive no
ClientAliveInterval 300
ClientAliveCountMax 2

# --- Misc ---
StrictModes yes
IgnoreRhosts yes
HostbasedAuthentication no
Compression no
Banner none

# SFTP subsystem: external sftp-server binary (path is per-distro, resolved
# by oslib's sftp_server_path). Enabled for backup retrieval, with logging to
# AUTHPRIV at INFO so file transfers are auditable.
Subsystem sftp ${SFTP_PATH} -f AUTHPRIV -l INFO

# Restrict who can SSH in. Add other users here if you create them.
AllowUsers root
EOF

# ----------------------------------------------------------------------------
# 6. Validate config
# ----------------------------------------------------------------------------
log "Validating sshd config..."
if ! sshd -t 2>/tmp/sshd-test.err; then
    cat /tmp/sshd-test.err >&2
    cp /etc/ssh/sshd_config.orig /etc/ssh/sshd_config
    die "sshd config invalid; restored original. NOT reloading."
fi
rm -f /tmp/sshd-test.err

# ----------------------------------------------------------------------------
# 7. sshguard (brute-force protection)
# ----------------------------------------------------------------------------
log "Configuring sshguard (backend + log source are OS-gated in oslib)..."
mkdir -p /etc/sshguard

WHITELIST=/etc/sshguard/whitelist
{
    echo "127.0.0.1"
    echo "::1"
    [[ -n "$ALLOWED_IP" ]] && echo "$ALLOWED_IP"
} > "$WHITELIST"

SSHGUARD_BACKEND="$(sshguard_backend)"
SSHGUARD_LOGREADER="$(sshguard_logreader)"
[[ -x "${SSHGUARD_BACKEND}" ]] || warn "sshguard backend not found at ${SSHGUARD_BACKEND}; brute-force blocking may be inactive."

cat > /etc/sshguard/sshguard.conf <<EOF
BACKEND="${SSHGUARD_BACKEND}"
${SSHGUARD_LOGREADER:+LOGREADER="${SSHGUARD_LOGREADER}"}
THRESHOLD=30
BLOCK_TIME=300
DETECTION_TIME=1800
PID_FILE=/run/sshguard.pid
WHITELIST_FILE=${WHITELIST}
EOF

# INPUT -> sshguard jump. When the host firewall (harden-firewall.sh) is enabled
# it owns the whole INPUT chain -- including this jump -- and persists it via the
# distro's native iptables package, so we install the firewall and skip the
# standalone boot hook. Otherwise fall back to the minimal init-agnostic boot
# hook that just (re)inserts the jump at every boot.
: "${ENABLE_FIREWALL:=1}"
if [[ "$ENABLE_FIREWALL" == "1" && -f "$SCRIPT_DIR/harden-firewall.sh" ]]; then
    log "Installing host firewall (deny-by-default INPUT; carries the sshguard jump)..."
    SSH_PORT="$SSH_PORT" OPEN_PORTS="${OPEN_PORTS:-}" \
    FW_SSH_SOURCE="${FW_SSH_SOURCE:-}" FW_ALLOW_PING="${FW_ALLOW_PING:-1}" \
    FORCE=1 bash "$SCRIPT_DIR/harden-firewall.sh" apply \
        || warn "harden-firewall.sh failed; INPUT left unfiltered. Re-run it manually."
else
    HOOK=$(mktemp)
    cat > "$HOOK" <<'EOF'
#!/bin/sh
# Ensure sshguard chain exists and INPUT jumps to it for the SSH port.
SSH_PORT=$(awk '/^Port / {print $2; exit}' /etc/ssh/sshd_config)
SSH_PORT=${SSH_PORT:-22}
for ipt in iptables ip6tables; do
    command -v "$ipt" >/dev/null 2>&1 || continue
    $ipt -N sshguard 2>/dev/null || true
    $ipt -C INPUT -p tcp --dport "$SSH_PORT" -j sshguard 2>/dev/null \
        || $ipt -I INPUT -p tcp --dport "$SSH_PORT" -j sshguard
done
EOF
    install_boot_hook sshguard-iptables "$HOOK"
    rm -f "$HOOK"
fi

svc_enable_start sshguard || warn "Could not start sshguard; check 'sshguard.conf' on this distro."

# ----------------------------------------------------------------------------
# 7b. Optional: SSH login notifier (pam_exec -> ntfy)
# ----------------------------------------------------------------------------
# Enabled when NTFY_URL is provided. Reports user + source IP + the key used,
# filtered by NOTIFY_GROUPS (empty here = every login on this host).
if [[ -n "${NTFY_URL:-}" ]]; then
    : "${NTFY_REGION:=$(host_region)}"
    log "Installing SSH login notifier (ntfy)..."
    install_login_notifier "$SCRIPT_DIR/ntfy-ssh-login.sh" || warn "Notifier install had issues."
fi

# ----------------------------------------------------------------------------
# 7c. Optional: daily unattended updates (set AUTO_UPDATE=1). New OS branches
# are reported, not auto-applied.
# ----------------------------------------------------------------------------
if [[ "${AUTO_UPDATE:-0}" == "1" && -f "$SCRIPT_DIR/auto-update.sh" ]]; then
    log "Scheduling daily auto-update..."
    AUTO_REBOOT="${AUTO_REBOOT:-0}" \
    ALLOW_RELEASE_UPGRADE="${ALLOW_RELEASE_UPGRADE:-0}" \
    NOTIFY="${NOTIFY:-1}" \
        bash "$SCRIPT_DIR/auto-update.sh" install || warn "Could not schedule auto-update."
fi

# ----------------------------------------------------------------------------
# 8. Enable sshd & reload (with safety prompt)
# ----------------------------------------------------------------------------
SSHD_SVC="$(sshd_service)"
log "Enabling ${SSHD_SVC} at boot..."
svc_enable "$SSHD_SVC"

# Print the private key BEFORE reloading sshd so a bad reload still leaves you
# with what you need to get back in via console.
cat <<EOF

================================================================
COPY THIS PRIVATE KEY TO YOUR CLIENT *NOW*

Save it as e.g. ~/.ssh/id_ed25519_host on your local machine,
chmod 600, and connect with:

  ssh -i ~/.ssh/id_ed25519_host -p ${SSH_PORT} root@<host>

----- BEGIN ROOT PRIVATE KEY (Ed25519) -----
${ROOT_PRIV}
----- END ROOT PRIVATE KEY -----

Public key (already in /root/.ssh/authorized_keys):
${ROOT_PUB}

Host fingerprint (verify on first connect):
$(ssh-keygen -l -f /etc/ssh/ssh_host_ed25519_key.pub)

================================================================
EOF

shred -u "$TMP_KEY" "${TMP_KEY}.pub" 2>/dev/null || rm -f "$TMP_KEY" "${TMP_KEY}.pub"

if [[ "$FORCE" != "1" ]]; then
    cat <<EOF
sshd config has passed validation. Ready to reload sshd.

If you are connected via SSH RIGHT NOW, open a SECOND session in another
terminal -- before answering yes -- to verify the new keys, port, and PQ
KEX work. If something is wrong, this reload will end your current session.

Test in another terminal first:
  ssh -i ~/.ssh/<your saved key> -p ${SSH_PORT} \\
      -o KexAlgorithms=${KEX_LIST} root@<host>

Reload sshd now? [y/N]
EOF
    read -r ans
    if [[ "${ans,,}" != "y" && "${ans,,}" != "yes" ]]; then
        warn "Skipping reload. Run 'svc reload of ${SSHD_SVC}' manually when ready."
        exit 0
    fi
fi

log "Reloading ${SSHD_SVC}..."
svc_reload "$SSHD_SVC"
log "Done. Your session, if any, should remain alive (reload preserves connections)."
log "Test from another machine before closing this session."