#!/bin/bash
# XDP CDN Filter - IP 맵 업데이트
# BunnyNet + Cloudflare IP를 cdn_allow_v4/v6 BPF 맵에 로드
# Usage: xdp-cdn-update [--force]
set -euo pipefail

CDN_IPS_FILE="/opt/haproxy/conf/cdn-ips.lst"
STATE_FILE="/var/lib/xdp-defense/cdn-ips.hash"

RED='\033[0;31m'
GREEN='\033[0;32m'
CYAN='\033[0;36m'
NC='\033[0m'

log_ok()   { echo -e "${GREEN}[OK]${NC} $1"; }
log_err()  { echo -e "${RED}[ERROR]${NC} $1" >&2; }
log_info() { echo -e "${CYAN}[INFO]${NC} $1"; }

get_map_id() {
    bpftool map show -j 2>/dev/null | python3 -c "
import sys, json
maps = json.load(sys.stdin)
for m in maps:
    if m.get('name') == '$1':
        print(m['id'])
        break
" 2>/dev/null
}

# Fetch fresh CDN IPs if the file doesn't exist or --force
fetch_ips() {
    local tmp="${CDN_IPS_FILE}.tmp"
    echo "# CDN IP allowlist - updated $(date -u '+%Y-%m-%d %H:%M UTC')" > "$tmp"

    echo "# BunnyNet IPv4" >> "$tmp"
    curl -sf --max-time 10 "https://api.bunny.net/system/edgeserverlist/plain" | tr -d '\r' >> "$tmp" 2>/dev/null || true
    echo "" >> "$tmp"

    echo "# BunnyNet IPv6" >> "$tmp"
    curl -sf --max-time 10 "https://bunnycdn.com/api/system/edgeserverlist/IPv6" | tr -d '[]"\r' | tr ',' '\n' | sed 's/^ *//' >> "$tmp" 2>/dev/null || true
    echo "" >> "$tmp"

    echo "# Cloudflare IPv4" >> "$tmp"
    curl -sf --max-time 10 "https://www.cloudflare.com/ips-v4/" | tr -d '\r' >> "$tmp" 2>/dev/null || true
    echo "" >> "$tmp"

    echo "# Cloudflare IPv6" >> "$tmp"
    curl -sf --max-time 10 "https://www.cloudflare.com/ips-v6/" | tr -d '\r' >> "$tmp" 2>/dev/null || true
    echo "" >> "$tmp"

    local total
    total=$(grep -cvE '^(#|$)' "$tmp" || echo 0)
    if [ "$total" -lt 50 ]; then
        log_err "Only $total IPs fetched, keeping old file"
        rm -f "$tmp"
        return 1
    fi

    mv "$tmp" "$CDN_IPS_FILE"
    log_info "Fetched $total IPs"
}

