#!/usr/bin/env bash
# block_range.sh — bloque un range IPv4 (/CIDR) contre flood (SYN) + purge conntrack + hardening kernel (optionnel)
# Usage:
#   sudo ./block_range.sh 5.181.17.0/24 apply
#   sudo ./block_range.sh 5.181.17.0/24 status
#   sudo ./block_range.sh 5.181.17.0/24 rollback
#
# Options:
#   --no-sysctl      : n'applique pas les sysctl anti-SYN (par défaut: appliqués en runtime)
#   --persist-sysctl : persiste les sysctl dans /etc/sysctl.d/99-anti-syn.conf
#   --no-output-drop : ne met pas de règle OUTPUT (par défaut: on coupe aussi toute réponse vers le range)
#   --dry-run        : affiche ce qui serait fait sans exécuter

set -euo pipefail

COLOR=1
if [[ -n "${NO_COLOR:-}" ]] || [[ ! -t 1 ]]; then COLOR=0; fi
red(){   [[ $COLOR -eq 1 ]] && printf "\033[31m%s\033[0m\n" "$*" || printf "%s\n" "$*"; }
green(){ [[ $COLOR -eq 1 ]] && printf "\033[32m%s\033[0m\n" "$*" || printf "%s\n" "$*"; }
yellow(){[[ $COLOR -eq 1 ]] && printf "\033[33m%s\033[0m\n" "$*" || printf "%s\n" "$*"; }
info(){  printf "%s\n" "$*"; }

die(){ red "ERROR: $*"; exit 1; }

need_root(){
  [[ "${EUID:-$(id -u)}" -eq 0 ]] || die "Lance en root: sudo $0 ..."
}

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

DRY=0
DO_SYSCTL=1
PERSIST_SYSCTL=0
DO_OUTPUT_DROP=1

usage(){
  cat <<'EOF'
block_range.sh — bloquer un range IPv4 contre flood (SYN) + purge conntrack + (optionnel) durcissement sysctl

Usage:
  sudo ./block_range.sh <CIDR> apply [options]
  sudo ./block_range.sh <CIDR> status
  sudo ./block_range.sh <CIDR> rollback

Exemples:
  sudo ./block_range.sh 5.181.17.0/24 apply
  sudo ./block_range.sh 5.181.17.0/24 apply --persist-sysctl
  sudo ./block_range.sh 5.181.17.0/24 status
  sudo ./block_range.sh 5.181.17.0/24 rollback

Options:
  --no-sysctl       n'applique pas les sysctl anti-SYN
  --persist-sysctl  persiste les sysctl dans /etc/sysctl.d/99-anti-syn.conf (en plus du runtime)
  --no-output-drop  ne bloque pas OUTPUT vers le range
  --dry-run         n'exécute rien, affiche seulement
EOF
}

# --- Args ---
CIDR="${1:-}"
ACTION="${2:-}"

shift $(( $#>0 ? 1 : 0 )) || true
shift $(( $#>0 ? 1 : 0 )) || true

while [[ $# -gt 0 ]]; do
  case "$1" in
    --dry-run) DRY=1 ;;
    --no-sysctl) DO_SYSCTL=0 ;;
    --persist-sysctl) PERSIST_SYSCTL=1 ;;
    --no-output-drop) DO_OUTPUT_DROP=0 ;;
    -h|--help) usage; exit 0 ;;
    *) die "Option inconnue: $1" ;;
  esac
  shift
done

if [[ -z "$CIDR" || -z "$ACTION" ]]; then
  usage
  exit 1
fi

# Simple validation CIDR IPv4 (0-32)
if [[ ! "$CIDR" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}/([0-9]|[12][0-9]|3[0-2])$ ]]; then
  die "CIDR invalide: $CIDR (ex: 5.181.17.0/24)"
fi

