#!/usr/bin/env bash
set -Eeuo pipefail
umask 077

readonly PROGRAM_NAME="${0##*/}"
readonly DEFAULT_CONFIG_FILE="/etc/abuseipdb-blocklist.conf"
readonly SET_V4="abuseipdb_ipv4" SET_V6="abuseipdb_ipv6"
readonly TEMP_SET_V4="abuseipdb_ipv4_new" TEMP_SET_V6="abuseipdb_ipv6_new"
CONFIG_FILE="$DEFAULT_CONFIG_FILE"
MODE="dry-run"
WORK_DIR=""

log() { printf '%s %s\n' "$(date --iso-8601=seconds)" "$*"; }
die() { log "ERROR: $*" >&2; exit 1; }
usage() {
    cat <<EOF
Usage: $PROGRAM_NAME [--config FILE] [--dry-run | --ensure-rules | --apply]
  --dry-run      Download and validate only; make no firewall changes (default)
  --ensure-rules Create owned empty sets and idempotent iptables rules; no API call
  --apply        Download, validate, ensure rules, and atomically swap both sets
  --config FILE  Read configuration from FILE
EOF
}
cleanup() {
    local exit_code=$?
    [[ -z "$WORK_DIR" || ! -d "$WORK_DIR" ]] || rm -rf -- "$WORK_DIR"
    exit "$exit_code"
}
trap cleanup EXIT
trap 'die "Update interrupted"' INT TERM

while (($#)); do
    case "$1" in
        --dry-run) MODE="dry-run" ;;
        --ensure-rules) MODE="ensure-rules" ;;
        --apply) MODE="apply" ;;
        --config) shift; (($#)) || die "--config requires a file"; CONFIG_FILE="$1" ;;
        --help|-h) usage; exit 0 ;;
        *) die "Unknown argument: $1" ;;
    esac
    shift
done

