#!/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: Location of the SSH daemon configuration file
# BACKUP_FILE: Location of the backup file with timestamp
SSHD_CONFIG="/etc/ssh/sshd_config"
BACKUP_FILE="${SSHD_CONFIG}_$(date +'%Y%m%d_%H%M%S').bak"

# Create a backup of the current SSH configuration
# This allows rollback if something goes wrong
cp "$SSHD_CONFIG" "$BACKUP_FILE"

# 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

# Configure IP restriction options
read -p "Do you want to restrict SSH access to a single IP? (yes/no): " restrict_ip
if [[ "$restrict_ip" == "yes" ]]; then
    read -p "Enter the IP address to allow SSH access: " allowed_ip
    allow_users="AllowUsers *@${allowed_ip}"
else
    allow_users="# AllowUsers configuration not set"
fi

# Configure SSH port options
read -p "Do you want to change the SSH port? (yes/no): " change_port
if [[ "$change_port" == "yes" ]]; then
    read -p "Enter the new SSH port: " new_port
    port_setting="Port $new_port"
else
    port_setting="# Port configuration not changed"
fi

# Create new SSH configuration file with secure settings
# This configuration follows Mozilla's Modern OpenSSH server recommendations
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

# 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

# IP restriction settings
$allow_users

# Port configuration
$port_setting
EOL

# Restart SSH service to apply new configuration
systemctl restart sshd

# Display completion message and warnings
echo "SSH configuration has been updated and backed up to $BACKUP_FILE"
if [[ "$auth_choice" == "1" ]]; then
    echo "WARNING: Make sure you have a valid SSH key configured before logging out!"
    echo "Your backup file is located at: $BACKUP_FILE"
fi

# Additional information
echo "
Important notes:
1. Keep your backup file ($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. For key-based authentication, ensure your public key is properly configured
"