#!/bin/bash

# Ensure the script is run as root
if [ "$EUID" -ne 0 ]; then
    echo "Please run the script as root."
    exit 1
fi

# Prompt user for account selection
echo "For which account do you want to grant access?"
echo "1) root (by default, recommended)"
echo "2) Custom (should be added to the sudo group)"
read -rp "Enter your choice (1 or 2): " choice

# Validate choice input
if [[ "$choice" != "1" && "$choice" != "2" ]]; then
    echo "Invalid choice. Please enter 1 or 2."
    exit 1
fi

# Determine user and SSH file path
if [[ "$choice" == "1" ]]; then
    user="root"
    file="/root/.ssh/authorized_keys"
    if grep -q "PermitRootLogin no" /etc/ssh/sshd_config; then
        echo "Warning: PermitRootLogin is set to 'no' in /etc/ssh/sshd_config."
    fi
else
    read -rp "Enter the username: " user
    if [[ ! "$user" =~ ^[a-zA-Z0-9]+$ ]]; then
        echo "Invalid username. Only letters and numbers are allowed."
        exit 1
    fi
    file="/home/$user/.ssh/authorized_keys"
    if ! id "$user" &>/dev/null; then
        echo "The specified username '$user' does not exist. Is the username correct? (y/n)"
        read -rp "Enter your response: " confirm
        if [[ "$confirm" != "y" ]]; then
            exit 1
        fi
    fi
fi

# Create the authorized_keys file if it doesn't exist
if [ ! -f "$file" ]; then
    echo "File $file does not exist. Create it? (y/n)"
    read -rp "Enter your response: " create
    if [[ "$create" == "y" ]]; then
        mkdir -p "$(dirname "$file")"
        touch "$file"
        chmod 600 "$file"
        chown "$user":"$user" "$file"
    else
        exit 1
    fi
fi

# Define SSH keys
keys=(
    "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIF2vuGttz6BqWypGumLH59ZqKRlPVGw2fBV38NI6QVwo aaron@shifthosting.com"
    "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIHsnt768TxPQF849VEC+2CV7Nu3jlGo8184Qkdel55lL hi@ssh.contact"
    "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIHYJvdT4juqDnfm08n3cucHU4UDo5HfV5zK3Fso0P8xn zach-gaming@Zach-GamingPC"
)

# Add keys to the authorized_keys file
echo "Adding keys..."
for key in "${keys[@]}"; do
    if grep -qF "$key" "$file"; then
        echo "Key already exists: $key"
    else
        echo "$key" >> "$file" && echo "Added key: $key"
    fi
done

# Check SSH port configuration
ssh_port=$(grep -E "^Port\s+[0-9]+" /etc/ssh/sshd_config | awk '{print $2}')
if [ -z "$ssh_port" ]; then
    ssh_port=22
    echo "Warning: SSH port is not specified. Using default port 22."
fi

if ss -tln | grep -q ":$ssh_port "; then
    echo "OK: SSH port $ssh_port is open and listening for connections."
else
    echo "Warning: SSH port $ssh_port is not open or listening. Check your firewall settings."
fi

# Final message
echo Complete