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

CONF_FILE=/etc/evobgp/firewall.conf
LOG_FILE=/var/log/evobgp-firewall.log
STATE_DIR=/var/lib/evobgp-firewall
HASH_FILE="${STATE_DIR}/last_hash"
PREFIX_FILE="${STATE_DIR}/last_prefixes.txt"

log() { echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) $*" | tee -a "$LOG_FILE"; }

if [[ ! -f "$CONF_FILE" ]]; then
  log "missing $CONF_FILE"
  exit 1
fi
# shellcheck disable=SC1090
source "$CONF_FILE"

: "${EVOBGP_CP_URL:?}"
: "${CLIENT_TOKEN:?}"
CLIENT_TOKEN="${CLIENT_TOKEN//$'\r'/}"
CLIENT_TOKEN="${CLIENT_TOKEN//$'\n'/}"

mkdir -p "$STATE_DIR"
BACKEND="${KERNEL_BACKEND:-auto}"

curl_get_blocklist_file() {
  local url="$1"
  local dest="$2"
  local code
  code=$(curl -sS -o "$dest" -w "%{http_code}" \
    -H "Authorization: Bearer ${CLIENT_TOKEN}" \
    -H "Accept: application/json" \
    "${url}/v1/firewall/blocklist") || return 1
  if [[ "$code" == "403" ]]; then
    log "pending approval"
    return 2
  fi
  if [[ "$code" != "200" ]]; then
    log "blocklist HTTP $code from $url"
    return 1
  fi
  return 0
}

try_fetch_blocklist() {
  local urls=()
  if [[ -n "${EVOBGP_FAILOVER_URLS:-}" ]]; then
    IFS=',' read -r -a urls <<<"$EVOBGP_FAILOVER_URLS"
  else
    urls=("${EVOBGP_CP_URL%/}")
  fi
  local u
  for u in "${urls[@]}"; do
    u="${u// /}"
    u="${u%/}"
    local rc=0
    curl_get_blocklist_file "$u" "$PREFIX_FILE" || rc=$?
    if [[ "$rc" == 2 ]]; then
      exit 0
    fi
    if [[ "$rc" == 0 ]]; then
      CP_HIT="$u"
      return 0
    fi
  done
  return 1
}

parse_blocklist_file() {
  local f="$1"
  if [[ ! -s "$f" ]]; then
    log "blocklist file empty: $f"
    return 1
  fi
  if command -v jq >/dev/null 2>&1; then
    HASH=$(jq -r '.hash // empty' "$f")
    TOTAL=$(jq -r '.total // 0' "$f")
    mapfile -t PREFIXES < <(jq -r '.prefixes[]? // empty' "$f")
    return 0
  fi
  if command -v python3 >/dev/null 2>&1; then
    local parsed
    parsed=$(python3 - "$f" <<'PY'
import json, sys
with open(sys.argv[1], encoding="utf-8") as fh:
    data = json.load(fh)
print(data.get("hash") or "")
print(data.get("total") or 0)
for p in data.get("prefixes") or []:
    if p:
        print(p)
PY
    )
    HASH=$(echo "$parsed" | sed -n '1p')
    TOTAL=$(echo "$parsed" | sed -n '2p')
    mapfile -t PREFIXES < <(echo "$parsed" | sed -n '3,$p')
    return 0
  fi
  HASH=$(grep -o '"hash"[[:space:]]*:[[:space:]]*"[^"]*"' "$f" | head -1 | sed 's/.*"\(sha256:[^"]*\)".*/\1/')
  TOTAL=$(grep -o '"total"[[:space:]]*:[[:space:]]*[0-9]*' "$f" | head -1 | grep -o '[0-9]*$' || true)
  mapfile -t PREFIXES < <(grep -oE '"[0-9]+(\.[0-9]+){3}/[0-9]+"' "$f" | tr -d '"' || true)
  return 0
}

nft_join_elements() {
  local out="" p
  for p in "$@"; do
    if [[ -n "$out" ]]; then
      out+=", "
    fi
    out+="$p"
  done
  printf '%s' "$out"
}

nft_add_v4_chunk() {
  local table=$1 name=$2
  shift 2
  local joined
  joined=$(nft_join_elements "$@")
  if nft add element "$table" "$name" v4 "{ ${joined} }" 2>>"$LOG_FILE"; then
    return 0
  fi
  log "nft batch add failed (chunk=$#), retrying one-by-one"
  local p ok=0
  for p in "$@"; do
    if nft add element "$table" "$name" v4 "{ $p }" 2>>"$LOG_FILE"; then
      ok=$((ok + 1))
    fi
  done
  [[ "$ok" -gt 0 ]]
}

if ! try_fetch_blocklist; then
  log "all endpoints failed"
  exit 1
fi

