#!/bin/bash
# =============================================================================
# VPS Setup & Hardening Script for Arch Linux
# Prepares a fresh box for OMATERM (docker-only) + Docker deployments
# Run as root on a fresh Arch Linux VPS
# =============================================================================

set -euo pipefail

# =============================================================================
# DEFAULTS
# =============================================================================

TIMEZONE="Europe/London"
LOCALE="en_GB.UTF-8"

# =============================================================================
# COLOURS
# =============================================================================

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

log()  { echo -e "${GREEN}[✓]${NC} $1"; }
warn() { echo -e "${YELLOW}[!]${NC} $1"; }
err()  { echo -e "${RED}[✗]${NC} $1"; }
step() { echo -e "\n${CYAN}━━━ $1 ━━━${NC}\n"; }

# =============================================================================
# PRE-FLIGHT CHECKS
# =============================================================================

if [ "$EUID" -ne 0 ]; then
    err "This script must be run as root."
    err "Usage: sudo bash arch-setup.sh"
    exit 1
fi

if ! grep -qi '^ID=arch' /etc/os-release 2>/dev/null; then
    warn "This script is designed for Arch Linux. Proceed with caution."
fi

# =============================================================================
# INTERACTIVE SETUP
# =============================================================================

echo -e "${CYAN}"
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║           VPS Setup & Hardening Script                      ║"
echo "║           Arch Linux · Docker · OMATERM                     ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo -e "${NC}"

# SSH alias
echo ""
echo "  This name is used in your local SSH config so you can"
echo "  connect with: ssh <alias>  (e.g. ssh titan, ssh nebula)"
read -rp "SSH alias for this server [vps]: " input_alias
SSH_ALIAS="${input_alias:-vps}"

# Username
echo ""
read -rp "Choose a username for the deploy user [deploy]: " input_user
DEPLOY_USER="${input_user:-deploy}"

# Validate username
if ! [[ "$DEPLOY_USER" =~ ^[a-z_][a-z0-9_-]*$ ]]; then
    err "Invalid username. Use lowercase letters, numbers, hyphens, underscores."
    exit 1
fi

# Confirm
echo ""
echo -e "${CYAN}━━━ Confirm settings ━━━${NC}"
echo ""
echo "  SSH alias:      $SSH_ALIAS"
echo "  Deploy user:    $DEPLOY_USER"
echo "  Timezone:       $TIMEZONE"
echo ""
read -p "Continue? (y/n) " -n 1 -r
echo ""
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "Aborted."
    exit 0
fi

# =============================================================================
# 1. SYSTEM UPDATES & ESSENTIALS
# =============================================================================

step "1/8 · System updates & essential packages"

pacman -Syu --noconfirm
# Minimal host package set — omaterm's container brings its own tooling
# (git, tmux, agents, tailscale all live inside it)
pacman -S --noconfirm --needed \
    ufw \
    fail2ban \
    curl \
    docker \
    docker-compose

# Docker must be running before omaterm installs — the new omaterm is
# docker-only (everything runs in the ghcr.io/omacom-io/omaterm container)
systemctl enable --now docker.service

log "System updated, packages installed, Docker running"

# =============================================================================
# 2. TIMEZONE & LOCALE
# =============================================================================

step "2/8 · Timezone & locale"

timedatectl set-timezone "$TIMEZONE"
log "Timezone set to $TIMEZONE"

# Generate and set locale (minimal Arch cloud images often ship C.UTF-8 only)
if ! locale -a 2>/dev/null | grep -qi "^${LOCALE/UTF-8/utf8}$"; then
    sed -i "s/^#${LOCALE}/${LOCALE}/" /etc/locale.gen
    locale-gen
fi
echo "LANG=$LOCALE" > /etc/locale.conf
log "Locale set to $LOCALE"

# =============================================================================
# 3. CREATE DEPLOY USER
# =============================================================================

step "3/8 · Deploy user"

if id "$DEPLOY_USER" &>/dev/null; then
    warn "User '$DEPLOY_USER' already exists — skipping creation"
else
    useradd -m -s /bin/bash "$DEPLOY_USER"
    usermod -aG wheel "$DEPLOY_USER"
    log "User '$DEPLOY_USER' created and added to wheel group"
fi

# Allow sudo without password for deploy user
echo "$DEPLOY_USER ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/$DEPLOY_USER
chmod 440 /etc/sudoers.d/$DEPLOY_USER
log "Passwordless sudo enabled for $DEPLOY_USER"

# Copy root's authorized_keys to deploy user if they exist
if [ -f /root/.ssh/authorized_keys ]; then
    mkdir -p /home/$DEPLOY_USER/.ssh
    cp /root/.ssh/authorized_keys /home/$DEPLOY_USER/.ssh/authorized_keys
    chown -R $DEPLOY_USER:$DEPLOY_USER /home/$DEPLOY_USER/.ssh
    chmod 700 /home/$DEPLOY_USER/.ssh
    chmod 600 /home/$DEPLOY_USER/.ssh/authorized_keys
    log "SSH keys copied from root to $DEPLOY_USER"
