#!/usr/bin/env bash

set -e

echo "================================="
echo "         Security Setup"
echo "================================="

read -rp "Disable root SSH login? (y/n): " DISABLE_ROOT
if [[ "$DISABLE_ROOT" == "y" ]]; then
    sed -i 's/^#PermitRootLogin.*/PermitRootLogin no/' /etc/ssh/sshd_config
fi

read -rp "Enable SSH key-only login (disable passwords)? (y/n): " KEY_ONLY
if [[ "$KEY_ONLY" == "y" ]]; then
    sed -i 's/^#PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config
fi

read -rp "Restart SSH service now? (y/n): " RESTART_SSH
if [[ "$RESTART_SSH" == "y" ]]; then
    systemctl restart ssh || systemctl restart sshd
fi

echo
read -rp "Enable UFW firewall? (y/n): " ENABLE_FIREWALL

if [[ "$ENABLE_FIREWALL" == "y" ]]; then
    apt update
    apt install -y ufw

    ufw allow OpenSSH
    ufw --force enable
fi

echo "Security setup complete"