#!/bin/bash

# Debian 12 Initial Setup Script
# Run as root: bash setup.sh

set -euo pipefail  # Exit on error, undefined vars, pipe failures

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

### === SETTINGS === ###
SYSADMIN_USER="sysadmin"
CREATE_ADDITIONAL_USER=""
ADDITIONAL_USER=""
USE_UFW=""
USE_FAIL2BAN=""
DISABLE_SSH_PASSWORD=""
SETUP_SSH_KEYS=""

# Logging functions
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[WARNING] $1${NC}"
}

error() {
    echo -e "${RED}[ERROR] $1${NC}"
    exit 1
}

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
       error "This script must be run as root"
    fi
}

# Check if running on Debian
check_debian() {
    if ! grep -q "Debian" /etc/os-release; then
        error "This script is designed for Debian systems"
    fi
}

# Check for required commands
check_commands() {
    # Check useradd (in /usr/sbin/)
    if ! command -v useradd >/dev/null 2>&1 && ! [ -x /usr/sbin/useradd ]; then
        error "useradd command not found. Please install the passwd package: apt install -y passwd"
    fi
    
    # Check passwd (usually in /usr/bin/)
    if ! command -v passwd >/dev/null 2>&1 && ! [ -x /usr/bin/passwd ]; then
        error "passwd command not found. Please install the passwd package: apt install -y passwd"
    fi
    
    # Check usermod (in /usr/sbin/)
    if ! command -v usermod >/dev/null 2>&1 && ! [ -x /usr/sbin/usermod ]; then
        warn "usermod command not found. Will attempt to use full path."
    fi
    
    # Check sudo (usually in /usr/bin/)
    if ! command -v sudo >/dev/null 2>&1 && ! [ -x /usr/bin/sudo ]; then
        warn "sudo command not found. Will install it during system setup."
    fi
}

# Ask user about creating an additional user
ask_additional_user() {
    echo ""
    echo -e "${BLUE}=== Additional User Creation ===${NC}"
    echo "This script will create the 'sysadmin' user by default."
    echo "You can also create an additional user account if needed."
    echo ""
    read -p "Would you like to create an additional user account? (y/N): " create_user_choice
    
    if [[ "$create_user_choice" =~ ^[Yy]$ ]]; then
        CREATE_ADDITIONAL_USER="yes"
        echo ""
        echo -e "${YELLOW}Enter username for the additional user:${NC}"
        echo "(This user will also have sudo privileges)"
        echo ""
        while true; do
            read -p "Username: " username_input
            
            # Validate username
            if [[ -z "$username_input" ]]; then
                echo -e "${RED}Username cannot be empty. Please try again.${NC}"
                continue
            elif [[ ! "$username_input" =~ ^[a-z][a-z0-9_-]*$ ]]; then
                echo -e "${RED}Invalid username. Use lowercase letters, numbers, hyphens, and underscores only.${NC}"
                echo -e "${RED}Username must start with a letter.${NC}"
                continue
            elif [[ "$username_input" == "$SYSADMIN_USER" ]]; then
                echo -e "${RED}Username cannot be the same as sysadmin user. Please choose a different name.${NC}"
                continue
            elif id "$username_input" &>/dev/null; then
                echo -e "${RED}User '$username_input' already exists. Please choose a different name.${NC}"
                continue
            else
                ADDITIONAL_USER="$username_input"
                log "Will create additional user: $ADDITIONAL_USER"
                break
            fi
        done
    else
        CREATE_ADDITIONAL_USER="no"
        log "Skipping additional user creation"
    fi
}

