#!/usr/bin/env bash
# =============================================================================
# onx-firewall-geoip-block — MaxMind GeoLite2 + nftables/ipset country block.
#
# Input (stdin JSON):
#   {
#     "country_code": "TR",   // ISO 3166-1 alpha-2 (case-insensitive)
#     "action":       "block" // block | unblock | list
#   }
#
# Output (stdout JSON):
#   {
#     "country":    "TR",
#     "action":     "block",
#     "ipset":      "onox-block-tr",
#     "ip_count":   2_345,
#     "applied":    true
#   }
#
# When action=list:
#   { "blocked": ["TR","CN","RU"], "count": 3 }
#
# Exit codes:
#   0  — success
#   1  — invalid input
#   2  — preflight fail (GeoLite2 missing, ipset missing, etc.)
#   3  — execution fail
#
# Production notes:
#   • Phase 2 dependency: /var/lib/onox/geoip/GeoLite2-Country.mmdb provided by
#     geoipupdate daemon (MaxMind license key in /etc/GeoIP.conf). Without it,
#     this script returns code 2 ("geolite2_not_installed") and the panel UI
#     should show a "configure GeoIP" message.
#   • Backend preference: nftables (preferred) → iptables+ipset (legacy).
#
# Sudoers:
#   apache ALL=(root) NOPASSWD: /usr/local/onoxsoft/bin/onx-firewall-geoip-block
#
# Deployed to: /usr/local/onoxsoft/bin/onx-firewall-geoip-block
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
# shellcheck source=_lib/common.sh
source "${SCRIPT_DIR}/_lib/common.sh"

onx_json_input

ACTION=$(onx_json_field 'action' 'block')
CODE_RAW=$(onx_json_field 'country_code' '')

GEOIP_DIR="${ONX_GEOIP_DIR:-/var/lib/onox/geoip}"
GEOIP_DB="${GEOIP_DIR}/GeoLite2-Country.mmdb"

case "$ACTION" in
    block|unblock|list) : ;;
    *) onx_die 1 "invalid action (must be block|unblock|list)" ;;
esac

# ── action=list → enumerate active onox-block-* ipsets ───────────────────────
if [[ "$ACTION" == "list" ]]; then
    if command -v ipset >/dev/null 2>&1; then
        # ipset names of the form onox-block-<lc>
        blocked=$(ipset list -name 2>/dev/null | grep -E '^onox-block-[a-z]{2}$' | sed 's/^onox-block-//' | tr '[:lower:]' '[:upper:]' || true)
    elif command -v nft >/dev/null 2>&1; then
        # nftables sets in family inet table named onox_block_<lc>
        blocked=$(nft list sets 2>/dev/null | grep -oE 'set onox_block_[a-z]{2}' | sed 's/^set onox_block_//' | tr '[:lower:]' '[:upper:]' | sort -u || true)
    else
        blocked=""
    fi

    if [[ -z "$blocked" ]]; then
        printf '{"blocked":[],"count":0}\n'
    else
        echo "$blocked" | jq -R . | jq -s '{blocked: .,count:length}'
    fi
    exit 0
fi

# block/unblock: country_code required
[[ "$CODE_RAW" =~ ^[A-Za-z]{2}$ ]] || onx_die 1 "country_code must be ISO 3166-1 alpha-2 (e.g. TR, CN)"
CODE_UC=$(printf '%s' "$CODE_RAW" | tr '[:lower:]' '[:upper:]')
CODE_LC=$(printf '%s' "$CODE_RAW" | tr '[:upper:]' '[:lower:]')
IPSET_NAME="onox-block-${CODE_LC}"

# ── Backend pick ─────────────────────────────────────────────────────────────
BACKEND=""
if command -v nft >/dev/null 2>&1; then
    BACKEND="nftables"
elif command -v ipset >/dev/null 2>&1 && command -v iptables >/dev/null 2>&1; then
    BACKEND="iptables-ipset"
else
    onx_die 2 "no supported firewall backend (need nft or ipset+iptables)"
fi

# ── GeoLite2 DB presence (block only — unblock can run without it) ───────────
if [[ "$ACTION" == "block" ]]; then
    if [[ ! -f "$GEOIP_DB" ]]; then
        onx_die 2 "geolite2_not_installed: $GEOIP_DB (run geoipupdate)"
    fi
fi

# ── Build country IP list (block only) ───────────────────────────────────────
extract_country_cidrs() {
    local cc="$1"
    # Preferred: mmdblookup over a precomputed CIDR list. In production a
    # geoipupdate-shipped converter (e.g. python-mmdb-tools) writes per-country
    # CIDR files to ${GEOIP_DIR}/cidrs/<cc>.txt. Use it if present.
    local cidr_file="${GEOIP_DIR}/cidrs/${cc}.txt"
    if [[ -r "$cidr_file" ]]; then
        cat "$cidr_file"
        return
    fi
    # Fallback: empty (caller will exit 3). Production deploy provides
    # /etc/onoxsoft/geoipupdate.conf + cron daily refresh.
    return 1
}

