#!/bin/bash

# LXS - Server Hardening
# Description: Apply baseline security hardening (UFW + fail2ban + SSH key-only + unattended-upgrades)
# Author: LXS
# Date: 2025

# Load LXS common library (colors, separator, run_spinner, loggers, helpers)
LXS_RAW_BASE="${LXS_RAW_BASE:-https://git.hyko.cx/hykocx/lxs/raw/branch/main}"
_lib=$(curl -fsSL "${LXS_RAW_BASE}/lib/common.sh") || { echo "Failed to fetch lib/common.sh" >&2; exit 1; }
eval "$_lib"
unset _lib
export LXS_LOG_FILE="/tmp/lxs_harden.log"

require_root "$0" "$@"

set -u

# ═══════════════════════════════════════════════════════════════════════════
# Configuration
# ═══════════════════════════════════════════════════════════════════════════

DO_UFW=1
DO_FAIL2BAN=1
DO_SSH=1
DO_UNATTENDED=1
ASSUME_YES=0

for arg in "$@"; do
    case "$arg" in
        --no-ufw)         DO_UFW=0 ;;
        --no-fail2ban)    DO_FAIL2BAN=0 ;;
        --no-ssh)         DO_SSH=0 ;;
        --no-unattended)  DO_UNATTENDED=0 ;;
        -y|--yes)         ASSUME_YES=1 ;;
        -h|--help)
            cat <<EOF
Usage: harden.sh [options]

Options:
  --no-ufw          Skip UFW firewall setup
  --no-fail2ban     Skip fail2ban setup
  --no-ssh          Skip SSH key-only enforcement
  --no-unattended   Skip unattended-upgrades setup
  -y, --yes         Skip confirmation prompt
  -h, --help        Show this help
EOF
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $arg${NC}" >&2
            exit 1
            ;;
    esac
done

# ═══════════════════════════════════════════════════════════════════════════
# Pre-checks
# ═══════════════════════════════════════════════════════════════════════════

require_debian_ubuntu || exit 1

SSH_PORT=$(awk '/^[[:space:]]*Port[[:space:]]+/ {print $2; exit}' /etc/ssh/sshd_config 2>/dev/null)
SSH_PORT=${SSH_PORT:-22}

echo -e "${WHITE}${BOLD}LXS Server Hardening${NC}"
echo -e "${GRAY}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo "The following actions will be applied to this server:"
[ $DO_UFW -eq 1 ]        && echo "  • UFW firewall: default deny incoming, allow SSH on port ${SSH_PORT}"
[ $DO_FAIL2BAN -eq 1 ]   && echo "  • fail2ban: enable sshd jail (bantime 1h, maxretry 5)"
[ $DO_SSH -eq 1 ]        && echo "  • SSH: disable password auth (key-only), prohibit-password root login"
[ $DO_UNATTENDED -eq 1 ] && echo "  • unattended-upgrades: enable automatic security updates"
echo -e "${GRAY}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"

if [ $ASSUME_YES -ne 1 ]; then
    read -r -p "Proceed? [y/N] " reply
    case "$reply" in
        [yY]|[yY][eE][sS]) ;;
        *) info "Aborted."; exit 0 ;;
    esac
fi

export DEBIAN_FRONTEND=noninteractive

# ═══════════════════════════════════════════════════════════════════════════
# UFW
# ═══════════════════════════════════════════════════════════════════════════

setup_ufw() {
    info "Installing and configuring UFW..."
    apt-get update -qq
    apt-get install -y -qq ufw
    ufw --force reset >/dev/null
    ufw default deny incoming
    ufw default allow outgoing
    ufw allow "${SSH_PORT}/tcp" comment 'SSH'
    ufw --force enable
    ok "UFW enabled (SSH on ${SSH_PORT}/tcp allowed)"
}

# ═══════════════════════════════════════════════════════════════════════════
# fail2ban
# ═══════════════════════════════════════════════════════════════════════════

setup_fail2ban() {
    info "Installing and configuring fail2ban..."
    apt-get install -y -qq fail2ban
    cat > /etc/fail2ban/jail.local <<EOF
[DEFAULT]
bantime = 1h
findtime = 10m
maxretry = 5
backend = systemd

[sshd]
enabled = true
port = ${SSH_PORT}
EOF
    systemctl enable --now fail2ban
    systemctl restart fail2ban
    ok "fail2ban enabled (sshd jail active)"
}

