#!/usr/bin/env bash
# 50-hardening — system defaults. Each sub-step y/N, all idempotent.

step_name="hardening"
step_desc="ssh(ufw-safe), ufw, updates, unattended, tz, hostname, swap, locale, fail2ban"
step_run="run_hardening"

run_hardening() {
  log "hardening system defaults"
  _h_ssh
  _h_ufw
  _h_updates
  _h_tz_hostname
  _h_swap
  _h_locale
  _h_fail2ban
  ok "hardening pass complete"
}

# ── SSH: permit root key login, disable password auth ──────────────────────
_h_ssh() {
  if ! command -v sshd >/dev/null 2>&1 && ! [[ -f /etc/ssh/sshd_config ]]; then
    warn "no sshd found — skipping SSH hardening"; return
  fi
  if ! yn "Harden sshd: PermitRootLogin prohibit-password, disable password auth?

  WARNING: only say yes if you can already log in with a public key,
  else you may lock yourself out. A backup will be made and sshd -t
  is run before reload. Continue?" n; then
    echo "  skip ssh hardening"; return
  fi

  local cfg="/etc/ssh/sshd_config"
  sudo_ cp -a "$cfg" "${cfg}.bak.$(date +%Y%m%d-%H%M%S)"

  sshd_set() { # key value — replace or append a top-level directive
    local k="$1" v="$2" cfg="$3"
    if grep -qiE "^\s*#?\s*${k}\b" "$cfg"; then
      sudo_ sed -i -E "s|^\s*#?\s*${k}\b.*|${k} ${v}|" "$cfg"
    else
      echo "${k} ${v}" | sudo_ tee -a "$cfg" >/dev/null
    fi
  }
  sshd_set PermitRootLogin prohibit-password "$cfg"
  sshd_set PasswordAuthentication no "$cfg"
  sshd_set KbdInteractiveAuthentication no "$cfg"
  sshd_set PubkeyAuthentication yes "$cfg"

  if sudo_ sshd -t 2>/dev/null; then
    if command -v systemctl >/dev/null 2>&1 && systemctl is-active ssh sshd 2>/dev/null | grep -q active; then
      sudo_ systemctl reload ssh 2>/dev/null || sudo_ systemctl reload sshd 2>/dev/null || sudo_ systemctl restart ssh sshd 2>/dev/null
    else
      sudo_ systemctl restart ssh 2>/dev/null || sudo_ systemctl restart sshd 2>/dev/null || \
        sudo_ service ssh restart 2>/dev/null || sudo_ service sshd restart 2>/dev/null || warn "could not reload sshd"
    fi
    ok "sshd hardened (root key-only, no passwords). Keep your current session open and test a new login before closing it."
  else
    err "sshd -t failed — restoring backup and aborting ssh hardening"
    sudo_ cp -a "${cfg}.bak.$(date +%Y%m%d-%H%M%S)" "$cfg" 2>/dev/null || true
    warn "ssh hardening aborted, config restored"
  fi
}

# ── ufw: deny incoming, allow 22 + mosh UDP + angie 80/443 ──────────────────
_h_ufw() {
  command -v ufw >/dev/null 2>&1 || { warn "ufw not installed — skipping"; return; }
  if ! yn "Enable ufw (deny incoming, allow 22, mosh 60000:61000/udp, 80+443/tcp)?" y; then
    echo "  skip ufw"; return
  fi
  sudo_ ufw --force reset >/dev/null 2>&1 || true
  sudo_ ufw default deny incoming
  sudo_ ufw default allow outgoing
  sudo_ ufw allow 22/tcp                     comment 'ssh'
  sudo_ ufw allow 60000:61000/udp            comment 'mosh'
  sudo_ ufw allow 80/tcp                     comment 'angie http'
  sudo_ ufw allow 443/tcp                    comment 'angie https'
  sudo_ ufw --force enable
  ok "ufw enabled. status:"
  sudo_ ufw status verbose | sed 's/^/  /'
}