HASH=""
TOTAL=0
PREFIXES=()
parse_blocklist_file "$PREFIX_FILE"
log "blocklist bytes=$(wc -c <"$PREFIX_FILE" | tr -d ' ') parsed=${#PREFIXES[@]} api_total=${TOTAL:-0}"

if [[ -z "${TOTAL// }" ]]; then
  TOTAL=${#PREFIXES[@]}
fi

PACKETS_DROPPED=0
PACKETS_ACCEPTED=0
KERNEL_METHOD=""
APPLIED_V4=0

count_ipv4_prefixes() {
  local n=0 p
  for p in "${PREFIXES[@]}"; do
    [[ "$p" == *:* ]] && continue
    n=$((n + 1))
  done
  APPLIED_V4=$n
}

nft_rule_packets() {
  local line=$1
  if [[ "$line" =~ counter[[:space:]]+packets[[:space:]]+([0-9]+) ]]; then
    echo "${BASH_REMATCH[1]}"
  else
    echo 0
  fi
}

ensure_nft_counters() {
  local table=inet name=evobgp_blocklist
  nft list chain "$table" "$name" input >/dev/null 2>&1 || return 0
  local drop_line
  drop_line=$(nft -a list chain "$table" "$name" input 2>/dev/null | grep 'ip saddr @v4' | grep drop | head -1 || true)
  if [[ -n "$drop_line" && "$drop_line" != *counter* ]]; then
    local handle
    handle=$(echo "$drop_line" | sed -n 's/.*# handle \([0-9]\+\).*/\1/p')
    if [[ -n "$handle" ]]; then
      nft delete rule "$table" "$name" input handle "$handle" 2>>"$LOG_FILE" || true
      drop_line=""
    fi
  fi
  if [[ -z "$drop_line" ]]; then
    nft add rule "$table" "$name" input ip saddr @v4 counter drop
  fi
  if ! nft list chain "$table" "$name" input 2>/dev/null | grep -qE '[[:space:]]counter[[:space:]]+accept'; then
    nft add rule "$table" "$name" input counter accept
  fi
}

collect_nft_packet_stats() {
  PACKETS_DROPPED=0
  PACKETS_ACCEPTED=0
  local line pkts
  while IFS= read -r line; do
    if [[ "$line" == *"ip saddr @v4"* && "$line" == *drop* ]]; then
      pkts=$(nft_rule_packets "$line")
      [[ -n "$pkts" ]] && PACKETS_DROPPED=$pkts
    elif [[ "$line" == *counter* && "$line" == *accept* && "$line" != *@v4* ]]; then
      pkts=$(nft_rule_packets "$line")
      [[ -n "$pkts" ]] && PACKETS_ACCEPTED=$pkts
    fi
  done < <(nft list chain inet evobgp_blocklist input 2>/dev/null || true)
}

collect_ipset_packet_stats() {
  PACKETS_DROPPED=0
  PACKETS_ACCEPTED=0
  local pkts
  pkts=$(iptables -L INPUT -v -n -x 2>/dev/null | awk '/match-set evobgp_blocklist_v4/ {print $1; exit}')
  [[ "$pkts" =~ ^[0-9]+$ ]] && PACKETS_DROPPED=$pkts
}

collect_packet_stats() {
  case "${KERNEL_METHOD:-$BACKEND}" in
    nft)
      ensure_nft_counters
      collect_nft_packet_stats
      ;;
    ipset)
      collect_ipset_packet_stats
      ;;
    iptables)
      PACKETS_DROPPED=$(iptables -L INPUT -v -n -x 2>/dev/null | awk '/DROP/ {s+=$1} END {print s+0}')
      PACKETS_ACCEPTED=0
      ;;
    *)
      if command -v nft >/dev/null 2>&1 && nft list chain inet evobgp_blocklist input >/dev/null 2>&1; then
        KERNEL_METHOD=nft
        ensure_nft_counters
        collect_nft_packet_stats
      elif iptables -L INPUT -v -n -x 2>/dev/null | grep -q 'evobgp_blocklist_v4'; then
        KERNEL_METHOD=ipset
        collect_ipset_packet_stats
      fi
      ;;
  esac
}

send_client_reports() {
  collect_packet_stats
  local km="${KERNEL_METHOD:-$BACKEND}"
  local report
  report=$(printf '{"status":"ok","prefix_count":%s,"ip_count":%s,"packets_dropped":%s,"packets_accepted":%s,"source":"cp","kernel_method":"%s"}' \
    "${TOTAL:-0}" "${APPLIED_V4:-0}" "${PACKETS_DROPPED:-0}" "${PACKETS_ACCEPTED:-0}" "$km")
  curl -fsS -X POST "${EVOBGP_CP_URL%/}/v1/firewall/apply-report" \
    -H "Authorization: Bearer ${CLIENT_TOKEN}" \
    -H "Content-Type: application/json" \
    -d "$report" >/dev/null 2>&1 || true
  curl -fsS -X POST "${EVOBGP_CP_URL%/}/v1/firewall/heartbeat" \
    -H "Authorization: Bearer ${CLIENT_TOKEN}" \
    -H "Content-Type: application/json" \
    -d '{"source":"cp"}' >/dev/null 2>&1 || true
}

