#!/bin/bash

# LXS - Server Hardening
# Description: Apply baseline security hardening (UFW + fail2ban + 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_UNATTENDED=1
ASSUME_YES=0

for arg in "$@"; do
    case "$arg" in
        --no-ufw)         DO_UFW=0 ;;
        --no-fail2ban)    DO_FAIL2BAN=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-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
require_disk_space 500 || exit 1

# Read the effective Port from sshd_config + any drop-in under sshd_config.d/
sshd_effective_port() {
    local files=(/etc/ssh/sshd_config)
    if compgen -G "/etc/ssh/sshd_config.d/*.conf" >/dev/null; then
        files+=(/etc/ssh/sshd_config.d/*.conf)
    fi
    awk '/^[[:space:]]*Port[[:space:]]+/ {print $2}' "${files[@]}" 2>/dev/null | tail -1
}

SSH_PORT=$(sshd_effective_port)
SSH_PORT=${SSH_PORT:-22}

echo -e "${WHITE}${BOLD}LXS Server Hardening${NC}"
show_separator
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_UNATTENDED -eq 1 ] && echo "  • unattended-upgrades: enable automatic security updates"
show_separator

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

apt_noninteractive
wait_for_apt || exit 1

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

# ═══════════════════════════════════════════════════════════════════════════
# 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_UNATTENDED -eq 1 ] && setup_unattended

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

echo ""
echo -e "${WHITE}${BOLD}Summary${NC}"
show_separator
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"
systemctl is-active unattended-upgrades >/dev/null 2>&1 && echo "unattended-upgrades: active" || echo "unattended-upgrades: inactive"
echo ""
ok "Hardening complete."