if [[ "$ACTION" == "block" ]]; then
    require_cmd jq
    CIDR_LIST=$(extract_country_cidrs "$CODE_UC" 2>/dev/null || true)
    if [[ -z "$CIDR_LIST" ]]; then
        onx_die 2 "geoip_cidrs_unavailable for $CODE_UC — provision ${GEOIP_DIR}/cidrs/${CODE_UC}.txt via geoipupdate"
    fi
    IP_COUNT=$(printf '%s\n' "$CIDR_LIST" | grep -cv '^$' || echo 0)
fi

# ── nftables backend ─────────────────────────────────────────────────────────
if [[ "$BACKEND" == "nftables" ]]; then
    SET_NAME="onox_block_${CODE_LC}"
    TABLE="inet onox"

    nft list table inet onox >/dev/null 2>&1 || nft add table inet onox

    if [[ "$ACTION" == "block" ]]; then
        # Clean set if already there, then recreate fresh
        nft delete set $TABLE "$SET_NAME" 2>/dev/null || true
        nft add set $TABLE "$SET_NAME" '{ type ipv4_addr; flags interval; }'

        # Bulk add CIDRs via nft -f stdin (fast)
        {
            printf 'add element %s %s { ' "$TABLE" "$SET_NAME"
            sep=""
            while IFS= read -r cidr; do
                [[ -z "$cidr" ]] && continue
                printf '%s%s' "$sep" "$cidr"
                sep=", "
            done <<<"$CIDR_LIST"
            printf ' }\n'
        } | nft -f - || onx_die 3 "nft_add_elements_failed"

        # Drop rule on INPUT
        nft list chain inet onox INPUT >/dev/null 2>&1 || nft 'add chain inet onox INPUT { type filter hook input priority 0 ; }'
        nft list chain inet onox INPUT 2>/dev/null | grep -q "@$SET_NAME" || \
            nft add rule inet onox INPUT ip saddr "@$SET_NAME" drop comment "\"onox geoip block $CODE_UC\""

        APPLIED=true
    else
        nft list table inet onox >/dev/null 2>&1 && {
            # Remove rule then set
            handle=$(nft -a list chain inet onox INPUT 2>/dev/null | awk -v s="@$SET_NAME" '$0 ~ s {for (i=1;i<=NF;i++) if ($i=="handle") print $(i+1)}' | head -1)
            [[ -n "$handle" ]] && nft delete rule inet onox INPUT handle "$handle" 2>/dev/null || true
            nft delete set inet onox "$SET_NAME" 2>/dev/null || true
        }
        APPLIED=true
        IP_COUNT=0
    fi
fi

# ── iptables+ipset backend ───────────────────────────────────────────────────
if [[ "$BACKEND" == "iptables-ipset" ]]; then
    if [[ "$ACTION" == "block" ]]; then
        ipset destroy "$IPSET_NAME" 2>/dev/null || true
        ipset create "$IPSET_NAME" hash:net family inet maxelem 131072
        while IFS= read -r cidr; do
            [[ -z "$cidr" ]] && continue
            ipset add "$IPSET_NAME" "$cidr" 2>/dev/null || true
        done <<<"$CIDR_LIST"

        # Idempotent insert at head of INPUT
        iptables -C INPUT -m set --match-set "$IPSET_NAME" src -j DROP 2>/dev/null || \
            iptables -I INPUT -m set --match-set "$IPSET_NAME" src -j DROP

        APPLIED=true
    else
        iptables -D INPUT -m set --match-set "$IPSET_NAME" src -j DROP 2>/dev/null || true
        ipset destroy "$IPSET_NAME" 2>/dev/null || true
        APPLIED=true
        IP_COUNT=0
    fi
fi

onx_audit "onx-firewall-geoip" "country=$CODE_UC action=$ACTION backend=$BACKEND applied=${APPLIED:-false} ip_count=${IP_COUNT:-0}"

jq -nc \
    --arg country "$CODE_UC" \
    --arg action "$ACTION" \
    --arg ipset "$IPSET_NAME" \
    --arg backend "$BACKEND" \
    --argjson applied "${APPLIED:-false}" \
    --argjson ip_count "${IP_COUNT:-0}" \
    '{country:$country, action:$action, ipset:$ipset, backend:$backend, applied:$applied, ip_count:$ip_count}'

exit 0