else
    warn "No SSH keys found for root — you'll need to add keys manually:"
    warn "  ssh-copy-id $DEPLOY_USER@YOUR_SERVER_IP"
fi

# =============================================================================
# 4. HARDEN SSH
# =============================================================================

step "4/8 · SSH hardening"

# Backup original config
cp /etc/ssh/sshd_config "/etc/ssh/sshd_config.backup.$(date +%Y%m%d)"

# Neutralise conflicting directives in the main config so our drop-in wins
sed -i 's/^\s*PermitRootLogin\b/#&/'          /etc/ssh/sshd_config
sed -i 's/^\s*PasswordAuthentication\b/#&/'    /etc/ssh/sshd_config
sed -i 's/^\s*PubkeyAuthentication\b/#&/'      /etc/ssh/sshd_config
sed -i 's/^\s*X11Forwarding\b/#&/'             /etc/ssh/sshd_config
sed -i 's/^\s*PermitEmptyPasswords\b/#&/'      /etc/ssh/sshd_config
log "Commented out conflicting directives in sshd_config"

# Remove cloud-init SSH overrides that can conflict (some VPS images use them)
rm -f /etc/ssh/sshd_config.d/50-cloud-init.conf

# Ensure the drop-in directory is included (Arch's stock sshd_config has this,
# but some provider images strip it)
mkdir -p /etc/ssh/sshd_config.d
if ! grep -q '^Include /etc/ssh/sshd_config.d/\*\.conf' /etc/ssh/sshd_config; then
    sed -i '1i Include /etc/ssh/sshd_config.d/*.conf' /etc/ssh/sshd_config
    warn "Added missing Include for sshd_config.d drop-ins"
fi

cat > /etc/ssh/sshd_config.d/hardened.conf << EOF
# Custom SSH hardening
PermitRootLogin no
PasswordAuthentication no
PubkeyAuthentication yes
MaxAuthTries 3
MaxSessions 5
AllowUsers $DEPLOY_USER
ClientAliveInterval 300
ClientAliveCountMax 2
X11Forwarding no
PermitEmptyPasswords no
EOF

log "SSH hardened — root disabled, key-only auth"

# Test config before restarting (service is 'sshd' on Arch, not 'ssh')
if sshd -t; then
    systemctl restart sshd
    log "SSH restarted successfully"
else
    err "SSH config test failed — reverting"
    rm /etc/ssh/sshd_config.d/hardened.conf
    exit 1
fi

# =============================================================================
# 5. FIREWALL (UFW)
# =============================================================================

step "5/8 · Firewall (UFW)"

ufw default deny incoming
ufw default allow outgoing
ufw allow 22/tcp comment 'SSH'

# Enable without prompt, and make it survive reboots
echo "y" | ufw enable
systemctl enable ufw

log "UFW enabled — port 22 only"
ufw status verbose

# =============================================================================
# 6. FAIL2BAN
# =============================================================================

step "6/8 · Fail2ban"

# Arch logs auth to the systemd journal — there is no /var/log/auth.log,
# so use the systemd backend instead of a logpath
cat > /etc/fail2ban/jail.local << EOF
[DEFAULT]
bantime = 3600
findtime = 600
maxretry = 3
banaction = ufw
backend = systemd

[sshd]
enabled = true
port = 22
filter = sshd
maxretry = 3
bantime = 3600
EOF

systemctl enable fail2ban
systemctl restart fail2ban
log "Fail2ban configured and running (systemd journal backend)"

# =============================================================================
# 7. AUTOMATIC UPDATES (weekly pacman -Syu timer)
# =============================================================================

step "7/8 · Automatic updates"

# Arch has no unattended-upgrades equivalent — use a weekly full update timer.
# NOTE: rolling release means an update can very occasionally need manual
# intervention (check https://archlinux.org/news/ if something breaks).
cat > /etc/systemd/system/pacman-update.service << EOF
[Unit]
Description=Weekly pacman system update
After=network-online.target
Wants=network-online.target

[Service]
Type=oneshot
ExecStart=/usr/bin/pacman -Syu --noconfirm
EOF

cat > /etc/systemd/system/pacman-update.timer << EOF
[Unit]
Description=Weekly pacman system update

[Timer]
OnCalendar=weekly
RandomizedDelaySec=1h
Persistent=true

[Install]
WantedBy=timers.target
EOF

systemctl daemon-reload
systemctl enable --now pacman-update.timer
log "Weekly auto-update timer enabled (pacman-update.timer)"
warn "Rolling release: check 'journalctl -u pacman-update' if anything misbehaves after an update"

# =============================================================================
# 8. OMATERM PREP (Docker already installed via pacman in step 1)
# =============================================================================

step "8/8 · Omaterm prep"

# Omaterm itself is installed interactively after first login as the deploy
# user — here we just make sure the box is ready for it:

# Deploy user needs docker group membership to run the omaterm container
if getent group docker &>/dev/null; then
    usermod -aG docker $DEPLOY_USER
    log "$DEPLOY_USER added to docker group"
fi

