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

# ------------------------------------------------------------
# PANIC GUARD (iptables) — anti-flood interactif & réversible
# - Menu + paramètres persistants
# - Dry-run + résumé + confirmation
# - Snapshot des règles ajoutées
# - Rollback (annule seulement les règles ajoutées par ce script)
# ------------------------------------------------------------

APP_NAME="panic-guard"
BASE_DIR="/var/lib/${APP_NAME}"
CONF_FILE="${BASE_DIR}/config.conf"
LOG_FILE="/var/log/${APP_NAME}.log"
SNAPSHOT_FILE="${BASE_DIR}/snapshot.rules"
LAST_RUN_FILE="${BASE_DIR}/last_run.txt"

DEFAULT_PORT="443"
DEFAULT_TOP="20"
DEFAULT_MIN_CONN="30"
DEFAULT_MODE="ip"        # ip | /24 | /16
DEFAULT_ACTION="DROP"    # DROP | REJECT
DEFAULT_CHAIN="INPUT"    # INPUT (entrant)
DEFAULT_DRYRUN="1"
DEFAULT_WHITELIST=""     # ex: "1.2.3.4,5.6.7.0/24"

# ---------- helpers ----------
need_root() { [[ "${EUID:-$(id -u)}" -eq 0 ]] || { echo "Run as root: sudo $0"; exit 1; }; }
cmd_exists() { command -v "$1" >/dev/null 2>&1; }

now() { date "+%Y-%m-%d %H:%M:%S"; }
log() { echo "[$(now)] $*" | tee -a "$LOG_FILE" >/dev/null; }
say() { echo "$*"; }
hr() { echo "------------------------------------------------------------"; }

ensure_dirs() {
  mkdir -p "$BASE_DIR"
  touch "$LOG_FILE"
  chmod 700 "$BASE_DIR"
}

# ---------- config ----------
load_config() {
  PORT="$DEFAULT_PORT"
  TOP="$DEFAULT_TOP"
  MIN_CONN="$DEFAULT_MIN_CONN"
  MODE="$DEFAULT_MODE"
  ACTION="$DEFAULT_ACTION"
  CHAIN="$DEFAULT_CHAIN"
  DRYRUN="$DEFAULT_DRYRUN"
  WHITELIST="$DEFAULT_WHITELIST"

  if [[ -f "$CONF_FILE" ]]; then
    # shellcheck disable=SC1090
    source "$CONF_FILE"
  fi
}

save_config() {
  cat > "$CONF_FILE" <<EOF
# ${APP_NAME} config (auto-saved)
PORT="${PORT}"
TOP="${TOP}"
MIN_CONN="${MIN_CONN}"
MODE="${MODE}"
ACTION="${ACTION}"
CHAIN="${CHAIN}"
DRYRUN="${DRYRUN}"
WHITELIST="${WHITELIST}"
EOF
  chmod 600 "$CONF_FILE"
  log "Saved config to $CONF_FILE"
}

show_config() {
  hr
  say "${APP_NAME} — current config"
  hr
  printf "PORT       : %s\n" "$PORT"
  printf "TOP        : %s\n" "$TOP"
  printf "MIN_CONN   : %s\n" "$MIN_CONN"
  printf "MODE       : %s\n" "$MODE"
  printf "ACTION     : %s\n" "$ACTION"
  printf "CHAIN      : %s\n" "$CHAIN"
  printf "DRYRUN     : %s\n" "$DRYRUN"
  printf "WHITELIST  : %s\n" "${WHITELIST:-<empty>}"
  hr
}

validate_config() {
  [[ "$PORT" =~ ^[0-9]+$ ]] || { say "Invalid PORT"; return 1; }
  [[ "$TOP" =~ ^[0-9]+$ ]] || { say "Invalid TOP"; return 1; }
  [[ "$MIN_CONN" =~ ^[0-9]+$ ]] || { say "Invalid MIN_CONN"; return 1; }
  [[ "$MODE" == "ip" || "$MODE" == "/24" || "$MODE" == "/16" ]] || { say "Invalid MODE"; return 1; }
  [[ "$ACTION" == "DROP" || "$ACTION" == "REJECT" ]] || { say "Invalid ACTION"; return 1; }
  [[ "$CHAIN" == "INPUT" || "$CHAIN" == "OUTPUT" ]] || { say "Invalid CHAIN"; return 1; }
  [[ "$DRYRUN" == "0" || "$DRYRUN" == "1" ]] || { say "Invalid DRYRUN"; return 1; }
  return 0
}

# ---------- whitelist ----------
is_whitelisted() {
  local ip="$1"
  [[ -z "${WHITELIST:-}" ]] && return 1
  IFS=',' read -r -a items <<< "$WHITELIST"
  for item in "${items[@]}"; do
    item="$(echo "$item" | xargs)"
    [[ -z "$item" ]] && continue
    if [[ "$item" == "$ip" ]]; then
      return 0
    fi
    if [[ "$item" =~ ^([0-9]+\.[0-9]+\.[0-9]+)\.0/24$ ]]; then
      [[ "$ip" == "${BASH_REMATCH[1]}."* ]] && return 0
    fi
    if [[ "$item" =~ ^([0-9]+\.[0-9]+)\.0\.0/16$ ]]; then
      [[ "$ip" == "${BASH_REMATCH[1]}."* ]] && return 0
    fi
  done
  return 1
}

to_target() {
  local ip="$1"
  case "$MODE" in
    ip) echo "$ip" ;;
    /24) echo "$(echo "$ip" | awk -F. '{print $1"."$2"."$3".0/24"}')" ;;
    /16) echo "$(echo "$ip" | awk -F. '{print $1"."$2".0.0/16"}')" ;;
  esac
}

# ---------- detection ----------
top_ips() {
  # Uses ss to list connections where local port == PORT
  # Extract peer IP, count, sort desc, show TOP.
  ss -Hant "sport = :$PORT" 2>/dev/null \
    | awk '{print