# ── updates + unattended-upgrades (debian) ─────────────────────────────────
_h_updates() {
  if ! yn "Run a system upgrade now?" y; then echo "  skip updates"; return; fi
  case "$DISTRO" in
    arch)
      sudo_ pacman -Syu --noconfirm
      ;;
    debian)
      sudo_ apt-get update -y
      sudo_ apt-get upgrade -y
      if yn "Enable unattended-upgrades (daily security auto-patches)?" y; then
        sudo_ apt-get install -y unattended-upgrades apt-listchanges
        sudo_ dpkg-reconfigure -f noninteractive unattended-upgrades 2>/dev/null || true
        ok "unattended-upgrades enabled"
      fi
      ;;
  esac
}

# ── timezone + hostname ─────────────────────────────────────────────────────
_h_tz_hostname() {
  if command -v timedatectl >/dev/null 2>&1; then
    if yn "Set timezone (default UTC)?" y; then
      local tz="${REPLY_TZ:-}"
      read -rp "  timezone [UTC]: " tz
      [[ -z "$tz" ]] && tz="UTC"
      sudo_ timedatectl set-timezone "$tz" && ok "timezone set to $tz"
    fi
  else
    echo "  timedatectl missing — skipping tz"
  fi

  if yn "Set a hostname now? (optional)" n; then
    read -rp "  hostname: " hn
    if [[ -n "$hn" ]] && command -v hostnamectl >/dev/null 2>&1; then
      sudo_ hostnamectl set-hostname "$hn" && ok "hostname set to $hn"
    fi
  fi
}

# ── swapfile if none and RAM is low ─────────────────────────────────────────
_h_swap() {
  if [[ "$(swapon --show --noheadings | wc -l)" -gt 0 ]]; then
    ok "swap already present — skipping"; return
  fi
  local mem_mb; mem_mb="$(awk '/MemTotal/ {printf "%d", $2/1024}' /proc/meminfo)"
  if (( mem_mb > 2048 )); then
    echo "  RAM ${mem_mb}MB > 2GB — skipping swapfile"; return
  fi
  if ! yn "Create a 2G swapfile (RAM is ${mem_mb}MB)?" y; then echo "  skip swap"; return; fi
  sudo_ fallocate -l 2G /swapfile || sudo_ dd if=/dev/zero of=/swapfile bs=1M count=2048
  sudo_ chmod 600 /swapfile
  sudo_ mkswap /swapfile
  sudo_ swapon /swapfile
  if ! grep -q '^/swapfile' /etc/fstab; then
    echo '/swapfile none swap sw 0 0' | sudo_ tee -a /etc/fstab >/dev/null
  fi
  ok "2G swapfile created and enabled"
}

# ── locale en_US.UTF-8 ───────────────────────────────────────────────────────
_h_locale() {
  case "$DISTRO" in
    arch)
      if ! grep -q '^en_US.UTF-8' /etc/locale.gen 2>/dev/null; then
        echo " localegen already has en_US.UTF-8 or file absent — touching"
      fi
      if [[ -f /etc/locale.gen ]]; then
        sudo_ sed -i 's/^#en_US.UTF-8/en_US.UTF-8/' /etc/locale.gen
        sudo_ locale-gen
      fi
      ;;
    debian)
      if ! locale -a 2>/dev/null | grep -qi 'en_US.utf8'; then
        echo 'en_US.UTF-8 UTF-8' | sudo_ tee -a /etc/locale.gen >/dev/null
        sudo_ locale-gen
        ok "generated en_US.UTF-8"
      else
        ok "en_US.UTF-8 already generated"
      fi
      ;;
  esac
}

# ── fail2ban (best-effort) ───────────────────────────────────────────────────
_h_fail2ban() {
  if ! yn "Install + enable fail2ban (sshd jail by default)?" y; then echo "  skip fail2ban"; return; fi
  case "$DISTRO" in
    arch)   sudo_ pacman -Sy --noconfirm --needed fail2ban ;;
    debian) sudo_ apt-get install -y fail2ban ;;
  esac
  # local jail for sshd (works across distros)
  sudo_ tee /etc/fail2ban/jail.local >/dev/null <<'EOF'
[DEFAULT]
backend = systemd
bantime = 1h
findtime = 10m
maxretry = 5

[sshd]
enabled = true
EOF
  if command -v systemctl >/dev/null 2>&1 && systemctl is-system-running >/dev/null 2>&1; then
    sudo_ systemctl enable --now fail2ban 2>/dev/null || warn "could not enable fail2ban"
  fi
  ok "fail2ban installed + sshd jail enabled"
}