[[ $EUID -eq 0 ]] || die "This script must run as root"
[[ -f "$CONFIG_FILE" ]] || die "Configuration file not found: $CONFIG_FILE"
config_owner="$(stat --format='%u' "$CONFIG_FILE")"
[[ $config_owner == 0 ]] || die "Configuration must be owned by root: $CONFIG_FILE"
config_mode="$(stat --format='%a' "$CONFIG_FILE")"
(( (8#$config_mode & 8#022) == 0 )) || die "Configuration must not be group/other-writable: $CONFIG_FILE"
# shellcheck disable=SC1090
source "$CONFIG_FILE"

readonly CONFIDENCE_MINIMUM="${CONFIDENCE_MINIMUM:-100}"
readonly LIST_LIMIT="${LIST_LIMIT:-10000}"
readonly MINIMUM_IPV4_COUNT="${MINIMUM_IPV4_COUNT:-1000}"
readonly MINIMUM_IPV6_COUNT="${MINIMUM_IPV6_COUNT:-1}"
readonly CURL_CONNECT_TIMEOUT="${CURL_CONNECT_TIMEOUT:-10}"
readonly CURL_MAX_TIME="${CURL_MAX_TIME:-60}"
readonly LOCK_FILE="${LOCK_FILE:-/run/lock/abuseipdb-blocklist.lock}"

[[ $CONFIDENCE_MINIMUM =~ ^[0-9]+$ ]] || die "CONFIDENCE_MINIMUM must be an integer"
((10#$CONFIDENCE_MINIMUM >= 25 && 10#$CONFIDENCE_MINIMUM <= 100)) || die "CONFIDENCE_MINIMUM must be from 25 to 100"
[[ $LIST_LIMIT =~ ^[1-9][0-9]*$ ]] || die "LIST_LIMIT must be a positive integer"
[[ $MINIMUM_IPV4_COUNT =~ ^[0-9]+$ ]] || die "MINIMUM_IPV4_COUNT must be a non-negative integer"
[[ $MINIMUM_IPV6_COUNT =~ ^[0-9]+$ ]] || die "MINIMUM_IPV6_COUNT must be a non-negative integer"

for command_name in awk date flock ip6tables ipset iptables mkdir mktemp rm stat; do
    command -v "$command_name" >/dev/null || die "Required command not found: $command_name"
done
if [[ $MODE != ensure-rules ]]; then
    : "${ABUSEIPDB_API_KEY:?ABUSEIPDB_API_KEY is required in $CONFIG_FILE}"
    [[ $ABUSEIPDB_API_KEY =~ ^[A-Za-z0-9]{40,128}$ ]] || die "ABUSEIPDB_API_KEY has an invalid format"
    readonly ABUSEIPDB_API_KEY
    for command_name in curl python3 wc; do
        command -v "$command_name" >/dev/null || die "Required command not found: $command_name"
    done
fi

mkdir -p -- "${LOCK_FILE%/*}"
exec 9>"$LOCK_FILE"
flock -n 9 || die "Another update is already running"

ensure_set() {
    local set_name="$1" family="$2"
    ipset create "$set_name" hash:ip family "$family" hashsize 16384 maxelem 524288 -exist
}

ensure_rule() {
    local tool="$1" chain="$2" position="$3" set_name="$4"
    local current_position
    if "$tool" --wait 5 --check "$chain" -m set --match-set "$set_name" src -m comment \
        --comment 'AbuseIPDB managed blocklist' -j DROP 2>/dev/null; then
        current_position="$(LC_ALL=C "$tool" --wait 5 -S "$chain" | awk -v set_name="$set_name" '
            $1 == "-A" { rule_number++ }
            index($0, "--match-set " set_name " src") { print rule_number; exit }
        ')"
        [[ $current_position == "$position" ]] && return
        "$tool" --wait 5 --delete "$chain" -m set --match-set "$set_name" src -m comment \
            --comment 'AbuseIPDB managed blocklist' -j DROP
    fi
    "$tool" --wait 5 --insert "$chain" "$position" -m set --match-set "$set_name" src \
        -m comment --comment 'AbuseIPDB managed blocklist' -j DROP
}

ensure_input_rule() {
    local tool="$1" set_name="$2" anchor_position current_position
    anchor_position="$(LC_ALL=C "$tool" --wait 5 -S INPUT | awk '
        $1 == "-A" { rule_number++ }
        $0 == "-A INPUT -j CROWDSEC_CHAIN" { print rule_number; exit }
    ')"
    [[ $anchor_position =~ ^[0-9]+$ ]] || die "$tool INPUT has no CROWDSEC_CHAIN anchor"

    if "$tool" --wait 5 --check INPUT -m set --match-set "$set_name" src -m comment \
        --comment 'AbuseIPDB managed blocklist' -j DROP 2>/dev/null; then
        current_position="$(LC_ALL=C "$tool" --wait 5 -S INPUT | awk -v set_name="$set_name" '
            $1 == "-A" { rule_number++ }
            index($0, "--match-set " set_name " src") { print rule_number; exit }
        ')"
        [[ $current_position == "$((anchor_position + 1))" ]] && return
        "$tool" --wait 5 --delete INPUT -m set --match-set "$set_name" src -m comment \
            --comment 'AbuseIPDB managed blocklist' -j DROP
        anchor_position="$(LC_ALL=C "$tool" --wait 5 -S INPUT | awk '
            $1 == "-A" { rule_number++ }
            $0 == "-A INPUT -j CROWDSEC_CHAIN" { print rule_number; exit }
        ')"
    fi

    "$tool" --wait 5 --insert INPUT "$((anchor_position + 1))" -m set \
        --match-set "$set_name" src -m comment --comment 'AbuseIPDB managed blocklist' -j DROP
}

ensure_firewall_objects() {
    ensure_set "$SET_V4" inet
    ensure_set "$SET_V6" inet6
    # Keep reputation blocking immediately after CrowdSec even when Fail2ban has
    # dynamically inserted an SSH-only jump before the CrowdSec anchor.
    ensure_input_rule iptables "$SET_V4"
    ensure_input_rule ip6tables "$SET_V6"
    # Docker evaluates administrator policy in DOCKER-USER before its accepts.
    ensure_rule iptables DOCKER-USER 1 "$SET_V4"
    ensure_rule ip6tables DOCKER-USER 1 "$SET_V6"
}

if [[ $MODE == ensure-rules ]]; then
    ensure_firewall_objects
    log "AbuseIPDB ipsets and iptables/ip6tables rules are present"
    exit 0
fi

WORK_DIR="$(mktemp -d /tmp/abuseipdb-blocklist.XXXXXX)"
download_lists() {
    local output_v4="$1" output_v6="$2" header_file="$3"
    log "Downloading mixed AbuseIPDB IPv4/IPv6 blacklist"
    # Stream curl configuration so the key exists only in shell memory and a
    # pipe: it is never present in curl's argv, environment, or a temporary file.
    printf 'header = "Key: %s"\nheader = "Accept: text/plain"\n' "$ABUSEIPDB_API_KEY" | \
    curl --config - --fail-with-body --silent --show-error --location \
        --retry 1 --retry-delay 2 \
        --connect-timeout "$CURL_CONNECT_TIMEOUT" --max-time "$CURL_MAX_TIME" \
        --get 'https://api.abuseipdb.com/api/v2/blacklist' \
        --data-urlencode "confidenceMinimum=$CONFIDENCE_MINIMUM" \
        --data-urlencode "limit=$LIST_LIMIT" \
        --dump-header "$header_file" --output "$WORK_DIR/blacklist.raw"
    python3 - "$WORK_DIR/blacklist.raw" "$output_v4" "$output_v6" <<'PY'
import ipaddress
import pathlib
import sys

source = pathlib.Path(sys.argv[1])
targets = {4: pathlib.Path(sys.argv[2]), 6: pathlib.Path(sys.argv[3])}
addresses = {4: set(), 6: set()}
for line_number, raw_line in enumerate(source.read_text(encoding="ascii").splitlines(), 1):
    value = raw_line.strip()
    if not value:
        continue
    try:
        address = ipaddress.ip_address(value)
    except ValueError as error:
        raise SystemExit(f"invalid address on line {line_number}: {error}")
    if not address.is_global:
        raise SystemExit(f"non-global address on line {line_number}: {address}")
    addresses[address.version].add(address)
for version, target in targets.items():
    target.write_text(
        "".join(f"{address}\n" for address in sorted(addresses[version])),
        encoding="ascii",
    )
PY
}

v4_file="$WORK_DIR/ipv4.txt"
v6_file="$WORK_DIR/ipv6.txt"
download_lists "$v4_file" "$v6_file" "$WORK_DIR/blacklist.headers"
v4_count="$(wc -l <"$v4_file")"
v6_count="$(wc -l <"$v6_file")"
total_count="$((v4_count + v6_count))"
((v4_count >= MINIMUM_IPV4_COUNT)) || die "IPv4 list has only $v4_count entries; refusing update"
((v6_count >= MINIMUM_IPV6_COUNT)) || die "IPv6 list has only $v6_count entries; refusing update"
((total_count <= LIST_LIMIT)) || die "Combined list exceeds requested limit ($total_count > $LIST_LIMIT)"
log "Candidate validated: IPv4=$v4_count IPv6=$v6_count total=$total_count confidence=$CONFIDENCE_MINIMUM limit=$LIST_LIMIT"

if [[ $MODE == dry-run ]]; then
    log "Dry run complete; live firewall was not changed"
    exit 0
fi

ensure_firewall_objects
ipset destroy "$TEMP_SET_V4" 2>/dev/null || true
ipset destroy "$TEMP_SET_V6" 2>/dev/null || true
ensure_set "$TEMP_SET_V4" inet
ensure_set "$TEMP_SET_V6" inet6
awk -v set_name="$TEMP_SET_V4" '{ print "add " set_name " " $0 }' "$v4_file" | ipset restore
awk -v set_name="$TEMP_SET_V6" '{ print "add " set_name " " $0 }' "$v6_file" | ipset restore

# Each family swaps atomically; both the old and new generation remain valid if
# the second family unexpectedly fails.
ipset swap "$SET_V4" "$TEMP_SET_V4"
ipset swap "$SET_V6" "$TEMP_SET_V6"
ipset destroy "$TEMP_SET_V4"
ipset destroy "$TEMP_SET_V6"
log "AbuseIPDB sets updated: IPv4=$v4_count IPv6=$v6_count"