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

# 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
        useradd -m -s /bin/bash "$SYSADMIN_USER"
        log "User $SYSADMIN_USER created successfully"
        
        # Set password for sysadmin user
        echo "Please set a password for user $SYSADMIN_USER:"
        passwd "$SYSADMIN_USER"
    fi
    
    # Add sysadmin to sudo group
    log "Adding $SYSADMIN_USER to sudo group..."
    usermod -aG sudo "$SYSADMIN_USER"
}

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"
        echo "AllowUsers $SYSADMIN_USER"
    } >> /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_sysadmin_ssh_key "$user_email"
  
  # Generate SSH key for root user
  generate_root_ssh_key "$user_email"
}

generate_sysadmin_ssh_key() {
  local user_email="$1"
  local sysadmin_home="/home/$SYSADMIN_USER"
  local ssh_dir="$sysadmin_home/.ssh"
  local ssh_key="$ssh_dir/id_ed25519"
  
  log "Setting up SSH key for user: $SYSADMIN_USER"
  
  # Create .ssh directory for sysadmin user if it doesn't exist
  if [ ! -d "$ssh_dir" ]; then
    sudo -u "$SYSADMIN_USER" mkdir -p "$ssh_dir"
    sudo -u "$SYSADMIN_USER" chmod 700 "$ssh_dir"
  fi
  
  if [ ! -f "$ssh_key.pub" ]; then
    log "Generating SSH key for $SYSADMIN_USER..."
    
    # Generate SSH key as sysadmin user
    sudo -u "$SYSADMIN_USER" ssh-keygen -t ed25519 -C "$user_email" -f "$ssh_key" -N ""
    
    # Set proper permissions
    sudo -u "$SYSADMIN_USER" chmod 600 "$ssh_key"
    sudo -u "$SYSADMIN_USER" chmod 644 "$ssh_key.pub"
    
    log "SSH key generated successfully for $SYSADMIN_USER"
  else
    warn "SSH key already exists for $SYSADMIN_USER"
  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 ""
  
  echo -e "${YELLOW}ROOT USER SSH KEY:${NC}"
  echo "----------------------------------------"
  cat "$root_ssh_key"
  echo "----------------------------------------"
  echo ""
  
  echo -e "${YELLOW}Instructions:${NC}"
  echo "1. Copy both keys above"
  echo "2. Add both keys to your git server and any other systems you need access to"
  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
    
    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. Generate SSH keys for both sysadmin and root users"
    echo ""
    read -p "Continue? (y/N): " confirm
    
    if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
        echo "Setup cancelled."
        exit 0
    fi
    
    # Ask for firewall preferences
    ask_firewall_preferences
    
    # System setup
    setup_system
    create_sysadmin_user
    configure_security
    
    # SSH key setup
    generate_ssh_keys
    show_ssh_keys
    
    # Finalize
    finalize_setup
}

main "$@"