#!/bin/bash
set -e

# Ensure script is run as root
if [[ $EUID -ne 0 ]]; then
   echo "Error: This script must be run as root (use sudo)."
   exit 1
fi

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


echo "⚠️ WARNING: This script disables password authentication."
echo "Make sure you have already copied your public SSH key to this server!"
read -p "Are you sure you want to continue? (y/n) " -r response

if [[ ! "$response" =~ ^[Yy]$ ]]; then
    echo -e "\nAborting."
    exit 1
fi

# Configuration
SSH_PORT=25382
SSH_CONFIG="/etc/ssh/sshd_config"
BACKUP_FILE="${SSH_CONFIG}.backup.$(date +%F-%H%M%S)"

if [[ ! -f "$SSH_CONFIG" ]]; then
    echo "Error: $SSH_CONFIG not found."
    exit 1
fi

echo "Backing up ssh config to $BACKUP_FILE..."
cp "$SSH_CONFIG" "$BACKUP_FILE"

echo "Configuring SSH security..."

set_config() {
    local KEY=$1
    local VALUE=$2

    # Remove existing occurrences (commented or not)
    sed -i -E "/^#?$KEY\b/d" "$SSH_CONFIG"
    
    # Append the new configuration at the end
    echo "$KEY $VALUE" >> "$SSH_CONFIG"
}

set_config Port "$SSH_PORT"
set_config PermitRootLogin no
set_config PasswordAuthentication no
set_config PubkeyAuthentication yes
set_config MaxAuthTries 3
set_config PermitEmptyPasswords no
set_config X11Forwarding no
set_config AllowTcpForwarding no


echo "Testing SSH configuration..."
if ! sshd -t; then
    echo "Error: sshd config test failed. Restoring backup."
    cp "$BACKUP_FILE" "$SSH_CONFIG"
    if ! sshd -t; then
        echo "Restored config is still invalid. Please check $SSH_CONFIG manually."
    fi
    exit 1
fi

echo "Restarting SSH service..."
systemctl restart sshd || systemctl restart ssh

echo "Installing and configuring UFW..."
if ! command -v ufw >/dev/null 2>&1; then
    apt update
    DEBIAN_FRONTEND=noninteractive apt install -y ufw
fi

echo "Setting firewall rules..."
ufw default deny incoming
ufw default allow outgoing
ufw allow "$SSH_PORT"/tcp

echo "Enabling UFW..."
ufw --force enable

echo "✅ Security hardening complete."
echo "New SSH port: $SSH_PORT"
echo "Please keep your current terminal session open and test logging in via a NEW terminal window to ensure you are not locked out."