#!/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=""

# 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 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
    
    # Additional SSH hardening
    log "Applying SSH security settings..."
    {
        echo "Protocol 2"
        echo "PasswordAuthentication no"
        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() {
    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"
}


### === SSH KEY GENERATION === ###
generate_ssh_keys() {
  echo "[+] Please enter your email for the SSH keys:"
  read -p "Enter your email: " user_email
  
  # Generate SSH key for sysadmin user
  generate_user_ssh_key "$SYSADMIN_USER" "$user_email"
  
  # Generate SSH key for additional user if created
  if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
    generate_user_ssh_key "$ADDITIONAL_USER" "$user_email"
  fi
  
  # Generate SSH key for root user
  generate_root_ssh_key "$user_email"
}

generate_user_ssh_key() {
  local username="$1"
  local user_email="$2"
  local user_home="/home/$username"
  local ssh_dir="$user_home/.ssh"
  local ssh_key="$ssh_dir/id_ed25519"
  
  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
  
  if [ ! -f "$ssh_key.pub" ]; then
    log "Generating SSH key for $username..."
    
    # Generate SSH key as user
    sudo -u "$username" ssh-keygen -t ed25519 -C "$user_email" -f "$ssh_key" -N ""
    
    # Set proper permissions
    sudo -u "$username" chmod 600 "$ssh_key"
    sudo -u "$username" chmod 644 "$ssh_key.pub"
    
    log "SSH key generated successfully for $username"
  else
    warn "SSH key already exists for $username"
  fi
}

generate_root_ssh_key() {
  local user_email="$1"
  local root_ssh_dir="/root/.ssh"
  local root_ssh_key="$root_ssh_dir/id_ed25519"
  
  log "Setting up SSH key for root user"
  
  # Create .ssh directory for root if it doesn't exist
  if [ ! -d "$root_ssh_dir" ]; then
    mkdir -p "$root_ssh_dir"
    chmod 700 "$root_ssh_dir"
  fi
  
  if [ ! -f "$root_ssh_key.pub" ]; then
    log "Generating SSH key for root..."
    
    # Generate SSH key as root
    ssh-keygen -t ed25519 -C "$user_email" -f "$root_ssh_key" -N ""
    
    # Set proper permissions
    chmod 600 "$root_ssh_key"
    chmod 644 "$root_ssh_key.pub"
    
    log "SSH key generated successfully for root"
  else
    warn "SSH key already exists for root"
  fi
}

show_ssh_keys() {
  local sysadmin_home="/home/$SYSADMIN_USER"
  local sysadmin_ssh_key="$sysadmin_home/.ssh/id_ed25519.pub"
  local root_ssh_key="/root/.ssh/id_ed25519.pub"
  
  echo ""
  echo -e "${BLUE}=== SSH PUBLIC KEYS ===${NC}"
  echo "Add the following public keys to your Git server and any other systems:"
  echo ""
  
  echo -e "${YELLOW}SYSADMIN USER ($SYSADMIN_USER) SSH KEY:${NC}"
  echo "----------------------------------------"
  cat "$sysadmin_ssh_key"
  echo "----------------------------------------"
  echo ""
  
  # Show additional user SSH key if created
  if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
    local additional_home="/home/$ADDITIONAL_USER"
    local additional_ssh_key="$additional_home/.ssh/id_ed25519.pub"
    
    echo -e "${YELLOW}ADDITIONAL USER ($ADDITIONAL_USER) SSH KEY:${NC}"
    echo "----------------------------------------"
    cat "$additional_ssh_key"
    echo "----------------------------------------"
    echo ""
  fi
  
  echo -e "${YELLOW}ROOT USER SSH KEY:${NC}"
  echo "----------------------------------------"
  cat "$root_ssh_key"
  echo "----------------------------------------"
  echo ""
  
  echo -e "${YELLOW}Instructions:${NC}"
  if [[ "$CREATE_ADDITIONAL_USER" == "yes" && -n "$ADDITIONAL_USER" ]]; then
    echo "1. Copy all keys above"
    echo "2. Add all keys to your git server and any other systems you need access to"
  else
    echo "1. Copy both keys above"
    echo "2. Add both keys to your git server and any other systems you need access to"
  fi
  echo "3. You can now use these keys for SSH authentication"
  echo ""
  read -p "Press Enter to continue..."
}


### === 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 firewall preferences
    ask_additional_user
    ask_firewall_preferences
    
    # System setup
    setup_system
    create_sysadmin_user
    create_additional_user
    configure_security
    
    # SSH key setup
    generate_ssh_keys
    show_ssh_keys
    
    # Finalize
    finalize_setup
}

main "$@"