# Ask user about SSH security preferences
ask_ssh_security() {
    echo ""
    echo -e "${BLUE}=== SSH Security Configuration ===${NC}"
    echo "This script can disable SSH password authentication for enhanced security."
    echo ""
    echo -e "${YELLOW}⚠️  WARNING: Disabling password authentication means you can ONLY log in with SSH keys!${NC}"
    echo "If you lose your SSH keys, you could be locked out of the server."
    echo ""
    echo -e "${YELLOW}Consider keeping password authentication if:${NC}"
    echo "• This is your first time setting up SSH keys"
    echo "• You don't have a reliable way to store SSH keys"
    echo "• You need emergency access options"
    echo ""
    read -p "Disable SSH password authentication? (y/N): " disable_password_choice
    
    if [[ "$disable_password_choice" =~ ^[Yy]$ ]]; then
        DISABLE_SSH_PASSWORD="yes"
        log "SSH password authentication will be disabled"
        
        echo ""
        echo -e "${BLUE}=== SSH Key Setup ===${NC}"
        echo "Since password authentication will be disabled, you'll need SSH keys to connect."
        echo "You need to add YOUR computer's SSH public key to the server."
        echo ""
        read -p "Add your SSH public key to authorized_keys for login? (Y/n): " setup_keys_choice
        
        if [[ "$setup_keys_choice" =~ ^[Nn]$ ]]; then
            SETUP_SSH_KEYS="no"
            warn "SSH keys will NOT be set up automatically"
            warn "You MUST manually add your public key to authorized_keys files after setup"
        else
            SETUP_SSH_KEYS="yes"
            log "You will be prompted to paste your SSH public key"
        fi
    else
        DISABLE_SSH_PASSWORD="no"
        SETUP_SSH_KEYS="no"
        log "SSH password authentication will remain enabled"
        
        echo ""
        read -p "Still want to add your SSH public key for convenient access? (Y/n): " add_keys_choice
        if [[ ! "$add_keys_choice" =~ ^[Nn]$ ]]; then
            SETUP_SSH_KEYS="yes"
            log "You will be prompted to paste your SSH public key"
        fi
    fi
}

# Ask user about firewall preferences
ask_firewall_preferences() {
    echo ""
    echo -e "${BLUE}=== Firewall Configuration ===${NC}"
    echo "This script can configure UFW (Uncomplicated Firewall) or you can manage iptables manually."
    echo ""
    echo -e "${YELLOW}UFW vs iptables:${NC}"
    echo "• UFW: Easy to use, good for basic setups"
    echo "• Manual iptables: More control, better for complex setups (VPN servers, etc.)"
    echo ""
    read -p "Would you like to install and configure UFW? (y/N): " ufw_choice
    
    if [[ "$ufw_choice" =~ ^[Yy]$ ]]; then
        USE_UFW="yes"
        log "UFW will be installed and configured"
    else
        USE_UFW="no"
        log "UFW will be skipped - you can configure iptables manually"
    fi
    
    echo ""
    echo -e "${BLUE}=== Fail2ban Configuration ===${NC}"
    echo "Fail2ban protects against brute-force attacks by monitoring logs and banning IPs."
    echo ""
    echo -e "${YELLOW}Consider skipping fail2ban if:${NC}"
    echo "• You're running a VPN server (WireGuard, OpenVPN, etc.)"
    echo "• You have complex iptables rules"
    echo "• You prefer managing IP banning manually"
    echo ""
    read -p "Would you like to install and configure fail2ban? (Y/n): " fail2ban_choice
    
    if [[ "$fail2ban_choice" =~ ^[Nn]$ ]]; then
        USE_FAIL2BAN="no"
        log "Fail2ban will be skipped"
    else
        USE_FAIL2BAN="yes"
        log "Fail2ban will be installed and configured"
    fi
}

### === DEBIAN SYSTEM SETUP === ###
setup_system() {
    log "Starting Debian 12 initial setup..."
    
    # Update system packages
    log "Updating system packages..."
    apt update && apt upgrade -y
    
    # Install essential packages
    log "Installing essential packages..."
    
    # Build package list based on user preferences
    local packages="sudo curl wget vim htop unzip git openssh-server"
    
    if [[ "$USE_UFW" == "yes" ]]; then
        packages="$packages ufw"
    fi
    
    if [[ "$USE_FAIL2BAN" == "yes" ]]; then
        packages="$packages fail2ban"
    fi
    
    log "Installing packages: $packages"
    apt install -y $packages
}