# Load IPs into BPF maps
load_maps() {
    local v4_id v6_id
    v4_id=$(get_map_id "cdn_allow_v4")
    v6_id=$(get_map_id "cdn_allow_v6")

    if [ -z "$v4_id" ] || [ -z "$v6_id" ]; then
        log_err "CDN BPF maps not found. Is xdp_cdn_filter loaded?"
        return 1
    fi

    # Flush existing entries
    python3 -c "
import sys, json, subprocess
def flush_map(map_id):
    dump = subprocess.run(['bpftool','map','dump','id', str(map_id), '-j'],
                         capture_output=True, text=True)
    if dump.returncode != 0:
        return 0
    entries = json.loads(dump.stdout)
    count = 0
    for e in entries:
        key = e.get('key', [])
        if isinstance(key, list):
            hex_bytes = [k.replace('0x','') for k in key]
            cmd = ['bpftool','map','delete','id', str(map_id),'key','hex'] + hex_bytes
            if subprocess.run(cmd, capture_output=True).returncode == 0:
                count += 1
    return count
flush_map($v4_id)
flush_map($v6_id)
" 2>/dev/null || true

    # Load IPs from file and count via BPF map
    while IFS= read -r line; do
        # Skip comments and empty lines
        [[ "$line" =~ ^[[:space:]]*# ]] && continue
        [[ -z "${line// /}" ]] && continue
        line="${line%%#*}"    # strip inline comments
        line="${line// /}"    # strip whitespace

        if [[ "$line" == *":"* ]]; then
            # IPv6
            local addr="$line"
            local prefix=128
            if [[ "$line" == *"/"* ]]; then
                addr="${line%/*}"
                prefix="${line#*/}"
            fi
            local hex_key
            hex_key=$(python3 -c "
import socket, struct
addr = socket.inet_pton(socket.AF_INET6, '$addr')
prefix = $prefix
key = struct.pack('<I', prefix) + addr
print(' '.join(f'{b:02x}' for b in key))
" 2>/dev/null) || continue
            bpftool map update id "$v6_id" key hex $hex_key value hex 01 2>/dev/null || true
        else
            # IPv4
            local addr="$line"
            local prefix=32
            if [[ "$line" == *"/"* ]]; then
                addr="${line%/*}"
                prefix="${line#*/}"
            fi
            local hex_key
            hex_key=$(python3 -c "
import socket, struct
addr = socket.inet_aton('$addr')
prefix = $prefix
key = struct.pack('<I', prefix) + addr
print(' '.join(f'{b:02x}' for b in key))
" 2>/dev/null) || continue
            bpftool map update id "$v4_id" key hex $hex_key value hex 01 2>/dev/null || true
        fi
    done < "$CDN_IPS_FILE"

    # Verify actual counts from BPF maps
    local v4_count v6_count
    v4_count=$(bpftool map dump id "$v4_id" -j 2>/dev/null | python3 -c "import sys,json; print(len(json.load(sys.stdin)))" 2>/dev/null || echo 0)
    v6_count=$(bpftool map dump id "$v6_id" -j 2>/dev/null | python3 -c "import sys,json; print(len(json.load(sys.stdin)))" 2>/dev/null || echo 0)

    log_ok "Loaded: IPv4=$v4_count, IPv6=$v6_count"
}

# Stats
show_stats() {
    local stats_id
    stats_id=$(get_map_id "cdn_stats")
    if [ -z "$stats_id" ]; then
        log_err "cdn_stats map not found"
        return 1
    fi
    echo "=== XDP CDN Filter Stats ==="
    bpftool map dump id "$stats_id" -j 2>/dev/null | python3 -c "
import sys, json
labels = ['CDN pass', 'Whitelist pass', 'Non-web pass', 'Dropped']
try:
    entries = json.load(sys.stdin)
    for e in entries:
        # Use formatted section if available (bpftool with BTF)
        fmt = e.get('formatted', {})
        if fmt:
            idx = int(fmt['key'])
            total = sum(v['value'] for v in fmt['values'])
        else:
            key = e.get('key', 0)
            if isinstance(key, list):
                idx = int(key[0], 16) if isinstance(key[0], str) else int(key[0])
            else:
                idx = int(key)
            vals = e.get('values', e.get('value', []))
            total = 0
            if isinstance(vals, list):
                for v in vals:
                    if isinstance(v, dict):
                        total += int(v.get('value', v.get('val', 0)))
                    else:
                        total += int(v)
            else:
                total = int(vals)
        if idx < len(labels):
            print(f'  {labels[idx]:<20} {total:>12}')
except Exception as ex:
    print(f'  Parse error: {ex}', file=sys.stderr)
"
}

main() {
    case "${1:-update}" in
        update|--force)
            if [ ! -f "$CDN_IPS_FILE" ] || [ "${1:-}" = "--force" ]; then
                fetch_ips
            fi
            # Check if hash changed
            local new_hash
            new_hash=$(md5sum "$CDN_IPS_FILE" | cut -d' ' -f1)
            local old_hash=""
            [ -f "$STATE_FILE" ] && old_hash=$(cat "$STATE_FILE")
            if [ "$new_hash" = "$old_hash" ] && [ "${1:-}" != "--force" ]; then
                log_info "No changes (hash: $new_hash)"
                return 0
            fi
            load_maps
            echo "$new_hash" > "$STATE_FILE"
            ;;
        fetch)
            fetch_ips
            ;;
        stats)
            show_stats
            ;;
        *)
            echo "Usage: $0 [update|fetch|stats|--force]"
            ;;
    esac
}

main "$@"