#!/bin/bash
# =============================================================================
# VPS Setup & Hardening Script for Arch Linux
# Takes a fresh Arch box to a clean, locked-down base — nothing else.
#
# Run as root on a fresh Arch Linux VPS:
#
#   ssh -t root@YOUR_SERVER 'bash <(curl -fsSL https://raw.githubusercontent.com/28allday/arch-boot-strap/main/arch-setup.sh)'
#
# or scp it over and:  ssh -t root@YOUR_SERVER bash arch-setup.sh
#
# Safe to re-run — completed steps skip themselves, and your answers survive
# a mid-run reboot. After the first run root SSH login is disabled, so
# re-run as the deploy user with sudo.
#
# Want hardening + OMATERM in one shot? That's the sibling project:
# https://github.com/28allday/omaserver
# =============================================================================

set -euo pipefail

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

TIMEZONE="${TIMEZONE:-Europe/London}"
LOCALE="${LOCALE:-en_GB.UTF-8}"
SELF_URL="https://raw.githubusercontent.com/28allday/arch-boot-strap/main/arch-setup.sh"

# =============================================================================
# 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: ssh -t root@YOUR_SERVER 'bash <(curl -fsSL $SELF_URL)'"
    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 · clean locked-down base                ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo -e "${NC}"

# Public IP — used in re-run instructions and the SSH config snippet
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")

# Answers survive the kernel-upgrade reboot — a re-run finds them here and
# doesn't ask again
STATE_FILE="/root/.arch-setup.conf"
SSH_ALIAS=""
DEPLOY_USER=""
if [ -f "$STATE_FILE" ]; then
    # shellcheck source=/dev/null
    . "$STATE_FILE"
fi

if [ -n "$SSH_ALIAS" ] && [ -n "$DEPLOY_USER" ]; then
    log "Using saved settings from the previous run"
else
    # Machine name — used as the hostname AND the alias in your local SSH
    # config (ssh <name>). Asked ONCE, used everywhere.
    echo ""
    echo "  This name becomes the hostname and your local SSH alias"
    echo "  (ssh <name>)."
    read -rp "Machine name 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}"
fi

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

# Confirm
echo ""
echo -e "${CYAN}━━━ Confirm settings ━━━${NC}"
echo ""
echo "  Machine name:   $SSH_ALIAS  (hostname + ssh alias)"
echo "  Deploy user:    $DEPLOY_USER"
echo "  Timezone:       $TIMEZONE"
echo ""
read -p "Continue? (Y/n) " -n 1 -r
echo ""
if [[ $REPLY =~ ^[Nn]$ ]]; then
    rm -f "$STATE_FILE"
    echo "Aborted — saved settings cleared. Re-run to start over."
    exit 0
fi

# Persist for re-runs (cleared on abort above)
printf 'SSH_ALIAS=%q\nDEPLOY_USER=%q\n' "$SSH_ALIAS" "$DEPLOY_USER" > "$STATE_FILE"
chmod 600 "$STATE_FILE"

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

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

pacman -Syu --noconfirm

# If -Syu upgraded the kernel, the running kernel's modules are gone and
# anything that loads modules (netfilter for ufw, tun, etc.) will fail —
# reboot and re-run (script is safe to re-run; completed steps skip
# themselves). This runs BEFORE SSH hardening, so root login still works
# for the re-run.
if [ ! -d "/usr/lib/modules/$(uname -r)" ]; then
    warn "Kernel was upgraded — running kernel $(uname -r) has no modules on disk."
    warn "The box must reboot before setup can continue."
    echo ""
    echo "  After the reboot, reconnect and re-run — it picks up where it left off:"
    echo ""
    if [ -f "$0" ] && [[ "$0" != /dev/fd/* ]]; then
        echo -e "    ${CYAN}ssh -t root@$SERVER_IP bash $(basename "$0")${NC}"
    else
        echo -e "    ${CYAN}ssh -t root@$SERVER_IP 'bash <(curl -fsSL $SELF_URL)'${NC}"
    fi
    echo ""
    read -p "Reboot now? (Y/n) " -n 1 -r
    echo ""
    if [[ ! $REPLY =~ ^[Nn]$ ]]; then
        systemctl reboot
    fi
    exit 1
fi

# Hardening essentials only — anything app-shaped (docker, etc.) is the
# consuming project's job.
# kitty-terminfo: without it, SSHing in from kitty gives
# "'xterm-kitty': unknown terminal type" and interactive tools die
pacman -S --noconfirm --needed \
    ufw \
    fail2ban \
    curl \
    kitty-terminfo \
    pacman-contrib

log "System updated, hardening packages installed"

# =============================================================================
# 2. HOSTNAME, TIMEZONE & LOCALE
# =============================================================================

step "2/8 · Hostname, timezone & locale"

hostnamectl set-hostname "$SSH_ALIAS"
log "Hostname set to $SSH_ALIAS"

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@$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'

# Dormant unless tailscale is installed later — without it, a future
# `tailscale up` silently drops all incoming tailnet traffic
ufw allow in on tailscale0 comment 'Tailscale'

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

log "UFW enabled — 22/tcp + tailscale0 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. HOUSEKEEPING — swap, services, journal, pacman cache
# =============================================================================

step "8/8 · Housekeeping"

# --- Swap (2GB) ---
if swapon --show --noheadings 2>/dev/null | grep -q .; then
    warn "Swap already active — skipping:"
    swapon --show
elif [ -f /swapfile ]; then
    warn "Swap file already exists — skipping"
else
    # fallocate swapfiles fail on btrfs (needs NOCOW) — use the proper tool
    if [ "$(findmnt -n -o FSTYPE /)" = "btrfs" ]; then
        btrfs filesystem mkswapfile --size 2g /swapfile
    else
        dd if=/dev/zero of=/swapfile bs=1M count=2048 status=none
        chmod 600 /swapfile
        mkswap /swapfile
    fi
    swapon /swapfile
    grep -q '/swapfile' /etc/fstab || echo '/swapfile none swap sw 0 0' >> /etc/fstab
    log "2GB swap created and enabled"
fi

# --- 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) ---
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"

# --- Pacman cache hygiene (prune old package versions weekly) ---
systemctl enable --now paccache.timer
log "paccache.timer enabled — old package versions pruned weekly"

# =============================================================================
# LOCAL SSH CONFIG SNIPPET
# =============================================================================

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

# Hardening complete — saved answers no longer needed
rm -f "$STATE_FILE"

# =============================================================================
# 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"
echo "  └─ tailscale0     allowed (dormant unless tailscale is installed)"
echo ""
echo -e "  ${GREEN}Services${NC}"
echo "  ├─ Fail2ban:      active (systemd backend)"
echo "  └─ Auto-updates:  weekly pacman-update.timer"
echo ""
echo "  On your local machine, 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 "  Then test:  ssh $SSH_ALIAS"
echo ""
echo -e "  ${YELLOW}A reboot is recommended to apply kernel updates:${NC}"
echo -e "  ${CYAN}  reboot${NC}"
echo ""