if [[ -f "$HASH_FILE" && "$(tr -d '\r\n' <"$HASH_FILE")" == "$HASH" && -n "$HASH" ]]; then
  count_ipv4_prefixes
  log "unchanged hash $HASH — skip kernel apply (ipv4=${APPLIED_V4})"
  send_client_reports
  exit 0
fi

apply_nft() {
  local table=inet
  local name=evobgp_blocklist
  local v4=()
  local p
  for p in "${PREFIXES[@]}"; do
    [[ "$p" == *:* ]] && continue
    v4+=("$p")
  done

  nft list table "$table" "$name" >/dev/null 2>&1 || nft add table "$table" "$name"
  nft list set "$table" "$name" v4 >/dev/null 2>&1 || \
    nft add set "$table" "$name" v4 '{ type ipv4_addr; flags interval; }'
  nft flush set "$table" "$name" v4

  if ((${#v4[@]})); then
    local batch=()
    local chunk=64
    for p in "${v4[@]}"; do
      batch+=("$p")
      if ((${#batch[@]} >= chunk)); then
        nft_add_v4_chunk "$table" "$name" "${batch[@]}" || log "nft chunk add partial failure"
        batch=()
      fi
    done
    if ((${#batch[@]})); then
      nft_add_v4_chunk "$table" "$name" "${batch[@]}" || log "nft tail chunk add partial failure"
    fi
  fi

  nft list chain "$table" "$name" input >/dev/null 2>&1 || {
    nft add chain "$table" "$name" input '{ type filter hook input priority 0; policy accept; }'
    nft add rule "$table" "$name" input ip saddr @v4 counter drop
    nft add rule "$table" "$name" input counter accept
  }
  ensure_nft_counters
  KERNEL_METHOD=nft
  APPLIED_V4=${#v4[@]}
}

apply_ipset() {
  local set=evobgp_blocklist_v4
  local n=0
  ipset list "$set" >/dev/null 2>&1 || ipset create "$set" hash:net family inet hashsize 4096 maxelem 1048576
  ipset flush "$set"
  local p
  for p in "${PREFIXES[@]}"; do
    [[ "$p" == *:* ]] && continue
    ipset add "$set" "$p" -exist
    n=$((n + 1))
  done
  iptables -C INPUT -m set --match-set "$set" src -j DROP 2>/dev/null || \
    iptables -I INPUT -m set --match-set "$set" src -j DROP
  KERNEL_METHOD=ipset
  APPLIED_V4=$n
}

apply_iptables_only() {
  iptables -D INPUT -m comment --comment evobgp-block -j DROP 2>/dev/null || true
  local n=0
  local p
  for p in "${PREFIXES[@]}"; do
    [[ "$p" == *:* ]] && continue
    iptables -C INPUT -s "$p" -j DROP 2>/dev/null || iptables -A INPUT -s "$p" -j DROP
    n=$((n + 1))
  done
  KERNEL_METHOD=iptables
  APPLIED_V4=$n
}

clear_block() {
  case "$BACKEND" in
    nft) nft delete table inet evobgp_blocklist 2>/dev/null || true ;;
    ipset)
      ipset destroy evobgp_blocklist_v4 2>/dev/null || true
      iptables -D INPUT -m set --match-set evobgp_blocklist_v4 src -j DROP 2>/dev/null || true
      ;;
    iptables) iptables -S INPUT | grep -i evobgp | sed 's/^-A /-D /' | while read -r line; do iptables $line 2>/dev/null || true; done ;;
  esac
  APPLIED_V4=0
  KERNEL_METHOD="${BACKEND:-auto}"
}

APPLIED_V4=0
KERNEL_METHOD=""
if [[ "${TOTAL:-0}" == "0" || ${#PREFIXES[@]} -eq 0 ]]; then
  clear_block
  log "cleared blocklist (api total=${TOTAL:-0}) backend=$BACKEND"
else
  case "$BACKEND" in
    nft|auto) if command -v nft >/dev/null 2>&1; then apply_nft; else apply_ipset; fi ;;
    ipset) apply_ipset ;;
    iptables) apply_iptables_only ;;
    *) apply_ipset ;;
  esac
  log "applied api_total=${TOTAL} ipv4_in_kernel=${APPLIED_V4} from ${CP_HIT:-$EVOBGP_CP_URL} backend=$BACKEND hash=${HASH:-empty}"
fi

echo "$HASH" >"$HASH_FILE"
send_client_reports