# ═══════════════════════════════════════════════════════════════════════════
# SSH key-only
# ═══════════════════════════════════════════════════════════════════════════

has_authorized_key() {
    local user_home
    for user_home in /root "/home/${SUDO_USER:-}"; do
        [ -z "$user_home" ] && continue
        [ "$user_home" = "/home/" ] && continue
        if [ -s "${user_home}/.ssh/authorized_keys" ]; then
            return 0
        fi
    done
    return 1
}

setup_ssh() {
    info "Enforcing SSH key-only authentication..."
    if ! has_authorized_key; then
        err "No authorized_keys found for root or ${SUDO_USER:-current user}."
        err "Refusing to disable password auth — you would lock yourself out."
        err "Add a public key first, then re-run: harden.sh --no-ufw --no-fail2ban --no-unattended"
        return 1
    fi

    local cfg=/etc/ssh/sshd_config
    cp -n "$cfg" "${cfg}.lxs-backup"

    sed -i -E \
        -e 's/^[#[:space:]]*PasswordAuthentication[[:space:]]+.*/PasswordAuthentication no/' \
        -e 's/^[#[:space:]]*ChallengeResponseAuthentication[[:space:]]+.*/ChallengeResponseAuthentication no/' \
        -e 's/^[#[:space:]]*KbdInteractiveAuthentication[[:space:]]+.*/KbdInteractiveAuthentication no/' \
        -e 's/^[#[:space:]]*PermitRootLogin[[:space:]]+.*/PermitRootLogin prohibit-password/' \
        "$cfg"

    grep -q '^PasswordAuthentication ' "$cfg" || echo 'PasswordAuthentication no' >> "$cfg"
    grep -q '^PermitRootLogin '         "$cfg" || echo 'PermitRootLogin prohibit-password' >> "$cfg"

    if sshd -t 2>/tmp/lxs_sshd_test; then
        systemctl reload ssh 2>/dev/null || systemctl reload sshd
        ok "SSH configured for key-only auth (backup: ${cfg}.lxs-backup)"
    else
        err "sshd config test failed; restoring backup."
        cp "${cfg}.lxs-backup" "$cfg"
        cat /tmp/lxs_sshd_test >&2
        return 1
    fi
}

# ═══════════════════════════════════════════════════════════════════════════
# unattended-upgrades
# ═══════════════════════════════════════════════════════════════════════════

setup_unattended() {
    info "Installing and enabling unattended-upgrades..."
    apt-get install -y -qq unattended-upgrades apt-listchanges
    cat > /etc/apt/apt.conf.d/20auto-upgrades <<'EOF'
APT::Periodic::Update-Package-Lists "1";
APT::Periodic::Unattended-Upgrade "1";
APT::Periodic::AutocleanInterval "7";
EOF
    systemctl enable --now unattended-upgrades
    ok "unattended-upgrades enabled (security updates only by default)"
}

# ═══════════════════════════════════════════════════════════════════════════
# Run
# ═══════════════════════════════════════════════════════════════════════════

[ $DO_UFW -eq 1 ]        && setup_ufw
[ $DO_FAIL2BAN -eq 1 ]   && setup_fail2ban
[ $DO_SSH -eq 1 ]        && setup_ssh || true
[ $DO_UNATTENDED -eq 1 ] && setup_unattended

# ═══════════════════════════════════════════════════════════════════════════
# Summary
# ═══════════════════════════════════════════════════════════════════════════

echo ""
echo -e "${WHITE}${BOLD}Summary${NC}"
echo -e "${GRAY}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
command -v ufw       >/dev/null && echo "UFW:                 $(ufw status | head -1)"
systemctl is-active fail2ban   >/dev/null 2>&1 && echo "fail2ban:            active" || echo "fail2ban:            inactive"
echo "SSH password auth:   $(grep -E '^PasswordAuthentication' /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo unknown)"
echo "SSH root login:      $(grep -E '^PermitRootLogin'         /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo unknown)"
systemctl is-active unattended-upgrades >/dev/null 2>&1 && echo "unattended-upgrades: active" || echo "unattended-upgrades: inactive"
echo ""
ok "Hardening complete."