create_sysadmin_user() {
    log "Creating user: $SYSADMIN_USER"
    
    if id "$SYSADMIN_USER" &>/dev/null; then
        warn "User $SYSADMIN_USER already exists, skipping creation"
    else
        # Create user with home directory
        if command -v useradd >/dev/null 2>&1; then
            useradd -m -s /bin/bash "$SYSADMIN_USER"
        elif [ -x /usr/sbin/useradd ]; then
            /usr/sbin/useradd -m -s /bin/bash "$SYSADMIN_USER"
        else
            error "useradd command not found. Please install the passwd package."
        fi
        log "User $SYSADMIN_USER created successfully"
        
        # Set password for sysadmin user
        echo "Please set a password for user $SYSADMIN_USER:"
        if command -v passwd >/dev/null 2>&1; then
            passwd "$SYSADMIN_USER"
        elif [ -x /usr/bin/passwd ]; then
            /usr/bin/passwd "$SYSADMIN_USER"
        else
            error "passwd command not found. Please install the passwd package."
        fi
    fi
    
    # Add sysadmin to sudo group
    log "Adding $SYSADMIN_USER to sudo group..."
    if command -v usermod >/dev/null 2>&1; then
        usermod -aG sudo "$SYSADMIN_USER"
    elif [ -x /usr/sbin/usermod ]; then
        /usr/sbin/usermod -aG sudo "$SYSADMIN_USER"
    else
        error "usermod command not found. Please install the passwd package."
    fi
}

create_additional_user() {
    if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
        log "Creating additional user: $ADDITIONAL_USER"
        
        if id "$ADDITIONAL_USER" &>/dev/null; then
            warn "User $ADDITIONAL_USER already exists, skipping creation"
        else
            # Create user with home directory
            if command -v useradd >/dev/null 2>&1; then
                useradd -m -s /bin/bash "$ADDITIONAL_USER"
            elif [ -x /usr/sbin/useradd ]; then
                /usr/sbin/useradd -m -s /bin/bash "$ADDITIONAL_USER"
            else
                error "useradd command not found. Please install the passwd package."
            fi
            log "User $ADDITIONAL_USER created successfully"
            
            # Set password for additional user
            echo "Please set a password for user $ADDITIONAL_USER:"
            if command -v passwd >/dev/null 2>&1; then
                passwd "$ADDITIONAL_USER"
            elif [ -x /usr/bin/passwd ]; then
                /usr/bin/passwd "$ADDITIONAL_USER"
            else
                error "passwd command not found. Please install the passwd package."
            fi
        fi
        
        # Add additional user to sudo group
        log "Adding $ADDITIONAL_USER to sudo group..."
        if command -v usermod >/dev/null 2>&1; then
            usermod -aG sudo "$ADDITIONAL_USER"
        elif [ -x /usr/sbin/usermod ]; then
            /usr/sbin/usermod -aG sudo "$ADDITIONAL_USER"
        else
            error "usermod command not found. Please install the passwd package."
        fi
    fi
}

configure_security() {
    # Disable root SSH login
    log "Disabling root SSH login..."
    sed -i 's/#PermitRootLogin yes/PermitRootLogin no/' /etc/ssh/sshd_config
    sed -i 's/PermitRootLogin yes/PermitRootLogin no/' /etc/ssh/sshd_config
    
    # Configure SSH settings (remove old custom settings first)
    log "Applying SSH security settings..."
    
    # Remove previous custom settings to avoid duplicates
    sed -i '/^Protocol 2$/d' /etc/ssh/sshd_config
    sed -i '/^PasswordAuthentication /d' /etc/ssh/sshd_config
    sed -i '/^PubkeyAuthentication /d' /etc/ssh/sshd_config
    sed -i '/^PermitEmptyPasswords /d' /etc/ssh/sshd_config
    sed -i '/^X11Forwarding /d' /etc/ssh/sshd_config
    sed -i '/^MaxAuthTries /d' /etc/ssh/sshd_config
    sed -i '/^ClientAliveInterval /d' /etc/ssh/sshd_config
    sed -i '/^ClientAliveCountMax /d' /etc/ssh/sshd_config
    sed -i '/^AllowUsers /d' /etc/ssh/sshd_config
    
    # Add current settings
    {
        echo "Protocol 2"
        if [[ "$DISABLE_SSH_PASSWORD" == "yes" ]]; then
            echo "PasswordAuthentication no"
            log "SSH password authentication disabled"
        else
            echo "PasswordAuthentication yes"
            log "SSH password authentication enabled"
        fi
        echo "PubkeyAuthentication yes"
        echo "PermitEmptyPasswords no"
        echo "X11Forwarding no"
        echo "MaxAuthTries 3"
        echo "ClientAliveInterval 300"
        echo "ClientAliveCountMax 2"
        if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
            echo "AllowUsers $SYSADMIN_USER $ADDITIONAL_USER"
        else
            echo "AllowUsers $SYSADMIN_USER"
        fi
    } >> /etc/ssh/sshd_config
    
    # Remove/lock root password
    log "Locking root password..."
    passwd -l root
    
    # Configure firewall based on user choice
    if [[ "$USE_UFW" == "yes" ]]; then
        configure_ufw
    else
        configure_basic_iptables
    fi
    
    # Configure fail2ban if requested
    if [[ "$USE_FAIL2BAN" == "yes" ]]; then
        configure_fail2ban
    fi
    
    # Set up automatic security updates
    log "Configuring automatic security updates..."
    apt install -y unattended-upgrades apt-listchanges
    
    # Configure unattended upgrades
    cat > /etc/apt/apt.conf.d/50unattended-upgrades << EOF
Unattended-Upgrade::Allowed-Origins {
    "\${distro_id}:\${distro_codename}-security";
    "\${distro_id}ESMApps:\${distro_codename}-apps-security";
    "\${distro_id}ESM:\${distro_codename}-infra-security";
};

Unattended-Upgrade::AutoFixInterruptedDpkg "true";
Unattended-Upgrade::MinimalSteps "true";
Unattended-Upgrade::Remove-Unused-Dependencies "true";
Unattended-Upgrade::Automatic-Reboot "false";
EOF
    
    # Enable automatic updates
    echo 'APT::Periodic::Update-Package-Lists "1";' > /etc/apt/apt.conf.d/20auto-upgrades
    echo 'APT::Periodic::Unattended-Upgrade "1";' >> /etc/apt/apt.conf.d/20auto-upgrades
    
    # Set timezone (optional)
    log "Current timezone: $(timedatectl show --property=Timezone --value)"
    echo "If you want to change timezone, run: timedatectl set-timezone Your/Timezone"
}