# Pre-pull the omaterm image so the installer doesn't have to
log "Pre-pulling omaterm image (this may take a while)..."
if docker pull ghcr.io/omacom-io/omaterm:latest; then
    log "Omaterm image pre-pulled"
else
    warn "Image pre-pull failed — the omaterm installer will pull it instead"
fi

# =============================================================================
# SWAP (2GB)
# =============================================================================

step "Swap file"

if [ ! -f /swapfile ]; then
    fallocate -l 2G /swapfile
    chmod 600 /swapfile
    mkswap /swapfile
    swapon /swapfile
    grep -q '/swapfile' /etc/fstab || echo '/swapfile none swap sw 0 0' >> /etc/fstab
    log "2GB swap created and enabled"
else
    warn "Swap file already exists — skipping"
fi

# =============================================================================
# DISABLE UNNECESSARY SERVICES
# =============================================================================

step "Disable unnecessary services"

for svc in ModemManager udisks2 multipathd; do
    if systemctl is-enabled "$svc" &>/dev/null; then
        systemctl stop "$svc"
        systemctl disable "$svc"
        log "Disabled $svc"
    fi
done

# =============================================================================
# CAP JOURNAL LOGS (50MB)
# =============================================================================

step "Cap journal logs"

mkdir -p /etc/systemd/journald.conf.d
cat > /etc/systemd/journald.conf.d/size.conf << EOF
[Journal]
SystemMaxUse=50M
EOF

systemctl restart systemd-journald
log "Journal capped at 50MB"

# =============================================================================
# CLEAN PACMAN CACHE (keep last 2 versions of each package)
# =============================================================================

step "Pacman cache hygiene"

pacman -S --noconfirm --needed pacman-contrib
systemctl enable --now paccache.timer
log "paccache.timer enabled — old package versions pruned weekly"

# =============================================================================
# GENERATE LOCAL SSH CONFIG FILE
# =============================================================================

step "SSH config for your local machine"

SERVER_IP=$(curl -4 -s --max-time 5 ifconfig.me 2>/dev/null || \
            curl -4 -s --max-time 5 icanhazip.com 2>/dev/null || \
            echo "YOUR_SERVER_IP")

cat > /home/$DEPLOY_USER/ssh-config-snippet.txt << SSHEOF
Host $SSH_ALIAS
    HostName $SERVER_IP
    User $DEPLOY_USER
    IdentityFile ~/.ssh/id_ed25519
SSHEOF

chown $DEPLOY_USER:$DEPLOY_USER /home/$DEPLOY_USER/ssh-config-snippet.txt
log "SSH config snippet saved to ~/ssh-config-snippet.txt"

# =============================================================================
# SUMMARY
# =============================================================================

echo ""
echo -e "${CYAN}"
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║                    SETUP COMPLETE                           ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo -e "${NC}"
echo ""
echo -e "  ${GREEN}SSH${NC}"
echo "  ├─ User:          $DEPLOY_USER"
echo "  ├─ Root login:    disabled"
echo "  └─ Auth:          key-only"
echo ""
echo -e "  ${GREEN}Firewall${NC}"
echo "  └─ 22/tcp          SSH (only open port)"
echo ""
echo -e "  ${GREEN}Services${NC}"
echo "  ├─ Docker:        $(docker --version 2>/dev/null || echo 'installed')"
echo "  ├─ Omaterm:       image pre-pulled — install after first login (see below)"
echo "  ├─ Fail2ban:      active (systemd backend)"
echo "  └─ Auto-updates:  weekly pacman-update.timer"
echo ""
echo -e "  ${GREEN}Paths${NC}"
echo "  └─ SSH config:    ~/ssh-config-snippet.txt"
echo ""
echo -e "  ${YELLOW}NEXT STEPS${NC}"
echo ""
echo "  1. On your local machine, run this to add the SSH config:"
echo ""
echo -e "  ${CYAN}────── COPY BELOW THIS LINE ──────${NC}"
echo ""
echo "cat >> ~/.ssh/config << 'EOF'"
echo ""
echo "Host $SSH_ALIAS"
echo "    HostName $SERVER_IP"
echo "    User $DEPLOY_USER"
echo "    IdentityFile ~/.ssh/id_ed25519"
echo "EOF"
echo ""
echo -e "  ${CYAN}────── COPY ABOVE THIS LINE ──────${NC}"
echo ""
echo "  2. Test SSH:  ssh $SSH_ALIAS"
echo "  3. Install omaterm (interactive SSH session, NOT a one-liner —"
echo "     the installer needs a real terminal with TERM set):"
echo ""
echo -e "       ${CYAN}ssh $SSH_ALIAS${NC}"
echo -e "       ${CYAN}curl -fsSL https://omaterm.org/install | bash${NC}"
echo ""
echo "     If it complains about the terminal: TERM=xterm-256color first."
echo ""
echo "  Once step 1 is done, Claude Code can control this server."
echo ""
echo -e "  ${YELLOW}A reboot is recommended to apply kernel updates:${NC}"
echo -e "  ${CYAN}  reboot${NC}"
echo ""