#!/usr/bin/env bash
#
# harden-ssh.sh
#
# SSH hardening for Alpine Linux. Run BEFORE deploy-simplex.sh on a fresh box.
#
# What this does:
#   1. Generates fresh Ed25519 host keys; removes RSA/ECDSA/DSA host keys
#   2. Generates an Ed25519 root keypair on the host, installs the public key
#      into /root/.ssh/authorized_keys, and PRINTS the private key to stdout
#      so you can copy it to your client. THIS IS YOUR ONLY CHANCE TO COPY IT.
#   3. Forces post-quantum hybrid KEX only:
#        mlkem768x25519-sha256        (the future default, NIST ML-KEM hybrid)
#        sntrup761x25519-sha512       (older PQ KEX, kept as fallback)
#      Drops every classical-only KEX. Connections that don't speak PQ KEX
#      will be rejected.
#   4. Modern ciphers and MACs only (chacha20-poly1305, aes256-gcm,
#      hmac-sha2-512-etm)
#   5. Disables everything not needed for an interactive terminal:
#      - password auth, root password login (key-only)
#      - challenge-response, GSSAPI, PAM, host-based auth
#      - X11 forwarding, agent forwarding
#      - TCP forwarding, stream-local forwarding (UNIX sockets)
#      - tunneling (PermitTunnel), gateway ports
#      - SFTP subsystem (kept ON — needed for backup retrieval)
#      - empty passwords, .ssh/rc execution, compression
#      Result: a session can run a shell. That's it. No -L, no -R, no -D, no
#      jump hosting, no sftp, no scp.
#   6. Optional non-default port (-p PORT)
#   7. Installs sshguard with iptables backend for brute-force protection
#   8. Validates config with `sshd -t` and prompts for confirmation before
#      reloading sshd (so a config error or a typo doesn't lock you out)
#
# A note on "quantum-safe":
#   Stock OpenSSH provides PQ KEY EXCHANGE (the session key, the thing that
#   matters for "store now, decrypt later"). It does NOT yet provide PQ
#   AUTHENTICATION KEYS -- there is no standardized PQ host or user key
#   algorithm in mainline OpenSSH yet. So:
#     - Your session is PQ-protected against SNDL: yes
#     - Your auth keypair (Ed25519) is classical: yes, and that's the best
#       practical choice today. PQ signature support exists only in the
#       open-quantum-safe/openssh fork, which breaks compatibility with
#       every standard SSH client.
#   This script gives you the strongest stock-OpenSSH posture available.
#
# Usage:
#   bash harden-ssh.sh                      # port stays 22, default
#   SSH_PORT=2222 bash harden-ssh.sh        # change port
#   ALLOWED_IP=1.2.3.4 bash harden-ssh.sh   # whitelist your client IP in sshguard
#   FORCE=1 bash harden-ssh.sh              # skip the "are you sure" prompt

set -euo pipefail

# ============================================================================
# CONFIG
# ============================================================================
: "${SSH_PORT:=22}"
: "${ALLOWED_IP:=}"        # optional: your client IP, will be sshguard-whitelisted
: "${KEY_COMMENT:=root@$(hostname)-$(date +%Y%m%d)}"
: "${FORCE:=0}"

log()  { printf '\033[1;32m[+]\033[0m %s\n' "$*"; }
warn() { printf '\033[1;33m[!]\033[0m %s\n' "$*" >&2; }
die()  { printf '\033[1;31m[x]\033[0m %s\n' "$*" >&2; exit 1; }

[[ $EUID -eq 0 ]] || die "Run as root."
[[ -f /etc/alpine-release ]] || die "This script targets Alpine Linux."

# ----------------------------------------------------------------------------
# 1. Pre-flight checks
# ----------------------------------------------------------------------------
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.
# OpenSSH 9.9+ also has mlkem768x25519-sha256.
HAS_MLKEM=0
HAS_SNTRUP=0
if [[ $SSH_MAJOR -gt 9 || ( $SSH_MAJOR -eq 9 && $SSH_MINOR -ge 0 ) ]]; then
    HAS_SNTRUP=1
fi
if [[ $SSH_MAJOR -gt 9 || ( $SSH_MAJOR -eq 9 && $SSH_MINOR -ge 9 ) ]]; then
    HAS_MLKEM=1
fi
[[ $HAS_SNTRUP -eq 1 || $HAS_MLKEM -eq 1 ]] \
    || die "OpenSSH ${SSH_VER} has no PQ KEX. Need >= 9.0. Upgrade Alpine first."

log "OpenSSH ${SSH_VER}: ML-KEM=${HAS_MLKEM} sntrup761=${HAS_SNTRUP}"

# Build the KEX list from what's actually available.
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
# ----------------------------------------------------------------------------
log "Installing openssh-server-pam, sshguard, iptables..."
# openssh-server-pam replaces openssh-server (PAM-enabled sshd). If the
# non-pam version was installed earlier, swap it out cleanly.
if apk info -e openssh-server >/dev/null 2>&1 && \
   ! apk info -e openssh-server-pam >/dev/null 2>&1; then
    apk del -q openssh-server || true
fi
apk add -q openssh openssh-server-pam linux-pam sshguard iptables ip6tables openrc

# ----------------------------------------------------------------------------
# 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*

# Keep existing ed25519 key if there is one (so the host fingerprint doesn't
# change unnecessarily on re-runs). Generate one if not.
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/^/    /'