configure_ufw() {
    log "Configuring UFW firewall..."
    
    # Reset UFW to defaults
    ufw --force reset
    
    # Set default policies
    ufw default deny incoming
    ufw default allow outgoing
    
    # Allow SSH (port 22)
    ufw allow ssh
    ufw allow 22/tcp
    
    # Allow HTTP (port 80)
    ufw allow http
    ufw allow 80/tcp
    
    # Allow HTTPS (port 443)
    ufw allow https
    ufw allow 443/tcp
    
    # Enable UFW
    ufw --force enable
    
    log "UFW status:"
    ufw status verbose
}

configure_basic_iptables() {
    log "Setting up basic iptables rules..."
    
    # Create a simple script for basic iptables rules
    cat > /etc/iptables-basic.sh << 'EOF'
#!/bin/bash
# Basic iptables rules for SSH, HTTP, HTTPS

# Flush existing rules
iptables -F
iptables -X
iptables -t nat -F
iptables -t nat -X
iptables -t mangle -F
iptables -t mangle -X

# Set default policies
iptables -P INPUT DROP
iptables -P FORWARD DROP
iptables -P OUTPUT ACCEPT

# Allow loopback
iptables -A INPUT -i lo -j ACCEPT
iptables -A OUTPUT -o lo -j ACCEPT

# Allow established and related connections
iptables -A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT

# Allow SSH (port 22)
iptables -A INPUT -p tcp --dport 22 -j ACCEPT

# Allow HTTP (port 80)
iptables -A INPUT -p tcp --dport 80 -j ACCEPT

# Allow HTTPS (port 443)
iptables -A INPUT -p tcp --dport 443 -j ACCEPT

# Save rules
iptables-save > /etc/iptables/rules.v4
EOF
    
    chmod +x /etc/iptables-basic.sh
    
    # Create iptables directory if it doesn't exist
    mkdir -p /etc/iptables
    
    # Run the basic rules script
    /etc/iptables-basic.sh
    
    # Install iptables-persistent to make rules persistent
    apt install -y iptables-persistent
    
    log "Basic iptables rules applied and saved"
    warn "You can customize /etc/iptables-basic.sh and run it to update rules"
    warn "Current rules are saved in /etc/iptables/rules.v4"
}

