#!/bin/bash
set -euo pipefail

###################
###    USAGE    ###
###################

# This script performs complete initial server setup tasks on a fresh Ubuntu installation,
# including creating a new sudo user, configuring SSH access, and setting up system basics.
# It is intended to be run as root on a fresh server before running the hardening script.
#
# Usage: ./initial-setup

# https://www.digitalocean.com/community/tutorials/initial-server-setup-with-ubuntu

if [ "$EUID" -ne 0 ]; then
    echo "This script must be run as root" >&2
    exit 1
fi

#############################
### APT LOCK CONFIGURATION ###
#############################

# Set APT lock timeout to prevent failures when other processes are using APT
cat > /etc/apt/apt.conf.d/90lock-timeout << 'EOF'
DPkg::Lock::Timeout "300";
EOF

########################
### SCRIPT VARIABLES ###
########################

DEFAULT_USERNAME="admin"

read -rp "Enter the username for the new sudo user [${DEFAULT_USERNAME}]: " USERNAME_INPUT
if [ -n "$USERNAME_INPUT" ]; then
    USERNAME=$USERNAME_INPUT
else
    USERNAME=${DEFAULT_USERNAME}
fi

if id -u "${USERNAME}" >/dev/null 2>&1; then
    echo "User ${USERNAME} already exists. Please choose a different username." >&2
    exit 1
elif [ -z "${USERNAME}" ]; then
    echo "No username provided. Please provide a valid username." >&2
    exit 1
fi

# Hostname configuration
CURRENT_HOSTNAME=$(hostname)
echo ""
read -rp "Enter a hostname for this server [${CURRENT_HOSTNAME}]: " HOSTNAME_INPUT
if [ -n "$HOSTNAME_INPUT" ]; then
    NEW_HOSTNAME=$HOSTNAME_INPUT
else
    NEW_HOSTNAME=${CURRENT_HOSTNAME}
fi

# Timezone configuration
CURRENT_TIMEZONE=$(timedatectl show -p Timezone --value 2>/dev/null || echo "UTC")
echo ""
echo "Current timezone: ${CURRENT_TIMEZONE}"
read -rp "Enter timezone (e.g., America/New_York, Europe/London) [${CURRENT_TIMEZONE}]: " TIMEZONE_INPUT
if [ -n "$TIMEZONE_INPUT" ]; then
    NEW_TIMEZONE=$TIMEZONE_INPUT
else
    NEW_TIMEZONE=${CURRENT_TIMEZONE}
fi

# Git configuration (optional)
echo ""
read -rp "Would you like to configure Git for this user? (y/N): " CONFIGURE_GIT
if [[ "$CONFIGURE_GIT" =~ ^[Yy]$ ]]; then
    read -rp "Enter Git user name: " GIT_USER_NAME
    read -rp "Enter Git user email: " GIT_USER_EMAIL
else
    GIT_USER_NAME=""
    GIT_USER_EMAIL=""
fi

COPY_AUTHORIZED_KEYS_FROM_ROOT=true

OTHER_PUBLIC_KEYS_TO_ADD=(
)

###############################
### SYSTEM UPDATE & PACKAGES ###
###############################

echo "Updating system packages..."
apt-get -y update
apt-get -y upgrade
PYTHONWARNINGS="ignore::SyntaxWarning" apt-get -y install fail2ban unattended-upgrades apt-listchanges

echo "Configuring fail2ban..."
cat > /etc/fail2ban/jail.local << 'EOF'
[DEFAULT]
bantime = 3600
findtime = 600
maxretry = 3

[sshd]
enabled = true
port = ssh
logpath = %(sshd_log)s
backend = %(sshd_backend)s
EOF

systemctl enable fail2ban
systemctl start fail2ban

