#!/bin/bash

# This script configures the SSH server according to Mozilla's best security practices.
# It backs up the original SSH configuration file, generates a new secure configuration file,
# and then restarts the SSH service to apply the changes.

# Usage:
# 1. Save this script as "secure_ssh.sh".
# 2. Make it executable with: `chmod +x secure_ssh.sh`.
# 3. Run it with root privileges using: `sudo ./secure_ssh.sh`.

# Variables
SSHD_CONFIG="/etc/ssh/sshd_config"
BACKUP_DIR="/etc/ssh/backup_sshd_config"
DATE_STR=$(date +"%Y%m%d_%H%M%S")
SSHD_SERVICE=""

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

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

    return $stat
}

# Function to validate the SSH port
validate_ssh_port() {
    local port=$1
    local stat=1

    if [[ "$port" =~ ^[0-9]+$ ]] && [[ "$port" -gt 0 && "$port" -le 65535 ]]; then
        stat=0
    fi

    return $stat
}

# Detect the name of the SSH service
detect_sshd_service() {
    if systemctl list-units --type=service | grep -q 'sshd'; then
        echo "sshd"
    elif systemctl list-units --type=service | grep -q 'ssh'; then
        echo "ssh"
    else
        echo ""
    fi
}

# Function to restart the SSH service
restart_sshd() {
    local service=$1
    systemctl restart "$service"
    return $?
}

# Create a backup of the SSH configuration file
backup_sshd_config() {
    mkdir -p "$BACKUP_DIR"
    BACKUP_FILE="${BACKUP_DIR}/sshd_config.${DATE_STR}.bak"
    cp "$SSHD_CONFIG" "$BACKUP_FILE"
    if [[ $? -ne 0 ]]; then
        echo "Error creating backup of the SSH configuration file."
        exit 1
    fi
    echo "SSH configuration backup created: $BACKUP_FILE"
}

# Prompt for the SSH port
prompt_ssh_port() {
    while true; do
        read -p "Please enter the SSH port (default 22): " SSH_PORT
        SSH_PORT=${SSH_PORT:-22}
        if validate_ssh_port "$SSH_PORT"; then
            echo "SSH port used: $SSH_PORT"
            break
        else
            echo "Invalid port. The port must be an integer between 1 and 65535."
        fi
    done
}

# Prompt to restrict SSH access to specific IP addresses
prompt_allowed_ips() {
    read -p "Do you want to restrict SSH access to specific IP addresses? (yes/no): " RESTRICT_ACCESS
    ALLOWED_IPS=""
    if [[ "$RESTRICT_ACCESS" =~ ^([oO]ui|[yY]es)$ ]]; then
        while true; do
            read -p "Enter a comma-separated list of allowed IP addresses (e.g. 192.168.1.1,10.0.0.1): " ALLOWED_IPS_INPUT
            IFS=',' read -ra IP_ARRAY <<< "$ALLOWED_IPS_INPUT"
            VALID=true
            ALLOWED_IPS_CLEAN=""
            for ip in "${IP_ARRAY[@]}"; do
                ip=$(echo "$ip" | xargs)  # Remove spaces
                if ! validate_ip "$ip"; then
                    echo "Invalid IP address: $ip"
                    VALID=false
                    break
                fi
                ALLOWED_IPS_CLEAN+="$ip "
            done
            if $VALID; then
                ALLOWED_IPS="${ALLOWED_IPS_CLEAN% }"  # Remove trailing space
                echo "SSH access allowed from: $ALLOWED_IPS"
                break
            else
                echo "Please enter valid IP addresses."
            fi
        done
    else
        echo "SSH access allowed from all IP addresses."
    fi
}

# Generate a new SSH configuration file
generate_new_sshd_config() {
    cat <<EOL > "$SSHD_CONFIG"
# Secure SSH configuration according to Mozilla's recommendations

# Protocol version
Protocol 2

# Secure ciphers
Ciphers chacha20-poly1305@openssh.com,aes256-gcm@openssh.com,aes128-gcm@openssh.com,aes256-ctr,aes192-ctr,aes128-ctr

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

# Secure MAC algorithms
MACs hmac-sha2-256-etm@openssh.com,hmac-sha2-512-etm@openssh.com,hmac-sha2-256,hmac-sha2-512

# Disable password authentication
PasswordAuthentication no

# Secure host keys (make sure they exist)
HostKey /etc/ssh/ssh_host_ed25519_key
HostKey /etc/ssh/ssh_host_rsa_key

# Restrict root access (root will not be able to connect)
PermitRootLogin no

# Strict login policy
MaxAuthTries 3
LoginGraceTime 30

# Specified SSH port
Port $SSH_PORT

# Other recommended configurations
PermitEmptyPasswords no
ChallengeResponseAuthentication no
UsePAM yes
AllowTcpForwarding no
X11Forwarding no

# IP address access restrictions (if specified)
EOL

    if [[ -n "$ALLOWED_IPS" ]]; then
        echo "" >> "$SSHD_CONFIG"
        echo "# IP address access restriction" >> "$SSHD_CONFIG"
        echo "Match Address $(echo "$ALLOWED_IPS" | tr ' ' ',')" >> "$SSHD_CONFIG"
        echo "    AllowUsers *" >> "$SSHD_CONFIG"
        echo "Match Address *,!$(echo "$ALLOWED_IPS" | tr ' ' ',')" >> "$SSHD_CONFIG"
        echo "    DenyUsers *" >> "$SSHD_CONFIG"
    fi
}

# Test the syntax of the new SSH configuration
test_sshd_config() {
    sshd -t
    return $?
}

# Detect the SSH service in use
detect_service() {
    SSHD_SERVICE=$(detect_sshd_service)
    if [[ -z "$SSHD_SERVICE" ]]; then
        echo "Error: Unable to detect the SSH service (sshd or ssh)."
        exit 1
    fi
}

# Main process
main() {
    # Warning message before proceeding with the script.
    echo "Warning: Before running this script, make sure you have created a user and an SSH key in the authorized_keys file."
    echo "The root user will not be able to connect."

    read -p "Do you want to continue? (yes/no): " CONTINUE

    if [[ ! "$CONTINUE" =~ ^([yY]es|[oO]ui)$ ]]; then
        echo "Exiting script."
        exit 0
    fi

    # Detect the SSH service
    detect_service

    # Backup current configuration
    backup_sshd_config

    # Prompt for the SSH port
    prompt_ssh_port

    # Prompt for allowed IP addresses
    prompt_allowed_ips

    # Generate the new configuration file
    generate_new_sshd_config
    echo "New SSH configuration file generated."

    # Test the syntax of the new configuration
    if test_sshd_config; then
        echo "The new SSH configuration is valid."
    else
        echo "Error: The new SSH configuration contains errors. Restoring original configuration."
        cp "$BACKUP_FILE" "$SSHD_CONFIG"
        restart_sshd "$SSHD_SERVICE"
        exit 1
    fi

    # Restart the SSH service
    if restart_sshd "$SSHD_SERVICE"; then
        echo "The SSH service has been restarted successfully."
        echo "The SSH configuration has been updated according to Mozilla's security recommendations."
    else
        echo "Error: Failed to restart the SSH service. Restoring original configuration."
        cp "$BACKUP_FILE" "$SSHD_CONFIG"
        restart_sshd "$SSHD_SERVICE"
        exit 1
    fi
}

# Execute the main process
main