configure_fail2ban() {
    if [[ "$USE_FAIL2BAN" == "yes" ]]; then
        log "Configuring fail2ban..."
        cat > /etc/fail2ban/jail.local << EOF
[DEFAULT]
bantime = 1h
findtime = 10m
maxretry = 3

[sshd]
enabled = true
port = ssh
filter = sshd
logpath = /var/log/auth.log
maxretry = 3
bantime = 1h
EOF
        
        # Start and enable fail2ban
        systemctl enable fail2ban
        systemctl start fail2ban
        
        log "Fail2ban configured and started"
    else
        # Remove/disable fail2ban if user chose not to use it
        if systemctl is-active --quiet fail2ban; then
            log "Stopping and disabling fail2ban..."
            systemctl stop fail2ban
            systemctl disable fail2ban
        fi
        
        if dpkg -l | grep -q fail2ban; then
            log "Removing fail2ban package..."
            apt remove --purge -y fail2ban
            apt autoremove -y
        fi
        
        log "Fail2ban has been removed/disabled"
    fi
}


### === SSH KEY SETUP === ###
setup_user_ssh_keys() {
  if [[ "$SETUP_SSH_KEYS" == "no" ]]; then
    log "SSH key setup skipped"
    return
  fi
  
  echo ""
  echo -e "${BLUE}=== SSH Public Key Setup ===${NC}"
  echo "To set up SSH key authentication, you need to provide your SSH public key."
  echo "You can find your public key on your computer by running:"
  echo -e "${YELLOW}  cat ~/.ssh/id_rsa.pub${NC}"
  echo -e "${YELLOW}  # or${NC}"
  echo -e "${YELLOW}  cat ~/.ssh/id_ed25519.pub${NC}"
  echo ""
  echo "If you don't have an SSH key pair, generate one on your computer first:"
  echo -e "${YELLOW}  ssh-keygen -t ed25519 -C \"your_email@example.com\"${NC}"
  echo ""
  
  while true; do
    echo "Please paste your SSH public key (starts with ssh-rsa, ssh-ed25519, etc.):"
    read -r user_public_key
    
    if [[ -z "$user_public_key" ]]; then
      echo -e "${RED}Public key cannot be empty. Please try again.${NC}"
      continue
    elif [[ "$user_public_key" =~ ^(ssh-rsa|ssh-dss|ssh-ed25519|ecdsa-sha2-) ]]; then
      break
    else
      echo -e "${RED}Invalid SSH public key format. Please ensure you copied the entire key.${NC}"
      continue
    fi
  done
  
  # Add SSH key to sysadmin user
  setup_ssh_key_for_user "$SYSADMIN_USER" "$user_public_key"
  
  # Add SSH key to additional user if created
  if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
    setup_ssh_key_for_user "$ADDITIONAL_USER" "$user_public_key"
  fi
}

setup_ssh_key_for_user() {
  local username="$1"
  local public_key="$2"
  local user_home="/home/$username"
  local ssh_dir="$user_home/.ssh"
  local authorized_keys="$ssh_dir/authorized_keys"
  
  log "Setting up SSH key for user: $username"
  
  # Create .ssh directory for user if it doesn't exist
  if [ ! -d "$ssh_dir" ]; then
    sudo -u "$username" mkdir -p "$ssh_dir"
    sudo -u "$username" chmod 700 "$ssh_dir"
  fi
  
  # Add public key to authorized_keys
  log "Adding public key to authorized_keys for $username..."
  
  # Check if key already exists in authorized_keys
  if [ -f "$authorized_keys" ] && grep -Fxq "$public_key" "$authorized_keys"; then
    warn "SSH key already exists in authorized_keys for $username"
  else
    echo "$public_key" | sudo -u "$username" tee -a "$authorized_keys" > /dev/null
    sudo -u "$username" chmod 600 "$authorized_keys"
    log "Public key added to authorized_keys for $username"
  fi
}

show_ssh_setup_summary() {
  if [[ "$SETUP_SSH_KEYS" == "yes" ]]; then
    echo ""
    echo -e "${GREEN}✓ SSH key authentication has been set up for:${NC}"
    echo "  - User: $SYSADMIN_USER"
    if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
      echo "  - User: $ADDITIONAL_USER"
    fi
    echo ""
    echo -e "${YELLOW}Your SSH public key has been added to authorized_keys files.${NC}"
    echo "You can now connect using your SSH private key."
  elif [[ "$DISABLE_SSH_PASSWORD" == "yes" ]]; then
    echo ""
    echo -e "${RED}⚠️  SSH password authentication is DISABLED${NC}"
    echo -e "${YELLOW}You MUST add your SSH public key manually:${NC}"
    echo ""
    echo "1. On your computer, get your public key:"
    echo "   cat ~/.ssh/id_ed25519.pub"
    echo ""
    echo "2. Add it to the server's authorized_keys:"
    echo "   echo 'YOUR_PUBLIC_KEY' >> /home/$SYSADMIN_USER/.ssh/authorized_keys"
    if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
      echo "   echo 'YOUR_PUBLIC_KEY' >> /home/$ADDITIONAL_USER/.ssh/authorized_keys"
    fi
    echo ""
  fi
}