echo "Configuring unattended-upgrades..."
cat > /etc/apt/apt.conf.d/50-unattended-upgrades << 'EOF'
Unattended-Upgrade::Allowed-Origins {
    "${distro_id}:${distro_codename}";
    "${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-Kernel-Packages "true";
Unattended-Upgrade::Remove-Unused-Dependencies "true";
Unattended-Upgrade::Automatic-Reboot "false";
EOF

cat > /etc/apt/apt.conf.d/20-auto-upgrades << 'EOF'
APT::Periodic::Update-Package-Lists "1";
APT::Periodic::Download-Upgradeable-Packages "1";
APT::Periodic::AutocleanInterval "7";
APT::Periodic::Unattended-Upgrade "1";
EOF

systemctl enable unattended-upgrades
systemctl start unattended-upgrades

###########################
### SYSTEM CONFIGURATION ###
###########################

# Set hostname
if [ "${NEW_HOSTNAME}" != "${CURRENT_HOSTNAME}" ]; then
    echo "Setting hostname to ${NEW_HOSTNAME}..."
    hostnamectl set-hostname "${NEW_HOSTNAME}"
    
    # Update /etc/hosts - remove old hostname entry and add new one
    sed -i "/^127\.0\.1\.1[[:space:]]/d" /etc/hosts
    echo "127.0.1.1 ${NEW_HOSTNAME}" >> /etc/hosts
fi

# Set timezone
if [ "${NEW_TIMEZONE}" != "${CURRENT_TIMEZONE}" ]; then
    echo "Setting timezone to ${NEW_TIMEZONE}..."
    if timedatectl set-timezone "${NEW_TIMEZONE}" 2>/dev/null; then
        echo "Timezone set successfully"
    else
        echo "Warning: Failed to set timezone to ${NEW_TIMEZONE}. Keeping ${CURRENT_TIMEZONE}"
    fi
fi

# Configure swap file
TOTAL_MEM=$(free -m | awk '/^Mem:/{print $2}')
SWAP_EXISTS=$(swapon --show=NAME --noheadings | wc -l)

if [ "${SWAP_EXISTS}" -eq 0 ]; then
    echo "No swap detected. Creating swap file..."
    
    # Determine swap size (equal to double the RAM up to 2GB, then 2GB for larger instances)
    if [ "${TOTAL_MEM}" -lt 2048 ]; then
        SWAP_SIZE="$((TOTAL_MEM * 2))M"
    else
        SWAP_SIZE="2G"
    fi
    
    fallocate -l "${SWAP_SIZE}" /swapfile
    chmod 600 /swapfile
    mkswap /swapfile
    swapon /swapfile
    
    # Make swap permanent
    echo '/swapfile none swap sw 0 0' >> /etc/fstab
    
    # Adjust swappiness for server use
    if ! grep -q "^vm.swappiness" /etc/sysctl.conf; then
        echo 'vm.swappiness=10' >> /etc/sysctl.conf
    fi
    sysctl -p > /dev/null 2>&1
    
    echo "Swap file created: ${SWAP_SIZE}"
else
    echo "Swap already configured, skipping..."
fi

###########################
### CREATE SUDO USER ###
###########################

echo "Creating sudo user: ${USERNAME}..."
useradd --create-home --shell "/bin/bash" --groups sudo "${USERNAME}"

# Lock root account and delete new user password
passwd --lock root
passwd --delete "${USERNAME}"
chage --lastday 0 "${USERNAME}"

# Configure Git for user if requested
if [ -n "${GIT_USER_NAME}" ] && [ -n "${GIT_USER_EMAIL}" ]; then
    echo "Configuring Git for ${USERNAME}..."
    sudo -u "${USERNAME}" bash << EOF
git config --global user.name "${GIT_USER_NAME}"
git config --global user.email "${GIT_USER_EMAIL}"
EOF
fi

###########################
### SSH CONFIGURATION ###
###########################

echo "Configuring SSH for ${USERNAME}..."
home_directory="$(eval echo ~${USERNAME})"
mkdir --parents "${home_directory}/.ssh"

if [ "${COPY_AUTHORIZED_KEYS_FROM_ROOT}" = true ] && [ -f /root/.ssh/authorized_keys ]; then
    cp /root/.ssh/authorized_keys "${home_directory}/.ssh"
fi

if [ ${#OTHER_PUBLIC_KEYS_TO_ADD[@]} -gt 0 ]; then
    for pub_key in "${OTHER_PUBLIC_KEYS_TO_ADD[@]}"; do
        echo "${pub_key}" >> "${home_directory}/.ssh/authorized_keys"
    done
fi

chmod 0700 "${home_directory}/.ssh"
if [ -f "${home_directory}/.ssh/authorized_keys" ]; then
    chmod 0600 "${home_directory}/.ssh/authorized_keys"
fi
chown --recursive "${USERNAME}":"${USERNAME}" "${home_directory}/.ssh"

# Verify SSH key setup to prevent lockout
if [ ! -f "${home_directory}/.ssh/authorized_keys" ] || [ ! -s "${home_directory}/.ssh/authorized_keys" ]; then
    echo ""
    echo "WARNING: No SSH keys have been configured for ${USERNAME}!"
    echo "You will need SSH keys to login after the hardening script runs."
    echo ""
    echo "Please add at least one SSH public key to continue."
    echo "You can:"
    echo "  1. Add keys to the OTHER_PUBLIC_KEYS_TO_ADD array in this script"
    echo "  2. Ensure root has authorized_keys that can be copied"
    echo ""
    read -rp "Press Enter to exit and fix this issue..."
    exit 1
fi

###########################
### FIREWALL SETUP ###
###########################

echo "Configuring basic UFW firewall..."
# Allow SSH on default port 22 for initial setup
# The hardening script will reconfigure this with proper port and restrictions
ufw allow OpenSSH comment "SSH (Initial Setup)"
ufw --force enable

###########################
### SSH KEY GENERATION ###
###########################

echo ""
echo "Generating SSH key for ${USERNAME}..."

sudo -u "${USERNAME}" bash << 'EOFSSH'
set -euo pipefail

if [ ! -f ~/.ssh/id_ed25519 ]; then
    echo "Generating ed25519 SSH key..."
    ssh-keygen -t ed25519 -f ~/.ssh/id_ed25519 -N "" -C "$(hostname)"
fi

echo ""
echo "SSH key generated successfully!"
echo "Public key:"
echo "=============================================="
cat ~/.ssh/id_ed25519.pub
echo "=============================================="
EOFSSH

###########################
### GITHUB SETUP (OPTIONAL) ###
###########################

echo ""
read -rp "Would you like to set up GitHub SSH authentication? (Y/n): " SETUP_GITHUB
if [[ ! "$SETUP_GITHUB" =~ ^[Nn]$ ]]; then
    echo "Installing GitHub CLI..."
    curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg | dd of=/usr/share/keyrings/githubcli-archive-keyring.gpg 2>/dev/null
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | tee /etc/apt/sources.list.d/github-cli.list > /dev/null
    apt-get update
    apt-get install -y gh

    echo ""
    echo "Switching to user ${USERNAME}..."
    echo "You will now authenticate with GitHub and add your SSH key."
    echo ""

    sudo -u "${USERNAME}" bash << 'EOFGITHUB'
set -euo pipefail

echo "Authenticating with GitHub..."
echo "Follow the prompts to complete authentication in your browser."
gh auth login -h github.com -s admin:public_key

echo ""
echo "Adding SSH key to GitHub..."
gh ssh-key add ~/.ssh/id_ed25519.pub --title "$(hostname)-$(date +%Y%m%d)"

echo ""
echo "Testing SSH connection to GitHub..."
ssh -T git@github.com || true

echo ""
echo "Logging out of GitHub CLI..."
gh auth logout -h github.com

echo ""
echo "GitHub SSH setup complete!"
EOFGITHUB

fi

###########################
### SETUP COMPLETE ###
###########################

echo ""
echo "=========================================="
echo "Initial server setup complete!"
echo "=========================================="
echo "User created: ${USERNAME}"
echo "Hostname: ${NEW_HOSTNAME}"
echo "Timezone: ${NEW_TIMEZONE}"
echo "SSH key: ed25519 generated for ${USERNAME}"
echo "Unattended upgrades: enabled"
echo "Fail2ban: configured and running"
echo "UFW firewall: enabled (SSH port 22 allowed)"
if [ "${SWAP_EXISTS}" -eq 0 ]; then
    echo "Swap: ${SWAP_SIZE} created"
else
    echo "Swap: already configured"
fi
echo ""
echo "IMPORTANT: You will be required to change your password on first login."
echo "Log in with: ssh ${USERNAME}@$(hostname -I | awk '{print $1}')"
echo ""
echo "=========================================="
echo "NEXT STEPS:"
echo "=========================================="
echo "1. Test SSH access with the new user before disconnecting"
echo "2. Run the hardening script: sudo hardening/harden"
echo "   This will configure:"
echo "   - System security hardening (sysctl, kernel params)"
echo "   - SSH hardening"
echo "   - UFW firewall with proper rules and restrictions"
echo "   - Fail2ban, AppArmor, auditd"
echo "   - File integrity monitoring (AIDE)"
echo "=========================================="

###########################
### DISCONNECT SESSION ###
###########################

echo ""
echo "Disconnecting SSH session now..."
sleep 1

# Kill the parent SSH session gracefully
kill -HUP $PPID