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

URL=""
INTERVAL="1h"
SET="abuseipdb"
WSET="abuseipdb_whitelist"
WHITELIST_CIDRS=""
URL_ONLY="0"
NO_UPDATER="0"

usage() {
  cat <<EOF
Usage: $0 --url <blocklist_url> [options]

Required:
  --url <url>                 URL to IPv4 blocklist

Options:
  --interval <systemd_time>   e.g. 15min, 1h, 6h, 1d (default: 1h)
  --set <name>                ipset name for blacklist (default: abuseipdb)
  --whitelist <cidr,...>      comma-separated CIDRs (only in full install)
  --url-only                  update ONLY systemd service/timer + ensure updater supports UA
  --no-updater                (with --url-only) do not touch updater script
EOF
}

while [ $# -gt 0 ]; do
  case "$1" in
    --url) URL="${2:-}"; shift 2;;
    --interval) INTERVAL="${2:-}"; shift 2;;
    --set) SET="${2:-}"; shift 2;;
    --whitelist) WHITELIST_CIDRS="${2:-}"; shift 2;;
    --url-only) URL_ONLY="1"; shift 1;;
    --no-updater) NO_UPDATER="1"; shift 1;;
    -h|--help) usage; exit 0;;
    *) echo "Unknown arg: $1"; usage; exit 2;;
  esac
done

[ -n "$URL" ] || { echo "ERROR: --url is required"; usage; exit 2; }
[ "$(id -u)" -eq 0 ] || { echo "ERROR: run as root (use sudo)"; exit 1; }

backup_if_exists() {
  local f="$1"
  if [ -f "$f" ]; then
    cp -a "$f" "$f.bak.$(date +%F_%H%M%S)"
  fi
}

write_updater() {
  backup_if_exists /usr/local/sbin/abuseipdb_ipset_update.sh

  cat >/usr/local/sbin/abuseipdb_ipset_update.sh <<'EOF'
#!/usr/bin/env bash
set -euo pipefail

URL="${1:-}"
SET="${2:-abuseipdb}"
[ -n "$URL" ] || { echo "Usage: $0 <url_to_list> [setname]"; exit 2; }

NEWSET="${SET}_new"

TMPDIR="$(mktemp -d)"
trap 'rm -rf "$TMPDIR"' EXIT

LIST_TXT="$TMPDIR/list.txt"
RESTORE_TXT="$TMPDIR/ipset.restore"

# Use browser-like UA because some servers block curl/wget by User-Agent
curl -fsSL -A 'Mozilla/5.0+(compatible; BlackListLNS/1.0; https://lns.bz/)' "$URL" -o "$LIST_TXT"

if [ ! -s "$LIST_TXT" ]; then
  echo "ERROR: downloaded list is empty, refusing to update."
  exit 1
fi

{
  echo "create $NEWSET hash:ip family inet hashsize 131072 maxelem 200000 -exist"
  awk '
    $0 ~ /^[[:space:]]*#/ {next}
    NF==0 {next}
    {
      ip=$1
      if (ip ~ /^([0-9]{1,3}\.){3}[0-9]{1,3}$/) {
        print "add '"$NEWSET"' " ip
      }
    }
  ' "$LIST_TXT"
} > "$RESTORE_TXT"

ipset restore -exist < "$RESTORE_TXT"

ipset create "$SET" hash:ip family inet hashsize 131072 maxelem 200000 -exist
ipset swap "$SET" "$NEWSET"
ipset destroy "$NEWSET"

echo "Updated ipset '$SET': $(ipset list "$SET" 2>/dev/null | awk -F': ' '/Number of entries/ {print $2}') entries"
EOF

  chmod +x /usr/local/sbin/abuseipdb_ipset_update.sh
}

ensure_updater_supports_ua() {
  if [ "$NO_UPDATER" = "1" ]; then
    return 0
  fi

  if [ ! -f /usr/local/sbin/abuseipdb_ipset_update.sh ]; then
    echo "Updater not found, installing..."
    write_updater
    return 0
  fi

  if grep -q "curl -fsSL .* -A 'Mozilla/5.0+(compatible; BlackListLNS/1.0; https://lns.bz/)" /usr/local/sbin/abuseipdb_ipset_update.sh; then
    return 0
  fi

  echo "Updater exists but without UA. Updating updater script (backup will be created)..."
  write_updater
}

install_systemd() {
  echo "Installing/updating systemd unit + timer..."

  backup_if_exists /etc/systemd/system/abuseipdb-ipset-update.service
  backup_if_exists /etc/systemd/system/abuseipdb-ipset-update.timer

  cat >/etc/systemd/system/abuseipdb-ipset-update.service <<EOF
[Unit]
Description=Update AbuseIPDB ipset blocklist
Wants=network-online.target
After=network-online.target

[Service]
Type=oneshot
ExecStart=/usr/local/sbin/abuseipdb_ipset_update.sh $URL $SET
EOF

  cat >/etc/systemd/system/abuseipdb-ipset-update.timer <<EOF
[Unit]
Description=Run AbuseIPDB ipset update periodically

[Timer]
OnBootSec=30sec
OnUnitActiveSec=$INTERVAL
Persistent=true

[Install]
WantedBy=timers.target
EOF

  systemctl daemon-reload
  systemctl start abuseipdb-ipset-update.service
  systemctl enable --now abuseipdb-ipset-update.timer
}

if [ "$URL_ONLY" = "1" ]; then
  ensure_updater_supports_ua
  install_systemd
  echo
  echo "URL-only update done."
  echo "Check: journalctl -u abuseipdb-ipset-update.service -n 50 --no-pager"
  exit 0
fi

# Full install (unchanged in spirit)
export DEBIAN_FRONTEND=noninteractive
apt-get update -y
apt-get install -y ipset ufw curl gawk

ipset create "$SET" hash:ip family inet hashsize 131072 maxelem 200000 -exist
if [ -n "$WHITELIST_CIDRS" ]; then
  ipset create "$WSET" hash:net family inet maxelem 65536 -exist
  IFS=',' read -ra _CIDRS <<<"$WHITELIST_CIDRS"
  for cidr in "${_CIDRS[@]}"; do
    cidr="$(echo "$cidr" | xargs)"
    [ -n "$cidr" ] || continue
    ipset add "$WSET" "$cidr" -exist
  done
fi

BEFORE="/etc/ufw/before.rules"
BACKUP="/etc/ufw/before.rules.bak.$(date +%F_%H%M%S)"
cp -a "$BEFORE" "$BACKUP"

ensure_rule_before_commit() {
  local rule="$1"
  grep -F -- "$rule" "$BEFORE" >/dev/null 2>&1 && return 0
  awk -v rule="$rule" '
    BEGIN{added=0}
    {
      if ($0=="COMMIT" && added==0) { print rule; added=1 }
      print $0
    }
    END{ if (added==0) { print rule } }
  ' "$BEFORE" > "${BEFORE}.tmp"
  mv "${BEFORE}.tmp" "$BEFORE"
}

if [ -n "$WHITELIST_CIDRS" ]; then
  ensure_rule_before_commit "# AbuseIPDB ipset whitelist (do not block)"
  ensure_rule_before_commit "-A ufw-before-input -m set --match-set $WSET src -j ACCEPT"
fi
ensure_rule_before_commit "# AbuseIPDB ipset blacklist"
ensure_rule_before_commit "-A ufw-before-input -m set --match-set $SET src -j DROP"

ufw allow 22/tcp >/dev/null || true
ufw allow 80/tcp >/dev/null || true
ufw allow 443/tcp >/dev/null || true
ufw --force enable
ufw reload

write_updater
install_systemd

echo "Done."