### === CLEANUP AND FINALIZATION === ###
finalize_setup() {
    # Restart SSH service
    log "Restarting SSH service..."
    systemctl restart sshd
    
    # Clean up
    log "Cleaning up..."
    apt autoremove -y
    apt autoclean
    
    # Final security check
    log "Final system status:"
    echo "================================"
    
    if [[ "$USE_UFW" == "yes" ]]; then
        echo "UFW Status:"
        ufw status
    else
        echo "Iptables Status:"
        iptables -L -n | head -20
    fi
    
    echo "================================"
    echo "SSH Configuration:"
    grep -E "PermitRootLogin|PasswordAuthentication|PubkeyAuthentication|AllowUsers" /etc/ssh/sshd_config
    echo "================================"
    
    if [[ "$USE_FAIL2BAN" == "yes" ]]; then
        echo "Fail2ban Status:"
        systemctl status fail2ban --no-pager -l
        echo "================================"
    fi
    
    echo "Sudo users:"
    grep -E "sudo|wheel" /etc/group
    echo "================================"
    
    log "Setup completed successfully!"
    echo ""
    echo -e "${BLUE}=== IMPORTANT NOTES ===${NC}"
    echo -e "${YELLOW}1. Root SSH login is now DISABLED${NC}"
    echo -e "${YELLOW}2. Root password is LOCKED${NC}"
    echo -e "${YELLOW}3. SSH password authentication is DISABLED${NC}"
    echo -e "${YELLOW}4. Use '$SYSADMIN_USER' user with SSH key to connect${NC}"
    
    if [[ "$USE_UFW" == "yes" ]]; then
        echo -e "${YELLOW}5. UFW firewall is active (SSH, HTTP, HTTPS allowed)${NC}"
    else
        echo -e "${YELLOW}5. Basic iptables rules are active (SSH, HTTP, HTTPS allowed)${NC}"
        echo -e "${YELLOW}   You can customize /etc/iptables-basic.sh for additional rules${NC}"
    fi
    
    local note_number=6
    if [[ "$USE_FAIL2BAN" == "yes" ]]; then
        echo -e "${YELLOW}$note_number. Fail2ban is protecting SSH${NC}"
        ((note_number++))
    else
        echo -e "${YELLOW}$note_number. Consider implementing IP banning manually if needed${NC}"
        ((note_number++))
    fi
    
    echo -e "${YELLOW}$note_number. Automatic security updates are enabled${NC}"
    ((note_number++))
    echo -e "${YELLOW}$note_number. SSH keys generated for both $SYSADMIN_USER and root users${NC}"
    
    echo ""
    echo -e "${RED}REBOOT RECOMMENDED${NC}"
    echo ""
    echo "To connect: ssh $SYSADMIN_USER@$(hostname -I | awk '{print $1}')"
}

### === MAIN === ###
main() {
    # Check prerequisites
    check_root
    check_debian
    check_commands
    
    echo -e "${BLUE}=== Debian 12 Initial Setup ===${NC}"
    echo "This script will:"
    echo "1. Set up a secure Debian 12 system"
    echo "2. Create a sysadmin user with SSH key authentication"
    echo "3. Optionally create an additional user account"
    echo "4. Generate SSH keys for all created users and root"
    echo ""
    read -p "Continue? (y/N): " confirm
    
    if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
        echo "Setup cancelled."
        exit 0
    fi
    
    # Ask for user and security preferences
    ask_additional_user
    ask_ssh_security
    ask_firewall_preferences
    
    # System setup
    setup_system
    create_sysadmin_user
    create_additional_user
    configure_security
    
    # SSH key setup
    setup_user_ssh_keys
    show_ssh_setup_summary
    
    # Finalize
    finalize_setup
}

main "$@"