#!/bin/bash

# This script configures SSH server security settings following Mozilla's guidelines
# It allows choosing between SSH key and password authentication methods
# Reference: https://infosec.mozilla.org/guidelines/openssh

# Usage:
# 1. Save this script as "secure_ssh.sh"
# 2. Make it executable: chmod +x secure_ssh.sh
# 3. Run with root privileges: sudo ./secure_ssh.sh
# 4. Follow the interactive prompts to configure your SSH security settings

# Check if the script is run as root
if [[ "$EUID" -ne 0 ]]; then
   echo "This script must be run as root. Please use sudo to execute it."
   exit 1
fi

# Define important variables
SSHD_CONFIG="/etc/ssh/sshd_config"
BACKUP_FILE="${SSHD_CONFIG}_$(date +'%Y%m%d_%H%M%S').bak"

# Function to validate IP address
validate_ip() {
    local ip=$1
    local stat=1

    if [[ $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        OIFS=$IFS
        IFS='.'
        ip=($ip)
        IFS=$OIFS
        [[ ${ip[0]} -le 255 && ${ip[1]} -le 255 && ${ip[2]} -le 255 && ${ip[3]} -le 255 ]]
        stat=$?
    fi
    return $stat
}

# Create a backup of the current SSH configuration
cp "$SSHD_CONFIG" "$BACKUP_FILE"
echo "Current SSH configuration backed up to $BACKUP_FILE"

# Configure SSH port (mandatory)
while true; do
    read -p "Enter the SSH port to use (between 1024 and 65535): " new_port
    if [[ "$new_port" =~ ^[0-9]+$ ]] && [ "$new_port" -ge 1024 ] && [ "$new_port" -le 65535 ]; then
        port_setting="Port $new_port"
        break
    else
        echo "Error: Please enter a valid port number between 1024 and 65535"
    fi
done

# Configure IP whitelist (mandatory)
while true; do
    read -p "Enter the IP address to whitelist for SSH access (e.g., 192.168.1.100): " allowed_ip
    if validate_ip "$allowed_ip"; then
        allow_users="AllowUsers *@${allowed_ip}"
        break
    else
        echo "Error: Please enter a valid IP address"
    fi
done

# Present authentication method options to the user
echo "Choose authentication method:"
echo "1) SSH key only (more secure)"
echo "2) Password authentication"
read -p "Enter your choice (1 or 2): " auth_choice

# Handle authentication method selection
case $auth_choice in
    1)
        # Configure SSH key authentication
        auth_method="PasswordAuthentication no"
        echo "SSH key authentication selected"
        
        # Check for existing SSH key configuration
        if [ ! -f "/root/.ssh/authorized_keys" ]; then
            read -p "No SSH key found. Would you like to add one now? (yes/no): " add_key
            if [[ "$add_key" == "yes" ]]; then
                # Create SSH directory with proper permissions
                mkdir -p /root/.ssh
                read -p "Paste your public SSH key: " ssh_key
                echo "$ssh_key" >> /root/.ssh/authorized_keys
                # Set proper permissions for SSH files
                chmod 700 /root/.ssh
                chmod 600 /root/.ssh/authorized_keys
            else
                echo "Warning: No SSH key configured. You might be locked out!"
                exit 1
            fi
        fi
        ;;
    2)
        # Configure password authentication
        auth_method="PasswordAuthentication yes"
        echo "Password authentication selected"
        ;;
    *)
        echo "Invalid choice. Exiting."
        exit 1
        ;;
esac

# Create new SSH configuration file with secure settings
cat <<EOL > "$SSHD_CONFIG"
# SSH Server Configuration
# Generated by secure_ssh.sh script
# Based on Mozilla's Modern OpenSSH server configuration
# Last modified: $(date)

# Protocol version (only SSH protocol 2 is secure)
Protocol 2

# Port configuration (mandatory)
$port_setting

# IP restriction settings (mandatory)
$allow_users

# Authentication configuration
$auth_method

# Cryptographic settings
# Only modern, secure ciphers are enabled
Ciphers chacha20-poly1305@openssh.com,aes256-gcm@openssh.com,aes128-gcm@openssh.com,aes256-ctr,aes192-ctr,aes128-ctr

# Key exchange algorithms
KexAlgorithms curve25519-sha256,curve25519-sha256@libssh.org,diffie-hellman-group-exchange-sha256

# Message Authentication Codes (MACs)
MACs hmac-sha2-256-etm@openssh.com,hmac-sha2-512-etm@openssh.com,hmac-sha2-256,hmac-sha2-512

# Host keys
HostKey /etc/ssh/ssh_host_ed25519_key
HostKey /etc/ssh/ssh_host_rsa_key

# Access control settings
PermitRootLogin no
MaxAuthTries 3
LoginGraceTime 30

# Additional security measures
AllowTcpForwarding no
MaxSessions 2
LogLevel VERBOSE
ClientAliveInterval 300
ClientAliveCountMax 2
EOL

# Configure UFW firewall if installed
if command -v ufw >/dev/null 2>&1; then
    echo "Configuring UFW firewall..."
    ufw allow from "$allowed_ip" to any port "$new_port" proto tcp
    ufw status
fi

# Restart SSH service to apply new configuration
systemctl restart sshd

# Display completion message and configuration summary
echo "
SSH Configuration Summary:
-------------------------
Port: $new_port
Whitelisted IP: $allowed_ip
Authentication Method: ${auth_choice == 1 ? 'SSH Key' : 'Password'}
Backup File: $BACKUP_FILE

Important notes:
1. Keep your backup file safe
2. Test your new SSH configuration in a new session before logging out
3. If you get locked out, use the backup file to restore the previous configuration
4. The following command will be needed to connect:
   ssh -p $new_port user@server_ip

Next steps:
1. Open a new terminal
2. Try to connect using the new configuration
3. Do NOT close this session until you confirm the new configuration works
"

if [[ "$auth_choice" == "1" ]]; then
    echo "WARNING: Make sure you have a valid SSH key configured before logging out!"
fi