#!/bin/bash
# =============================================================================
# SSH Hardening for Security Scanner VM
# Run this inside the VM as root
# =============================================================================
set -euo pipefail

echo "============================================"
echo " SSH Hardening - Security Scanner VM"
echo "============================================"

SSHD_CONFIG="/etc/ssh/sshd_config"
SSHD_CUSTOM="/etc/ssh/sshd_config.d/99-scanner-hardening.conf"
SSHD_BACKUP="${SSHD_CONFIG}.bak.$(date +%Y%m%d_%H%M%S)"

# --- 1. Backup original config ---
cp "${SSHD_CONFIG}" "${SSHD_BACKUP}"

# --- 2. Write hardened SSH config ---
echo "[+] Writing hardened SSH configuration..."
cat > "${SSHD_CUSTOM}" << 'EOF'
# =============================================================================
# SSH Hardening - Security Scanner VM
# =============================================================================

# --- Authentication ---
PermitRootLogin no
PasswordAuthentication no
PubkeyAuthentication yes
AuthenticationMethods publickey
PermitEmptyPasswords no
MaxAuthTries 10
MaxSessions 3
LoginGraceTime 30

# --- Ciphers ---
KexAlgorithms sntrup761x25519-sha512@openssh.com,curve25519-sha256,curve25519-sha256@libssh.org
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

# --- Network ---
Port 22
AddressFamily inet
ListenAddress 0.0.0.0
TCPKeepAlive yes
ClientAliveInterval 300
ClientAliveCountMax 2

# --- Access control ---
# Uncomment and set your admin user
# AllowUsers scanner-admin
# AllowGroups ssh-users

# --- Restrictions ---
# DisableForwarding covers X11, agent, TCP forwarding, and tunneling
DisableForwarding yes
GatewayPorts no

# --- Logging ---
SyslogFacility AUTH
LogLevel VERBOSE

# --- Misc ---
PrintMotd no
PrintLastLog yes
UseDNS no
Compression no
StrictModes yes
IgnoreRhosts yes
HostbasedAuthentication no

# --- Banner ---
Banner /etc/issue.net
EOF

# --- 3. Create scanner admin user ---
ADMIN_USER="scanner-admin"
# Also allow the cloud-init default user
CLOUDINIT_USER="${SCANNER_SSH_USER:-kai}"
echo "[+] Creating admin user: ${ADMIN_USER}..."

if ! id "${ADMIN_USER}" &>/dev/null; then
    useradd -m -s /bin/bash -G sudo,docker "${ADMIN_USER}"
    mkdir -p "/home/${ADMIN_USER}/.ssh"
    chmod 700 "/home/${ADMIN_USER}/.ssh"
    touch "/home/${ADMIN_USER}/.ssh/authorized_keys"
    chmod 600 "/home/${ADMIN_USER}/.ssh/authorized_keys"
    chown -R "${ADMIN_USER}:${ADMIN_USER}" "/home/${ADMIN_USER}/.ssh"
    echo "[!] Add your SSH public key to: /home/${ADMIN_USER}/.ssh/authorized_keys"
else
    echo "    User ${ADMIN_USER} already exists."
fi

# Set AllowUsers (idempotent: replace commented or active line)
ALLOW_USERS="${ADMIN_USER} ${CLOUDINIT_USER}"
if grep -q "^AllowUsers" "${SSHD_CUSTOM}"; then
    sed -i "s/^AllowUsers.*/AllowUsers ${ALLOW_USERS}/" "${SSHD_CUSTOM}"
elif grep -q "^# AllowUsers" "${SSHD_CUSTOM}"; then
    sed -i "s/^# AllowUsers.*/AllowUsers ${ALLOW_USERS}/" "${SSHD_CUSTOM}"
else
    echo "AllowUsers ${ALLOW_USERS}" >> "${SSHD_CUSTOM}"
fi

# --- 4. Configure fail2ban for SSH ---
echo "[+] Configuring fail2ban for SSH..."
cat > /etc/fail2ban/jail.d/sshd.conf << 'EOF'
[sshd]
enabled = true
port = ssh
filter = sshd
backend = systemd
maxretry = 3
findtime = 600
bantime = 3600
banaction = nftables[type=allports]
EOF

systemctl enable fail2ban
systemctl restart fail2ban

# --- 5. Generate new host keys (replace defaults) ---
echo "[+] Regenerating SSH host keys..."
rm -f /etc/ssh/ssh_host_*
ssh-keygen -t ed25519 -f /etc/ssh/ssh_host_ed25519_key -N ""
ssh-keygen -t rsa -b 4096 -f /etc/ssh/ssh_host_rsa_key -N ""

# Remove small DH moduli
echo "[+] Hardening DH moduli..."
awk '$5 >= 3071' /etc/ssh/moduli > /etc/ssh/moduli.safe
mv /etc/ssh/moduli.safe /etc/ssh/moduli

# --- 6. Validate and restart ---
echo "[+] Validating SSH configuration..."
if sshd -t; then
    systemctl restart ssh 2>/dev/null || systemctl restart sshd
    echo "[+] SSH hardening complete."
else
    echo "[!] SSH config validation failed. Restoring backup..."
    cp "${SSHD_BACKUP}" "${SSHD_CONFIG}"
    rm -f "${SSHD_CUSTOM}"
    exit 1
fi

echo ""
echo "[!] IMPORTANT: Before closing this session:"
echo "    1. Add your SSH public key to /home/${ADMIN_USER}/.ssh/authorized_keys"
echo "    2. Test login with: ssh ${ADMIN_USER}@<this-vm-ip>"
echo "    3. Verify you can sudo: sudo whoami"
echo ""
echo "[!] WARNING: Password login is DISABLED. Ensure key-based access works first!"