# Check octets range 0-255
IFS='/' read -r IP _MASK <<<"$CIDR"
IFS='.' read -r o1 o2 o3 o4 <<<"$IP"
for o in "$o1" "$o2" "$o3" "$o4"; do
  [[ "$o" -ge 0 && "$o" -le 255 ]] || die "IP invalide dans CIDR: $CIDR"
done

need_root
have iptables || die "iptables introuvable"

TAG="IDEO_BLOCKRANGE"
# state file (per CIDR)
STATE_DIR="/var/tmp"
SAFE_CIDR="${CIDR//\//_}"
STATE_FILE="$STATE_DIR/blockrange_${SAFE_CIDR}.state"

run(){
  if [[ $DRY -eq 1 ]]; then
    yellow "[dry-run] $*"
    return 0
  fi
  eval "$@"
}

save_rule_cmd(){
  # Stocke les commandes de rollback exactes
  # Ex: iptables -D INPUT ... (on les construira à l'ajout)
  echo "$*" >> "$STATE_FILE"
}

iptables_has_rule(){
  # $1 = chain, rest = rule spec
  local chain="$1"; shift
  iptables -C "$chain" "$@" >/dev/null 2>&1
}

add_rule_once(){
  # $1 chain, then rule-spec (without -A/-I)
  local chain="$1"; shift
  if iptables_has_rule "$chain" "$@"; then
    info "• Règle déjà présente: $chain $*"
    return 0
  fi
  # Insert at top of chain for priority
  run "iptables -I $chain 1 $*"
  # Prepare rollback command (delete the exact same spec)
  save_rule_cmd "iptables -D $chain $*"
  green "✓ Ajout: $chain $*"
}

apply_sysctl(){
  [[ $DO_SYSCTL -eq 1 ]] || { info "• Sysctl ignoré (--no-sysctl)"; return 0; }
  info "• Appliquer sysctl anti-SYN (runtime)"
  run "sysctl -w net.ipv4.tcp_syncookies=1 >/dev/null"
  run "sysctl -w net.ipv4.tcp_max_syn_backlog=8192 >/dev/null"
  run "sysctl -w net.ipv4.tcp_synack_retries=2 >/dev/null"
  run "sysctl -w net.ipv4.tcp_abort_on_overflow=1 >/dev/null"

  if [[ $PERSIST_SYSCTL -eq 1 ]]; then
    info "• Persistance sysctl dans /etc/sysctl.d/99-anti-syn.conf"
    local f="/etc/sysctl.d/99-anti-syn.conf"
    if [[ $DRY -eq 1 ]]; then
      yellow "[dry-run] écriture $f"
    else
      cat > "$f" <<'EOF'
# Anti-SYN flood hardening (runtime configurable via sysctl)
net.ipv4.tcp_syncookies=1
net.ipv4.tcp_max_syn_backlog=8192
net.ipv4.tcp_synack_retries=2
net.ipv4.tcp_abort_on_overflow=1
EOF
      sysctl --system >/dev/null || true
    fi
    green "✓ Sysctl persisté"
  fi
}

purge_conntrack(){
  if have conntrack; then
    info "• Purge conntrack (sessions existantes) pour $CIDR"
    run "conntrack -D -s $CIDR >/dev/null 2>&1 || true"
    run "conntrack -D -d $CIDR >/dev/null 2>&1 || true"
    green "✓ Conntrack purgé"
  else
    yellow "• conntrack introuvable (installe: apt install conntrack) — purge ignorée"
  fi
}

