#!/usr/bin/env bash
set -euo pipefail

LOG_FILE="/var/log/ai-chatbot-hardening.log"
SSH_DROPIN_DIR="/etc/ssh/sshd_config.d"
SSH_DROPIN_FILE="$SSH_DROPIN_DIR/99-ai-chatbot-hardening.conf"
FAIL2BAN_JAIL_FILE="/etc/fail2ban/jail.d/ai-chatbot.local"
SYSCTL_FILE="/etc/sysctl.d/99-ai-chatbot-hardening.conf"

HARDEN_SSH_STRICT="${HARDEN_SSH_STRICT:-0}"
HARDEN_SKIP_PACKAGE_INSTALL="${HARDEN_SKIP_PACKAGE_INSTALL:-0}"
HARDEN_FIREWALL_MODE="${HARDEN_FIREWALL_MODE:-enforce}"
HARDEN_ALLOW_PORTS="${HARDEN_ALLOW_PORTS:-}"

mkdir -p "$(dirname "$LOG_FILE")"
touch "$LOG_FILE"
chmod 600 "$LOG_FILE" || true

log() {
  printf '[%s] [hardening] %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$*" | tee -a "$LOG_FILE"
}

require_root() {
  if [ "${EUID:-$(id -u)}" -ne 0 ]; then
    log "must run as root"
    exit 1
  fi
}

command_exists() {
  command -v "$1" >/dev/null 2>&1
}

detect_package_manager() {
  if command_exists apt-get; then
    echo "apt"
    return
  fi
  if command_exists dnf; then
    echo "dnf"
    return
  fi
  if command_exists yum; then
    echo "yum"
    return
  fi
  echo "none"
}

install_packages() {
  if [ "$HARDEN_SKIP_PACKAGE_INSTALL" = "1" ]; then
    log "skip package install (HARDEN_SKIP_PACKAGE_INSTALL=1)"
    return
  fi

  local pm
  pm="$(detect_package_manager)"
  case "$pm" in
    apt)
      export DEBIAN_FRONTEND=noninteractive
      apt-get update -y
      apt-get install -y --no-install-recommends fail2ban ufw ca-certificates curl jq
      ;;
    dnf)
      dnf install -y fail2ban firewalld ca-certificates curl jq
      ;;
    yum)
      yum install -y fail2ban firewalld ca-certificates curl jq
      ;;
    *)
      log "unsupported package manager, skip package install"
      ;;
  esac
}

configure_registry() {
  if command_exists npm; then
    npm config set registry https://registry.npmjs.org >/dev/null 2>&1 || true
  fi
  if command_exists pnpm; then
    pnpm config set registry https://registry.npmjs.org >/dev/null 2>&1 || true
  fi
  cat > /etc/npmrc << 'EOF'
registry=https://registry.npmjs.org/
audit=true
fund=false
EOF
  chmod 644 /etc/npmrc
  log "configured npm/pnpm registry to npmjs.org"
}

collect_allowed_ports() {
  local ports=()
  local ssh_ports=()
  local current_ssh_port=""

  if [ -n "${SSH_CONNECTION:-}" ]; then
    current_ssh_port="$(echo "$SSH_CONNECTION" | awk '{print $4}')"
    if [ -n "$current_ssh_port" ]; then
      ssh_ports+=("$current_ssh_port")
    fi
  fi

  if command_exists sshd; then
    while IFS= read -r port; do
      [ -n "$port" ] && ssh_ports+=("$port")
    done < <(sshd -T 2>/dev/null | awk '/^port / {print $2}' | sort -u)
  fi

  if [ "${#ssh_ports[@]}" -eq 0 ]; then
    ssh_ports=("22")
  fi

  ports+=("${ssh_ports[@]}")
  ports+=("80" "443")

  if [ -n "$HARDEN_ALLOW_PORTS" ]; then
    IFS=',' read -r -a extra <<< "$HARDEN_ALLOW_PORTS"
    ports+=("${extra[@]}")
  fi

  printf '%s\n' "${ports[@]}" | sed '/^$/d' | sed 's/ //g' | sort -u
}

configure_firewall() {
  if [ "$HARDEN_FIREWALL_MODE" = "skip" ]; then
    log "skip firewall hardening (HARDEN_FIREWALL_MODE=skip)"
    return
  fi

  local ports
  ports="$(collect_allowed_ports)"

  if command_exists ufw; then
    ufw default deny incoming || true
    ufw default allow outgoing || true
    while IFS= read -r port; do
      [ -n "$port" ] && ufw allow "${port}/tcp" || true
    done <<< "$ports"
    ufw --force enable || true
    log "ufw enabled with allowed TCP ports: $(echo "$ports" | tr '\n' ',' | sed 's/,$//')"
    return
  fi

  if command_exists firewall-cmd; then
    systemctl enable --now firewalld >/dev/null 2>&1 || true
    firewall-cmd --permanent --set-default-zone=public >/dev/null 2>&1 || true
    firewall-cmd --permanent --add-service=ssh >/dev/null 2>&1 || true
    firewall-cmd --permanent --add-service=http >/dev/null 2>&1 || true
    firewall-cmd --permanent --add-service=https >/dev/null 2>&1 || true
    while IFS= read -r port; do
      [ -n "$port" ] && firewall-cmd --permanent --add-port="${port}/tcp" >/dev/null 2>&1 || true
    done <<< "$ports"
    firewall-cmd --reload >/dev/null 2>&1 || true
    log "firewalld enabled with allowed TCP ports: $(echo "$ports" | tr '\n' ',' | sed 's/,$//')"
    return
  fi

  log "no supported firewall command found"
}

