#!/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 safely
    log "Applying SSH security settings..."
    
    # Backup original sshd_config
    cp /etc/ssh/sshd_config /etc/ssh/sshd_config.backup.$(date +%Y%m%d_%H%M%S)
    
    # Configure PasswordAuthentication
    if [[ "$DISABLE_SSH_PASSWORD" == "yes" ]]; then
        sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config
        log "SSH password authentication disabled"
    else
        sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication yes/' /etc/ssh/sshd_config
        log "SSH password authentication enabled"
    fi
    
    # Configure other SSH settings more safely
    sed -i 's/^#*PubkeyAuthentication.*/PubkeyAuthentication yes/' /etc/ssh/sshd_config
    sed -i 's/^#*PermitEmptyPasswords.*/PermitEmptyPasswords no/' /etc/ssh/sshd_config
    sed -i 's/^#*X11Forwarding.*/X11Forwarding no/' /etc/ssh/sshd_config
    sed -i 's/^#*MaxAuthTries.*/MaxAuthTries 3/' /etc/ssh/sshd_config
    
    # Add settings that might not exist
    if ! grep -q "^ClientAliveInterval" /etc/ssh/sshd_config; then
        echo "ClientAliveInterval 300" >> /etc/ssh/sshd_config
    else
        sed -i 's/^#*ClientAliveInterval.*/ClientAliveInterval 300/' /etc/ssh/sshd_config
    fi
    
    if ! grep -q "^ClientAliveCountMax" /etc/ssh/sshd_config; then
        echo "ClientAliveCountMax 2" >> /etc/ssh/sshd_config
    else
        sed -i 's/^#*ClientAliveCountMax.*/ClientAliveCountMax 2/' /etc/ssh/sshd_config
    fi
    
    # Configure AllowUsers (remove old entries first)
    sed -i '/^AllowUsers/d' /etc/ssh/sshd_config
    if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
        echo "AllowUsers $SYSADMIN_USER $ADDITIONAL_USER" >> /etc/ssh/sshd_config
    else
        echo "AllowUsers $SYSADMIN_USER" >> /etc/ssh/sshd_config
    fi
    
    # Test SSH configuration before proceeding
    log "Testing SSH configuration..."
    if ! sshd -t; then
        error "SSH configuration is invalid! Restoring backup..."
        mv /etc/ssh/sshd_config.backup.$(date +%Y%m%d_%H%M%S) /etc/ssh/sshd_config
        exit 1
    fi
    
    log "SSH configuration is valid"
    
    # 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 safely
    log "Restarting SSH service..."
    if ! systemctl restart sshd; then
        error "Failed to restart SSH service! Check configuration and try again."
        log "You can restore SSH configuration from backup if needed:"
        log "ls /etc/ssh/sshd_config.backup.*"
        exit 1
    fi
    
    # Verify SSH service is running
    if ! systemctl is-active --quiet sshd; then
        error "SSH service is not running after restart!"
        exit 1
    fi
    
    log "SSH service restarted successfully"
    
    # 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 "$@"