status(){
  info "=== STATUS pour $CIDR ==="
  info "--- iptables (règles qui matchent le range) ---"
  iptables -L INPUT -n -v --line-numbers | grep -F "$CIDR" || true
  iptables -L FORWARD -n -v --line-numbers | grep -F "$CIDR" || true
  iptables -L OUTPUT -n -v --line-numbers | grep -F "$CIDR" || true

  if have conntrack; then
    info "--- conntrack (compte rapide) ---"
    conntrack -L 2>/dev/null | grep -F "${CIDR%/*}" | head -n 5 || true
  fi

  info "--- ss (SYN-SENT count) ---"
  if have ss; then
    ss -ant state syn-sent 2>/dev/null | wc -l | awk '{print "SYN-SENT:", $1}'
  else
    yellow "ss introuvable"
  fi

  if [[ -f "$STATE_FILE" ]]; then
    info "--- state file ---"
    echo "State: $STATE_FILE"
    wc -l "$STATE_FILE" | awk '{print "Rollback commands:", $1}'
  else
    info "State: (aucun) — script pas appliqué (ou déjà rollback)."
  fi
}

apply(){
  info "=== APPLY blocage pour $CIDR ==="
  # Fresh state file (append-safe)
  if [[ $DRY -eq 0 ]]; then
    : > "$STATE_FILE"
    chmod 600 "$STATE_FILE" || true
  else
    yellow "[dry-run] state file: $STATE_FILE"
  fi

  # 1) Priorité: drop SYN en premier (évite tracking/charge)
  add_rule_once INPUT  -s "$CIDR" -p tcp --syn -m comment --comment "$TAG syn" -j DROP

  # 2) Drop global (tous protocoles) en entrée + forward
  add_rule_once INPUT   -s "$CIDR" -m comment --comment "$TAG in" -j DROP
  add_rule_once FORWARD -s "$CIDR" -m comment --comment "$TAG fwd" -j DROP

  # 3) Optionnel: éviter toute réponse vers eux (réduit amplification)
  if [[ $DO_OUTPUT_DROP -eq 1 ]]; then
    add_rule_once OUTPUT  -d "$CIDR" -m comment --comment "$TAG out" -j DROP
  else
    info "• OUTPUT non bloqué (--no-output-drop)"
  fi

  # 4) Purge conntrack (tuer l'existant)
  purge_conntrack

  # 5) Sysctl anti-flood (runtime + option persist)
  apply_sysctl

  green "✅ APPLY terminé. Pour vérifier: sudo $0 $CIDR status"
  info "Rollback: sudo $0 $CIDR rollback"
}

rollback(){
  info "=== ROLLBACK pour $CIDR ==="
  if [[ ! -f "$STATE_FILE" ]]; then
    yellow "• Pas de state file: $STATE_FILE"
    yellow "  -> Je tente quand même un rollback best-effort via suppression des règles connues."
    # Best-effort deletion (ignore errors)
    run "iptables -D INPUT -s $CIDR -p tcp --syn -m comment --comment \"$TAG syn\" -j DROP >/dev/null 2>&1 || true"
    run "iptables -D INPUT -s $CIDR -m comment --comment \"$TAG in\" -j DROP >/dev/null 2>&1 || true"
    run "iptables -D FORWARD -s $CIDR -m comment --comment \"$TAG fwd\" -j DROP >/dev/null 2>&1 || true"
    run "iptables -D OUTPUT -d $CIDR -m comment --comment \"$TAG out\" -j DROP >/dev/null 2>&1 || true"
    green "✓ Best-effort rollback terminé"
    return 0
  fi

  # Execute stored rollback commands in reverse order for safety
  if [[ $DRY -eq 1 ]]; then
    yellow "[dry-run] exécuter rollback depuis $STATE_FILE (ordre inverse)"
    tac "$STATE_FILE"
    exit 0
  fi

  tac "$STATE_FILE" | while IFS= read -r cmd; do
    [[ -z "$cmd" ]] && continue
    bash -c "$cmd" >/dev/null 2>&1 || true
    green "✓ $cmd"
  done

  # Purge conntrack again (optional)
  purge_conntrack

  rm -f "$STATE_FILE" || true
  green "✅ ROLLBACK terminé."
}

case "$ACTION" in
  apply) apply ;;
  rollback|revert|undo) rollback ;;
  status|check) status ;;
  *) die "Action invalide: $ACTION (apply|status|rollback)" ;;
esac