has_root_authorized_keys() {
  [ -s /root/.ssh/authorized_keys ]
}

configure_sshd() {
  mkdir -p "$SSH_DROPIN_DIR"

  local permit_root="yes"
  local password_auth="yes"

  if [ "$HARDEN_SSH_STRICT" = "1" ]; then
    if has_root_authorized_keys; then
      permit_root="prohibit-password"
      password_auth="no"
      log "strict SSH mode enabled (key login required)"
    else
      log "strict SSH mode requested but no /root/.ssh/authorized_keys, fallback to password auth"
    fi
  fi

  local tmp_file
  tmp_file="$(mktemp)"
  cat > "$tmp_file" << EOF
# Managed by ai-chatbot scripts/security/harden-server.sh
Protocol 2
PubkeyAuthentication yes
PasswordAuthentication $password_auth
PermitRootLogin $permit_root
PermitEmptyPasswords no
ChallengeResponseAuthentication no
UsePAM yes
MaxAuthTries 3
LoginGraceTime 30
ClientAliveInterval 300
ClientAliveCountMax 2
X11Forwarding no
AllowTcpForwarding no
AllowAgentForwarding no
EOF

  local rollback_file=""
  if [ -f "$SSH_DROPIN_FILE" ]; then
    cp -a "$SSH_DROPIN_FILE" "${SSH_DROPIN_FILE}.bak.$(date '+%Y%m%d%H%M%S')" || true
    rollback_file="$(mktemp)"
    cp -a "$SSH_DROPIN_FILE" "$rollback_file"
  fi
  install -m 644 "$tmp_file" "$SSH_DROPIN_FILE"
  rm -f "$tmp_file"

  if command_exists sshd && sshd -t; then
    systemctl reload sshd >/dev/null 2>&1 || systemctl reload ssh >/dev/null 2>&1 || systemctl restart sshd >/dev/null 2>&1 || systemctl restart ssh >/dev/null 2>&1 || true
    log "sshd hardened and reloaded"
  else
    if [ -n "$rollback_file" ] && [ -f "$rollback_file" ]; then
      install -m 644 "$rollback_file" "$SSH_DROPIN_FILE"
    else
      rm -f "$SSH_DROPIN_FILE"
    fi
    log "sshd config test failed, rollback applied"
  fi
  [ -n "$rollback_file" ] && rm -f "$rollback_file" || true
}

configure_fail2ban() {
  if ! command_exists fail2ban-client; then
    log "fail2ban not installed, skip"
    return
  fi

  local nginx_jail_enabled="false"
  if [ -f /var/log/nginx/access.log ] || [ -f /var/log/nginx/error.log ]; then
    nginx_jail_enabled="true"
  fi

  mkdir -p "$(dirname "$FAIL2BAN_JAIL_FILE")"
  cat > "$FAIL2BAN_JAIL_FILE" << EOF
[DEFAULT]
bantime = 1h
findtime = 10m
maxretry = 5
backend = auto
bantime.increment = true
bantime.factor = 2
bantime.maxtime = 24h
ignoreip = 127.0.0.1/8 ::1

[sshd]
enabled = true
backend = systemd
journalmatch = _COMM=sshd
mode = normal

[nginx-http-auth]
enabled = ${nginx_jail_enabled}

[nginx-botsearch]
enabled = ${nginx_jail_enabled}
EOF

  systemctl enable fail2ban >/dev/null 2>&1 || true
  systemctl restart fail2ban >/dev/null 2>&1 || true

  if systemctl is-active --quiet fail2ban; then
    log "fail2ban configured"
  else
    log "fail2ban failed to start, check: journalctl -u fail2ban -n 100 --no-pager"
  fi
}

configure_sysctl() {
  cat > "$SYSCTL_FILE" << 'EOF'
net.ipv4.conf.all.accept_redirects = 0
net.ipv4.conf.default.accept_redirects = 0
net.ipv4.conf.all.send_redirects = 0
net.ipv4.conf.default.send_redirects = 0
net.ipv4.conf.all.rp_filter = 1
net.ipv4.conf.default.rp_filter = 1
net.ipv4.tcp_syncookies = 1
net.ipv4.icmp_echo_ignore_broadcasts = 1
net.ipv4.icmp_ignore_bogus_error_responses = 1
kernel.kptr_restrict = 2
kernel.dmesg_restrict = 1
fs.protected_hardlinks = 1
fs.protected_symlinks = 1
EOF
  sysctl --system >/dev/null 2>&1 || true
  log "sysctl hardening applied"
}

main() {
  require_root
  log "start server hardening"
  install_packages
  configure_registry
  configure_sshd
  configure_fail2ban
  configure_firewall
  configure_sysctl
  log "server hardening done"
}

main "$@"