# Stop Alpine's sshd init from regenerating RSA/ECDSA keys on every start.
# /etc/conf.d/sshd: pin sshd_disable_keygen=no but only generate ed25519 by
# overriding the keygen line in the init script via a drop-in.
if [[ -f /etc/conf.d/sshd ]]; then
    if ! grep -q '^sshd_disable_keygen=' /etc/conf.d/sshd; then
        echo 'sshd_disable_keygen="yes"' >> /etc/conf.d/sshd
    else
        sed -i 's/^sshd_disable_keygen=.*/sshd_disable_keygen="yes"/' /etc/conf.d/sshd
    fi
fi

# ----------------------------------------------------------------------------
# 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

# Always create a brand-new key pair in a temp location so we can show the
# private key to the user and then add the public key to 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.
if ! grep -qxF "$ROOT_PUB" /root/.ssh/authorized_keys; then
    echo "$ROOT_PUB" >> /root/.ssh/authorized_keys
fi

# ----------------------------------------------------------------------------
# 5. sshd_config
# ----------------------------------------------------------------------------
log "Writing /etc/ssh/sshd_config..."
# Back up whatever was there before, once.
[[ -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 -- $(date -u +%FT%TZ)
# Original config preserved at /etc/ssh/sshd_config.orig

Port ${SSH_PORT}
AddressFamily any
ListenAddress 0.0.0.0
ListenAddress ::

# --- 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. This
# protects against "store now, decrypt later" because the session key is
# derived from a PQ KEM hybrid.
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 (for the host proving itself) ---
HostKeyAlgorithms ssh-ed25519,ssh-ed25519-cert-v01@openssh.com

# --- Public key algorithms accepted from clients ---
# Ed25519 + RSA-4096 (RSA kept for older YubiKey firmware that lacks Ed25519
# in the PIV applet -- pre-5.7). Plain ssh-rsa (SHA-1) is NOT included.
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
# PAM is enabled so the session stack runs on every login (used by
# pam_exec.so for post-login notification hooks). Auth still requires a
# pubkey -- PAM only handles account/session, not auth, because
# PasswordAuthentication and KbdInteractiveAuthentication are off.
UsePAM yes
AuthenticationMethods publickey
MaxAuthTries 3
MaxSessions 4
LoginGraceTime 30s

# --- Session restrictions: terminal access only, no forwarding/tunneling ---
# This server runs SimpleX relays. The only legitimate SSH use is an admin
# terminal session. Every forwarding feature below is therefore disabled --
# they're just attack surface and lateral-movement aids if the host is ever
# compromised.
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               # don't run ~/.ssh/rc on login
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: enabled for backup retrieval only.
# Note: this allows sftp and scp to work. Only the forwarding/tunneling
# features above are disabled -- file transfer via SFTP is legitimate.
Subsystem sftp internal-sftp

# 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 with iptables backend..."
mkdir -p /etc/sshguard

# Whitelist: localhost always, plus optional caller-supplied IP.
WHITELIST=/etc/sshguard/whitelist
{
    echo "127.0.0.1"
    echo "::1"
    [[ -n "$ALLOWED_IP" ]] && echo "$ALLOWED_IP"
} > "$WHITELIST"

# sshguard.conf: explicitly point it at iptables backend.
cat > /etc/sshguard/sshguard.conf <<EOF
BACKEND="/usr/libexec/sshg-fw-iptables"
LOGREADER="LANG=C journalctl -afb -p info -n1 -u sshd -o cat"
THRESHOLD=30
BLOCK_TIME=300
DETECTION_TIME=1800
PID_FILE=/run/sshguard.pid
WHITELIST_FILE=${WHITELIST}
EOF

# sshguard inserts rules into chain "sshguard"; we need a jump from INPUT.
# awall doesn't manage this chain, so we add it once and make it persist via
# a small startup hook.
cat > /etc/local.d/sshguard-iptables.start <<'EOF'
#!/bin/sh
# Ensure sshguard chain exists and INPUT jumps to it for tcp/22 (and PORT).
SSH_PORT=$(awk '/^Port / {print $2; exit}' /etc/ssh/sshd_config)
SSH_PORT=${SSH_PORT:-22}
for ipt in iptables ip6tables; do
    $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
chmod +x /etc/local.d/sshguard-iptables.start
rc-update add local default 2>/dev/null || true
/etc/local.d/sshguard-iptables.start

rc-update add sshguard default
rc-service sshguard restart || rc-service sshguard start

# ----------------------------------------------------------------------------
# 8. SSHD enable & reload (with safety prompt)
# ----------------------------------------------------------------------------
log "Enabling sshd at boot..."
rc-update add sshd default

# Print the private key BEFORE reloading sshd so even if reload locks the
# user out, they have what they need to come back in via console.
cat <<EOF

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

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

  ssh -i ~/.ssh/id_ed25519_simplex -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

# Wipe the temp files holding the private key.
shred -u "$TMP_KEY" "${TMP_KEY}.pub" 2>/dev/null || rm -f "$TMP_KEY" "${TMP_KEY}.pub"

# Final guard: confirm before reloading sshd. A bad reload is recoverable from
# console; a bad reload while you assumed everything was fine is not.
if [[ "$FORCE" != "1" ]]; then
    cat <<EOF
sshd config has passed validation. Ready to reload sshd.

If you are connected via SSH RIGHT NOW, opening a SECOND ssh session in
another terminal -- before answering yes -- to verify the new keys, port,
and PQ KEX work is the safest path. If your new key/port/KEX 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 'rc-service sshd reload' manually when ready."
        exit 0
    fi
fi

log "Reloading sshd..."
rc-service sshd reload || rc-service sshd restart

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