#!/bin/bash
################################################################################
# Script Name: ufw-blocklists.sh
# Version: 1.23
# Description: Per-feed UFW threat intelligence blocking with ipset
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
################################################################################
# Changelog:
#   1.23 - Fix ipset restore: write batch to temp file instead of pipe to avoid
#          silent failures; log errors instead of suppressing; clean up temp
#          ipsets on failure
#        - Fix IPv6 parser: reject timestamps (e.g. 14:34:21) misidentified
#          as IPv6 addresses by requiring hex letters, 3+ groups, or ::
#   1.22 - Flush temp ipsets before loading to avoid stale entries after crashes
#        - Always load ipsets even when feed is empty (clears stale blocks)
#        - show-stats: cache journalctl once, count both v4 and v6 blocks,
#          use Members: header for ipset counting
#        - add-feed: validate that URL is provided
################################################################################
# Don't use 'set -e' - it causes silent failures when log file has permission issues

CONFIG_DIR="/etc/ufw-threats"
FEEDS_CONFIG="$CONFIG_DIR/feeds.conf"
CACHE_DIR="$CONFIG_DIR/cache"
LOG_FILE="/var/log/ufw-threats.log"
SSH_PORT="22"
ENABLE_AUTO_UPDATE=true
UPDATE_INTERVAL="daily"
ENABLE_IPV6=true
UFW_RULES_FILE="/etc/ufw/before.rules"
UFW_RULES_V6_FILE="/etc/ufw/before6.rules"
IPSET_PREFIX="ufw-feed"
WHITELIST_IPSET="ufw-whitelist"
WHITELIST_IPSET_V6="ufw-whitelist-v6"
MAX_BACKUPS=10


show_usage() {
    cat <<EOF
Usage: $0 [OPTIONS] [COMMAND]

PER-FEED VERSION: Each threat feed gets its own ipset and iptables rule.
Provides detailed per-feed blocking statistics and metrics.

COMMANDS:
    install                 Install and configure threat feed blocking
    update                  Update all enabled feeds now (ipsets only, no UFW reload)
    apply-rules             Regenerate and apply UFW rules (use with caution!)
    test-rules              Test rule generation without applying
    add-feed NAME URL       Add a custom feed
    remove-feed NAME        Remove a feed
    enable-feed NAME        Enable a disabled feed
    disable-feed NAME       Disable a feed
    list-feeds              List all configured feeds
    show-stats              Show blocking statistics per feed
    whitelist-add IP        Add IP/CIDR to whitelist
    whitelist-init          Initialize whitelist with RFC1918/Docker networks
    whitelist-list          Show all whitelisted IPs
    clean-cache             Remove cache files for disabled feeds

OPTIONS:
    -h, --help              Show this help message
    -s, --ssh-port PORT     SSH port (default: 22)
    --no-auto-update        Disable automatic daily updates
    --no-ipv6               Disable IPv6 support
    --update-interval TIME  Update interval: hourly, daily, weekly (default: daily)

EXAMPLES:
    sudo $0 install
    sudo $0 update          # Safe - only updates ipsets
    sudo $0 test-rules      # Safe - validates without applying
    sudo $0 apply-rules     # DANGER - regenerates UFW config
    sudo $0 show-stats

EOF
    exit 0
}

log_message() {
    local msg
    msg="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    echo "$msg"
    echo "$msg" >> "$LOG_FILE" 2>/dev/null || true
}

# Iterate over enabled feeds in $FEEDS_CONFIG, calling the provided callback
# function with arguments: name url type description
# Usage: for_each_enabled_feed my_callback_function
for_each_enabled_feed() {
    local callback="$1"
    [ -f "$FEEDS_CONFIG" ] || return 0

    local enabled name url type description
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue
        "$callback" "$name" "$url" "$type" "$description"
    done < "$FEEDS_CONFIG"
}

# Iterate over ALL feeds (enabled + disabled), calling the provided callback
# function with arguments: enabled name url type description
for_each_feed() {
    local callback="$1"
    [ -f "$FEEDS_CONFIG" ] || return 0

    local enabled name url type description
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        "$callback" "$enabled" "$name" "$url" "$type" "$description"
    done < "$FEEDS_CONFIG"
}

parse_args() {
    COMMAND=""
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            -s|--ssh-port) SSH_PORT="$2"; shift 2 ;;
            --no-auto-update) ENABLE_AUTO_UPDATE=false; shift ;;
            --no-ipv6) ENABLE_IPV6=false; shift ;;
            --update-interval) UPDATE_INTERVAL="$2"; shift 2 ;;
            install|update|apply-rules|test-rules|list-feeds|show-stats|whitelist-init|whitelist-list|clean-cache) COMMAND="$1"; shift ;;
            add-feed) COMMAND="add-feed"; FEED_NAME="$2"; FEED_URL="$3"; shift 3 ;;
            remove-feed|enable-feed|disable-feed) COMMAND="$1"; FEED_NAME="$2"; shift 2 ;;
            whitelist-add) COMMAND="whitelist-add"; WHITELIST_IP="$2"; shift 2 ;;
            *) echo "Unknown option: $1"; exit 1 ;;
        esac
    done
    [ -z "$COMMAND" ] && COMMAND="install"
}

cleanup_old_backups() {
    local max_keep=${MAX_BACKUPS:-10}

    find "$(dirname "$UFW_RULES_FILE")" -maxdepth 1 -name "$(basename "$UFW_RULES_FILE").backup-*" -printf '%T@ %p\n' 2>/dev/null \
        | sort -rn | tail -n +$((max_keep + 1)) | cut -d' ' -f2- | xargs -r rm -f 2>/dev/null || true

    if [ "$ENABLE_IPV6" = true ]; then
        find "$(dirname "$UFW_RULES_V6_FILE")" -maxdepth 1 -name "$(basename "$UFW_RULES_V6_FILE").backup-*" -printf '%T@ %p\n' 2>/dev/null \
            | sort -rn | tail -n +$((max_keep + 1)) | cut -d' ' -f2- | xargs -r rm -f 2>/dev/null || true
    fi

    rm -f "${UFW_RULES_FILE}.backup-"*.clean "${UFW_RULES_V6_FILE}.backup-"*.clean 2>/dev/null || true
}

check_requirements() {
    local enable_ufw="${1:-true}"

    [ "$EUID" -ne 0 ] && { echo "Please run as root"; exit 1; }

    if ! command -v ufw >/dev/null 2>&1; then
        apt-get update && apt-get install -y ufw ipset curl 2>/dev/null || \
        dnf install -y ufw ipset curl 2>/dev/null || \
        yum install -y ufw ipset curl 2>/dev/null
    fi

    command -v ipset >/dev/null 2>&1 || apt-get install -y ipset
    command -v curl >/dev/null 2>&1 || { echo "ERROR: curl required"; exit 1; }

    # CRITICAL: Ensure all ipsets referenced by before.rules exist BEFORE enabling UFW.
    # If ipsets are missing (e.g., after reboot, failed persistence), UFW enable will fail
    # with "Set ufw-feed-XXX doesn't exist" and block ALL traffic including DNS.
    ensure_ipsets_exist

    # NOTE: Do NOT enable UFW here. UFW should only be enabled/reloaded in
    # apply_ufw_rules() AFTER ipsets are verified. Enabling too early causes
    # lockouts when ipsets aren't fully loaded yet.
    if [ "$enable_ufw" = true ]; then
        if ! ufw status | grep -q "Status: active"; then
            log_message "UFW is not active — it will be enabled after rules are applied"
        fi
    fi

    cleanup_old_backups
}

_ensure_feed_ipset() {
    local name="$1"

    ipset list "${IPSET_PREFIX}-${name}" >/dev/null 2>&1 || \
        ipset create "${IPSET_PREFIX}-${name}" hash:net family inet hashsize 4096 maxelem 200000 2>/dev/null || true

    if [ "$ENABLE_IPV6" = true ]; then
        ipset list "${IPSET_PREFIX}-${name}-v6" >/dev/null 2>&1 || \
            ipset create "${IPSET_PREFIX}-${name}-v6" hash:net family inet6 hashsize 4096 maxelem 200000 2>/dev/null || true
    fi
}

ensure_ipsets_exist() {
    if [ -f /etc/ipset.conf ]; then
        ipset restore -f /etc/ipset.conf 2>/dev/null || true
    fi

    ipset list "$WHITELIST_IPSET" >/dev/null 2>&1 || \
        ipset create "$WHITELIST_IPSET" hash:net family inet hashsize 1024 maxelem 10000 2>/dev/null || true

    if [ "$ENABLE_IPV6" = true ]; then
        ipset list "$WHITELIST_IPSET_V6" >/dev/null 2>&1 || \
            ipset create "$WHITELIST_IPSET_V6" hash:net family inet6 hashsize 1024 maxelem 10000 2>/dev/null || true
    fi

    for_each_enabled_feed _ensure_feed_ipset
}

validate_feed_name() {
    local name="$1"
    if [ -z "$name" ]; then
        echo "ERROR: Feed name cannot be empty"; return 1
    fi
    if [[ ! "$name" =~ ^[a-zA-Z0-9_-]+$ ]]; then
        echo "ERROR: Feed name '$name' contains invalid characters (only a-z, 0-9, _, - allowed)"; return 1
    fi
    if [ "${#name}" -gt 20 ]; then
        echo "ERROR: Feed name '$name' too long (max 20 chars, ipset name limit)"; return 1
    fi
}

create_directory_structure() {
    mkdir -p "$CONFIG_DIR" "$CACHE_DIR"
    touch "$LOG_FILE"
    chmod 700 "$CONFIG_DIR"
    chmod 600 "$LOG_FILE"
}

initialize_feeds_config() {
    local has_feeds
    has_feeds=$(grep -c '^[01]|' "$FEEDS_CONFIG" 2>/dev/null || true)

    if [ -f "$FEEDS_CONFIG" ] && [ "$has_feeds" -gt 0 ]; then
        log_message "Feeds configuration already exists with $has_feeds feeds"
        return
    fi

    log_message "Creating feeds configuration..."

    [ -f "$FEEDS_CONFIG" ] && mv "$FEEDS_CONFIG" "${FEEDS_CONFIG}.old-$(date +%Y%m%d-%H%M%S)"

    cat > "$FEEDS_CONFIG" <<'EOF'
# Threat Intelligence Feeds Configuration
# Format: ENABLED|NAME|URL|TYPE|DESCRIPTION
#
# ENABLED: 1 (enabled) or 0 (disabled)
# NAME: Unique feed identifier
# URL: Feed URL (http/https) or local file (file:///path/to/file)
# TYPE: Format type (plain, cidr, commented)
#   plain     - One IP/CIDR per line, no comments
#   cidr      - IP/CIDR with optional inline comments/fields
#   commented - Lines starting with # or ; are ignored, IPs extracted
# DESCRIPTION: Feed description

1|cinsarmy|http://cinsscore.com/list/ci-badguys.txt|plain|CINS Army Malicious IPs
1|firehol-level1|https://raw.githubusercontent.com/ktsaou/blocklist-ipsets/master/firehol_level1.netset|cidr|FireHOL Level 1 - Most aggressive attackers
1|firehol-level2|https://raw.githubusercontent.com/ktsaou/blocklist-ipsets/master/firehol_level2.netset|cidr|FireHOL Level 2 - Attacks in last 48h
0|firehol-level3|https://raw.githubusercontent.com/ktsaou/blocklist-ipsets/master/firehol_level3.netset|cidr|FireHOL Level 3 - Attacks in last 30d
1|ipsum-1|https://raw.githubusercontent.com/stamparm/ipsum/master/levels/1.txt|plain|IPsum Level 1 - Most dangerous
0|ipsum-2|https://raw.githubusercontent.com/stamparm/ipsum/master/levels/2.txt|plain|IPsum Level 2 - Dangerous
0|ipsum-3|https://raw.githubusercontent.com/stamparm/ipsum/master/levels/3.txt|plain|IPsum Level 3 - Suspicious
0|spamhaus-drop|https://www.spamhaus.org/drop/drop.txt|commented|Spamhaus DROP List
0|spamhaus-edrop|https://www.spamhaus.org/drop/edrop.txt|commented|Spamhaus EDROP List
1|spamhaus-dropv6|https://www.spamhaus.org/drop/dropv6.txt|commented|Spamhaus DROP V6 List
0|feodo-tracker|https://feodotracker.abuse.ch/downloads/ipblocklist.txt|commented|Feodo Tracker C2 IPs
0|sslbl-aggressive|https://sslbl.abuse.ch/blacklist/sslipblacklist_aggressive.txt|commented|SSL Blacklist Aggressive
0|sslbl-all|https://sslbl.abuse.ch/blacklist/sslipblacklist.txt|commented|SSL Blacklist All
1|blocklist-de|https://lists.blocklist.de/lists/all.txt|plain|Blocklist.de All Attacks
0|greensnow|https://blocklist.greensnow.co/greensnow.txt|plain|GreenSnow Blacklist
0|emergingthreats|https://rules.emergingthreats.net/fwrules/emerging-Block-IPs.txt|plain|Emerging Threats IPs
0|bruteforce-ssh|https://lists.blocklist.de/lists/ssh.txt|plain|SSH Bruteforce Attempts
1|binarydefense|https://www.binarydefense.com/banlist.txt|plain|Binary Defense Blacklist
1|bruteforce-bl|https://danger.rulez.sk/projects/bruteforceblocker/blist.php|commented|BruteForce Blocker
0|dshield-top|https://www.dshield.org/block.txt|commented|DShield Top Attackers
1|dshield-fhol|https://iplists.firehol.org/files/dshield.netset|commented|Dshield FireHol top 20
0|tor-exit|https://check.torproject.org/torbulkexitlist|plain|TOR Exit Nodes (optional)
0|abuseipdb-1d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-1d.ipv4|commented|AbuseIPDB with confidence score 100 1 day
0|abuseipdb-3d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-3d.ipv4|commented|AbuseIPDB with confidence score 100 3 day
0|abuseipdb-7d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-7d.ipv4|commented|AbuseIPDB with confidence score 100 7 day
1|abuseipdb-14d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-14d.ipv4|commented|AbuseIPDB with confidence score 100 14 day
0|abuseipdb-30d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-30d.ipv4|commented|AbuseIPDB with confidence score 100 30 day


# Add custom feeds below this line
EOF
    chmod 600 "$FEEDS_CONFIG"
}

_setup_feed_ipset() {
    local name="$1"

    if ! ipset list "${IPSET_PREFIX}-${name}" >/dev/null 2>&1; then
        ipset create "${IPSET_PREFIX}-${name}" hash:net family inet hashsize 4096 maxelem 200000
        log_message "  Created ipset: ${IPSET_PREFIX}-${name}"
    fi

    if [ "$ENABLE_IPV6" = true ] && ! ipset list "${IPSET_PREFIX}-${name}-v6" >/dev/null 2>&1; then
        ipset create "${IPSET_PREFIX}-${name}-v6" hash:net family inet6 hashsize 4096 maxelem 200000
        log_message "  Created ipset: ${IPSET_PREFIX}-${name}-v6"
    fi
}

setup_ipsets() {
    log_message "Setting up ipsets (per-feed mode)..."

    if ! ipset list "$WHITELIST_IPSET" >/dev/null 2>&1; then
        ipset create "$WHITELIST_IPSET" hash:net family inet hashsize 1024 maxelem 10000
        ipset add "$WHITELIST_IPSET" 127.0.0.1
    fi

    if [ "$ENABLE_IPV6" = true ] && ! ipset list "$WHITELIST_IPSET_V6" >/dev/null 2>&1; then
        ipset create "$WHITELIST_IPSET_V6" hash:net family inet6 hashsize 1024 maxelem 10000
        ipset add "$WHITELIST_IPSET_V6" ::1
    fi

    for_each_enabled_feed _setup_feed_ipset
    setup_ipset_persistence
}

setup_ipset_persistence() {
    local script_path
    script_path=$(readlink -f "$0")

    cat > /etc/systemd/system/ipset-persistent.service <<EOF
[Unit]
Description=ipset persistent configuration
Before=network-pre.target ufw.service
Wants=network-pre.target

[Service]
Type=oneshot
RemainAfterExit=yes
# Restore saved ipsets, then ensure all feed ipsets exist (creates empty ones
# for any that failed to restore). This prevents UFW from failing on boot
# when before.rules references ipsets that don't exist.
ExecStart=/bin/sh -c '/sbin/ipset restore -f /etc/ipset.conf 2>/dev/null || true; ${script_path} _ensure-ipsets 2>/dev/null || true'
ExecStop=/sbin/ipset save -f /etc/ipset.conf

[Install]
WantedBy=multi-user.target
EOF

    ipset save > /etc/ipset.conf
    systemctl daemon-reload 2>/dev/null || true
    systemctl enable ipset-persistent.service 2>/dev/null || true
}

download_feed() {
    local url="$1" output="$2"

    # Local file support: file:///path or file://path
    if [[ "$url" == file://* ]]; then
        local local_path="${url#file://}"
        if [ ! -f "$local_path" ]; then
            log_message "    Local file not found: $local_path"
            return 1
        fi
        cp "$local_path" "$output" 2>/dev/null || return 1
        return 0
    fi

    local http_code
    http_code=$(curl -f -s -m 60 --connect-timeout 10 -L \
        -A "ufw-threat-feeds-per-feed/1.0" \
        -w "%{http_code}" -o "$output" "$url" 2>/dev/null) || true

    if [ ! -s "$output" ]; then
        log_message "    Download failed for $url (HTTP $http_code, empty response)"
        return 1
    fi
    return 0
}

parse_feed() {
    local file="$1" type="$2" output_v4="$3" output_v6="$4"

    : > "$output_v4"
    : > "$output_v6"

    # Strip carriage returns from all feed types (common in downloaded feeds)
    local cleaned
    cleaned=$(mktemp)
    tr -d '\r' < "$file" > "$cleaned"

    # IPv6 filter: require either a hex letter [a-fA-F], or 3+ colon-separated
    # groups, or ::. This excludes timestamps like 14:34:21 which only have
    # digits and exactly two colon-separated groups.
    local v6_filter='([a-fA-F]|:.*:.*:|::)'

    case "$type" in
        plain)
            grep -E '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?$' "$cleaned" >> "$output_v4" 2>/dev/null || true
            if [ "$ENABLE_IPV6" = true ]; then
                grep -E '^[0-9a-fA-F:]+(/[0-9]+)?$' "$cleaned" | grep ':' \
                    | grep -E "$v6_filter" >> "$output_v6" 2>/dev/null || true
            fi
            ;;
        cidr)
            grep -oE '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?' "$cleaned" \
                | grep -v '^$' >> "$output_v4" 2>/dev/null || true
            if [ "$ENABLE_IPV6" = true ]; then
                grep -oE '^[0-9a-fA-F:]+(/[0-9]+)?' "$cleaned" \
                    | grep ':' | grep -E "$v6_filter" | grep -v '^$' >> "$output_v6" 2>/dev/null || true
            fi
            ;;
        commented)
            grep -v -E '^[#;]|^$' "$cleaned" \
                | grep -oE '[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?' >> "$output_v4" 2>/dev/null || true
            if [ "$ENABLE_IPV6" = true ]; then
                grep -v -E '^[#;]|^$' "$cleaned" \
                    | grep -oE '[0-9a-fA-F:]+(/[0-9]+)?' \
                    | grep -E '^[0-9a-fA-F]{1,4}:[0-9a-fA-F:]+' \
                    | grep -E "$v6_filter" >> "$output_v6" 2>/dev/null || true
            fi
            ;;
    esac

    rm -f "$cleaned"
}

_clean_stale_cache() {
    local enabled_feeds="$1"
    local cleaned=0

    for cache_file in "$CACHE_DIR"/*.raw "$CACHE_DIR"/*-v4.parsed "$CACHE_DIR"/*-v6.parsed; do
        [ -f "$cache_file" ] || continue
        local bn feed_name
        bn=$(basename "$cache_file")
        feed_name="${bn%%.raw}"
        feed_name="${feed_name%%-v4.parsed}"
        feed_name="${feed_name%%-v6.parsed}"

        if ! grep -q "^${feed_name}$" <<< "$enabled_feeds"; then
            rm -f "$cache_file" && cleaned=$((cleaned + 1))
        fi
    done

    [ "$cleaned" -gt 0 ] && log_message "  Cleaned $cleaned stale cache files"
}

_load_ipset_v4() {
    local name="$1" v4_file="$2"
    local batch_file restore_err
    batch_file=$(mktemp /tmp/ipset-v4-XXXXXX)

    {
        echo "create ${IPSET_PREFIX}-${name}-tmp hash:net family inet hashsize 4096 maxelem 200000"
        echo "flush ${IPSET_PREFIX}-${name}-tmp"
        while IFS= read -r ip; do
            [ -z "$ip" ] && continue
            echo "add ${IPSET_PREFIX}-${name}-tmp $ip"
        done < "$v4_file"
        echo "swap ${IPSET_PREFIX}-${name} ${IPSET_PREFIX}-${name}-tmp"
        echo "destroy ${IPSET_PREFIX}-${name}-tmp"
    } > "$batch_file"

    if ! restore_err=$(ipset restore -exist < "$batch_file" 2>&1); then
        log_message "WARNING: ipset restore failed for ${name} (v4): $restore_err"
        ipset destroy "${IPSET_PREFIX}-${name}-tmp" 2>/dev/null || true
    fi
    rm -f "$batch_file"
}

_load_ipset_v6() {
    local name="$1" v6_file="$2"
    local batch_file restore_err
    batch_file=$(mktemp /tmp/ipset-v6-XXXXXX)

    {
        echo "create ${IPSET_PREFIX}-${name}-v6-tmp hash:net family inet6 hashsize 4096 maxelem 200000"
        echo "flush ${IPSET_PREFIX}-${name}-v6-tmp"
        while IFS= read -r ip; do
            [ -z "$ip" ] && continue
            echo "add ${IPSET_PREFIX}-${name}-v6-tmp $ip"
        done < "$v6_file"
        echo "swap ${IPSET_PREFIX}-${name}-v6 ${IPSET_PREFIX}-${name}-v6-tmp"
        echo "destroy ${IPSET_PREFIX}-${name}-v6-tmp"
    } > "$batch_file"

    if ! restore_err=$(ipset restore -exist < "$batch_file" 2>&1); then
        log_message "WARNING: ipset restore failed for ${name} (v6): $restore_err"
        ipset destroy "${IPSET_PREFIX}-${name}-v6-tmp" 2>/dev/null || true
    fi
    rm -f "$batch_file"
}

update_feeds() {
    log_message "Starting per-feed update..."

    if [ ! -f "$FEEDS_CONFIG" ]; then
        echo "ERROR: Feeds config not found: $FEEDS_CONFIG"
        echo "Run 'install' command first"
        exit 1
    fi

    local enabled_count
    enabled_count=$(grep -c '^1|' "$FEEDS_CONFIG" 2>/dev/null || true)
    if [ "$enabled_count" -eq 0 ]; then
        echo "ERROR: No enabled feeds found in $FEEDS_CONFIG"
        echo "Check the config file format"
        exit 1
    fi

    log_message "Found $enabled_count enabled feeds"

    local enabled_feeds
    enabled_feeds=$(grep '^1|' "$FEEDS_CONFIG" 2>/dev/null | cut -d'|' -f2)

    # NOTE: Do NOT destroy ipsets for disabled feeds here. The before.rules may still
    # reference them (if apply-rules hasn't been re-run). Destroying in-use ipsets causes
    # "Set doesn't exist" on next UFW reload, which blocks all traffic.
    # Ipset cleanup happens safely in cmd_disable_feed/cmd_remove_feed after rules are regenerated.
    _clean_stale_cache "$enabled_feeds"

    local total_feeds=0
    local failed_feeds=0

    local enabled name url type description
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue

        total_feeds=$((total_feeds + 1))
        log_message "Processing feed: $name"

        local raw="$CACHE_DIR/${name}.raw"
        local v4_file="$CACHE_DIR/${name}-v4.parsed"
        local v6_file="$CACHE_DIR/${name}-v6.parsed"

        if download_feed "$url" "$raw" && parse_feed "$raw" "$type" "$v4_file" "$v6_file"; then
            local count_v4 count_v6
            count_v4=$(wc -l < "$v4_file" 2>/dev/null || echo 0)
            count_v6=0
            [ "$ENABLE_IPV6" = true ] && count_v6=$(wc -l < "$v6_file" 2>/dev/null || echo 0)

            _load_ipset_v4 "$name" "$v4_file"
            [ "$ENABLE_IPV6" = true ] && _load_ipset_v6 "$name" "$v6_file"

            log_message "  $name: $count_v4 IPv4, $count_v6 IPv6"
        else
            log_message "  FAILED: $name"
            failed_feeds=$((failed_feeds + 1))
        fi
    done < "$FEEDS_CONFIG"

    ipset save > /etc/ipset.conf
    log_message "Updated $total_feeds feeds ($failed_feeds failed)"
}

# Build iptables rules block for IPv4 or IPv6
# Args: v4|v6 output_file
_build_rules_block() {
    local family="$1" output="$2"
    local chain_prefix whitelist_set set_suffix log_tag

    if [ "$family" = "v4" ]; then
        chain_prefix="ufw-before-input"
        whitelist_set="$WHITELIST_IPSET"
        set_suffix=""
        log_tag="THREAT"
    else
        chain_prefix="ufw6-before-input"
        whitelist_set="$WHITELIST_IPSET_V6"
        set_suffix="-v6"
        log_tag="THREAT-v6"
    fi

    cat > "$output" <<EOF

# UFW THREAT FEEDS - PER-FEED MODE - START
# Whitelist bypass
-A ${chain_prefix} -m set --match-set ${whitelist_set} src -j ACCEPT

EOF

    local enabled name url type description
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue

        cat >> "$output" <<EOF
# $description
-A ${chain_prefix} -m set --match-set ${IPSET_PREFIX}-${name}${set_suffix} src -m limit --limit 5/min -j LOG --log-prefix "[${log_tag}:${name}] "
-A ${chain_prefix} -m set --match-set ${IPSET_PREFIX}-${name}${set_suffix} src -j DROP
EOF
    done < "$FEEDS_CONFIG"

    echo "# UFW THREAT FEEDS - PER-FEED MODE - END" >> "$output"
}

# Insert rules into a UFW template file and validate
# Args: template_file rules_file output_file
# Returns 0 on success, 1 on validation failure
_insert_and_validate_rules() {
    local template="$1" rules_file="$2" output="$3"
    local insert_line

    insert_line=$(grep -n "^# End required lines" "$template" | cut -d: -f1)
    if [ -z "$insert_line" ]; then
        log_message "ERROR: Could not find '# End required lines' in $template"
        return 1
    fi

    head -n "$insert_line" "$template" > "$output"
    cat "$rules_file" >> "$output"
    tail -n +"$((insert_line + 1))" "$template" >> "$output"

    local filter_count
    filter_count=$(grep -c '^\*filter' "$output" 2>/dev/null || true)
    if [ "$filter_count" -ne 1 ]; then
        log_message "ERROR: Generated rules file has $filter_count *filter blocks (expected 1)"
        return 1
    fi

    return 0
}

_verify_ipsets_callback() {
    local name="$1"

    if ! ipset list "${IPSET_PREFIX}-${name}" >/dev/null 2>&1; then
        log_message "ERROR: Required ipset ${IPSET_PREFIX}-${name} is missing"
        _MISSING_SETS=$((_MISSING_SETS + 1))
    fi
    if [ "$ENABLE_IPV6" = true ] && ! ipset list "${IPSET_PREFIX}-${name}-v6" >/dev/null 2>&1; then
        log_message "ERROR: Required ipset ${IPSET_PREFIX}-${name}-v6 is missing"
        _MISSING_SETS=$((_MISSING_SETS + 1))
    fi
}

apply_ufw_rules() {
    log_message "Applying UFW rules (per-feed)..."

    if [ ! -f /usr/share/ufw/before.rules ]; then
        log_message "ERROR: UFW default template /usr/share/ufw/before.rules not found"
        return 1
    fi

    local tmpdir
    tmpdir=$(mktemp -d)
    trap 'rm -rf "$tmpdir"' RETURN

    [ -f "$UFW_RULES_FILE" ] && cp "$UFW_RULES_FILE" "${UFW_RULES_FILE}.backup-$(date +%Y%m%d-%H%M%S)"
    [ "$ENABLE_IPV6" = true ] && [ -f "$UFW_RULES_V6_FILE" ] && \
        cp "$UFW_RULES_V6_FILE" "${UFW_RULES_V6_FILE}.backup-$(date +%Y%m%d-%H%M%S)"

    cp /usr/share/ufw/before.rules "$UFW_RULES_FILE"
    [ "$ENABLE_IPV6" = true ] && cp /usr/share/ufw/before6.rules "$UFW_RULES_V6_FILE"

    log_message "  Starting from clean UFW templates"

    # Build and insert IPv4 rules
    local v4_rules="$tmpdir/v4_rules"
    local v4_output="$tmpdir/v4_output"
    _build_rules_block "v4" "$v4_rules"

    if ! _insert_and_validate_rules "$UFW_RULES_FILE" "$v4_rules" "$v4_output"; then
        log_message "  Aborting to prevent corruption."
        return 1
    fi
    mv "$v4_output" "$UFW_RULES_FILE"
    log_message "  IPv4 rules generated and validated"

    # Build and insert IPv6 rules
    if [ "$ENABLE_IPV6" = true ]; then
        local v6_rules="$tmpdir/v6_rules"
        local v6_output="$tmpdir/v6_output"
        _build_rules_block "v6" "$v6_rules"

        if _insert_and_validate_rules "$UFW_RULES_V6_FILE" "$v6_rules" "$v6_output"; then
            mv "$v6_output" "$UFW_RULES_V6_FILE"
            log_message "  IPv6 rules generated and validated"
        else
            log_message "  Aborting IPv6 rules. Keeping IPv4 only."
        fi
    fi

    # CRITICAL: Ensure all ipsets exist BEFORE reloading UFW
    log_message "  Verifying ipsets exist..."
    ensure_ipsets_exist
    setup_ipsets

    _MISSING_SETS=0
    for_each_enabled_feed _verify_ipsets_callback

    if [ "$_MISSING_SETS" -gt 0 ]; then
        log_message "ERROR: $_MISSING_SETS required ipsets missing. Aborting UFW reload to prevent lockout."
        return 1
    fi

    ipset save > /etc/ipset.conf

    log_message "  Reloading UFW..."
    if ufw status | grep -q "Status: active"; then
        ufw reload
    else
        ufw --force enable
    fi

    cleanup_old_backups
    log_message "UFW rules applied and validated successfully"
}

setup_auto_update() {
    [ "$ENABLE_AUTO_UPDATE" = false ] && return

    local script_path
    script_path=$(readlink -f "$0")

    cat > /etc/systemd/system/ufw-threat-feeds-update.service <<EOF
[Unit]
Description=Update UFW threat feeds (per-feed)
After=network-online.target

[Service]
Type=oneshot
ExecStart=$script_path update
ExecStartPost=/bin/sh -c 'ipset save > /etc/ipset.conf'
EOF

    cat > /etc/systemd/system/ufw-threat-feeds-update.timer <<EOF
[Unit]
Description=Update UFW threat feeds $UPDATE_INTERVAL

[Timer]
Unit=ufw-threat-feeds-update.service
OnCalendar=$UPDATE_INTERVAL
Persistent=true
RandomizedDelaySec=1800

[Install]
WantedBy=timers.target
EOF

    systemctl daemon-reload
    systemctl enable --now ufw-threat-feeds-update.timer
}

create_management_commands() {
    cat > /usr/local/bin/ufw-whitelist <<'EOF'
#!/bin/bash
[ -z "$1" ] && { echo "Usage: ufw-whitelist <IP|CIDR>"; exit 1; }
if [[ "$1" == *:* ]]; then
    ipset add ufw-whitelist-v6 "$1" && echo "Whitelisted IPv6: $1"
else
    ipset add ufw-whitelist "$1" && echo "Whitelisted IPv4: $1"
fi
ipset save > /etc/ipset.conf
EOF

    local script_path
    script_path=$(readlink -f "$0")
    cat > /usr/local/bin/ufw-threat-reload <<EOF
#!/bin/bash
$script_path apply-rules
EOF

    chmod +x /usr/local/bin/ufw-{whitelist,threat-reload}
}

cmd_show_stats() {
    echo "=========================================="
    echo "Per-Feed Blocking Statistics"
    echo "=========================================="
    printf "%-25s %10s %10s %12s\n" "FEED" "IPv4 IPs" "IPv6 IPs" "BLOCKS (1h)"
    echo "-------------------------------------------------------------------"

    if [ ! -f "$FEEDS_CONFIG" ]; then
        echo "ERROR: Config not found. Run 'install' first."
        return 1
    fi

    local journal_cache
    journal_cache=$(journalctl --since "1 hour ago" 2>/dev/null || true)

    local enabled name url type description
    local v4_count v6_count blocks_v4 blocks_v6 blocks
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue

        v4_count=$(ipset list "${IPSET_PREFIX}-${name}" 2>/dev/null | awk '/^Members:/{found=1; next} found' | wc -l)
        v4_count=${v4_count:-0}

        v6_count=0
        if [ "$ENABLE_IPV6" = true ]; then
            v6_count=$(ipset list "${IPSET_PREFIX}-${name}-v6" 2>/dev/null | awk '/^Members:/{found=1; next} found' | wc -l)
            v6_count=${v6_count:-0}
        fi

        blocks_v4=$(grep -c "\[THREAT:${name}\]" <<< "$journal_cache" 2>/dev/null || true)
        blocks_v6=$(grep -c "\[THREAT-v6:${name}\]" <<< "$journal_cache" 2>/dev/null || true)
        blocks=$(( ${blocks_v4:-0} + ${blocks_v6:-0} ))

        printf "%-25s %10d %10d %12d\n" "$name" "$v4_count" "$v6_count" "$blocks"
    done < "$FEEDS_CONFIG"
}

_list_feed_entry() {
    local feed_enabled="$1" name="$2" url="$3" type="$4" description="$5"
    local status="DISABLED"
    [ "$feed_enabled" = "1" ] && status="ENABLED"
    printf "%-10s %-25s %s\n" "$status" "$name" "$description"
}

cmd_list_feeds() {
    printf "%-10s %-25s %s\n" "STATUS" "NAME" "DESCRIPTION"
    echo "-------------------------------------------------------------------"
    for_each_feed _list_feed_entry
}

cmd_add_feed() {
    validate_feed_name "$FEED_NAME" || exit 1
    [ -z "$FEED_URL" ] && { echo "Usage: $0 add-feed <NAME> <URL>"; exit 1; }
    grep -q "^[01]|${FEED_NAME}|" "$FEEDS_CONFIG" 2>/dev/null && { echo "Feed exists"; exit 1; }
    echo "1|${FEED_NAME}|${FEED_URL}|plain|Custom: ${FEED_NAME}" >> "$FEEDS_CONFIG"
    log_message "Added feed: $FEED_NAME"
}

cmd_remove_feed() {
    validate_feed_name "$FEED_NAME" || exit 1
    sed -i "/^[01]|${FEED_NAME}|/d" "$FEEDS_CONFIG"
    log_message "Removed feed: $FEED_NAME"

    log_message "Regenerating UFW rules..."
    apply_ufw_rules || return 1

    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}" 2>/dev/null || true
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}-v6" 2>/dev/null || true
}

cmd_enable_feed() {
    validate_feed_name "$FEED_NAME" || exit 1
    sed -i "s/^0|${FEED_NAME}|/1|${FEED_NAME}|/" "$FEEDS_CONFIG"
    log_message "Enabled: $FEED_NAME"

    log_message "Regenerating UFW rules..."
    apply_ufw_rules
}

cmd_disable_feed() {
    validate_feed_name "$FEED_NAME" || exit 1
    sed -i "s/^1|${FEED_NAME}|/0|${FEED_NAME}|/" "$FEEDS_CONFIG"
    log_message "Disabled: $FEED_NAME"

    log_message "Regenerating UFW rules..."
    apply_ufw_rules || return 1

    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}" 2>/dev/null || true
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}-v6" 2>/dev/null || true
}

cmd_whitelist_add() {
    [ -z "$WHITELIST_IP" ] && { echo "Usage: $0 whitelist-add <IP|CIDR>"; exit 1; }

    if [[ "$WHITELIST_IP" == *:* ]]; then
        if ipset add "$WHITELIST_IPSET_V6" "$WHITELIST_IP" 2>/dev/null; then
            log_message "Added to IPv6 whitelist: $WHITELIST_IP"
        else
            echo "Failed to add $WHITELIST_IP"; exit 1
        fi
    else
        if ipset add "$WHITELIST_IPSET" "$WHITELIST_IP" 2>/dev/null; then
            log_message "Added to IPv4 whitelist: $WHITELIST_IP"
        else
            echo "Failed to add $WHITELIST_IP"; exit 1
        fi
    fi

    ipset save > /etc/ipset.conf
}

cmd_whitelist_init() {
    log_message "Initializing whitelist with private networks..."

    local private_networks=(
        "10.0.0.0/8"
        "172.16.0.0/12"
        "192.168.0.0/16"
        "169.254.0.0/16"
        "127.0.0.0/8"
    )

    local private_networks_v6=(
        "fc00::/7"
        "fe80::/10"
        "::1"
    )

    echo "Adding IPv4 private networks to whitelist..."
    for net in "${private_networks[@]}"; do
        if ipset add "$WHITELIST_IPSET" "$net" 2>/dev/null; then
            echo "  + $net"
        else
            echo "  - $net (already exists or error)"
        fi
    done

    if [ "$ENABLE_IPV6" = true ]; then
        echo "Adding IPv6 private networks to whitelist..."
        for net in "${private_networks_v6[@]}"; do
            if ipset add "$WHITELIST_IPSET_V6" "$net" 2>/dev/null; then
                echo "  + $net"
            else
                echo "  - $net (already exists or error)"
            fi
        done
    fi

    ipset save > /etc/ipset.conf
    log_message "Whitelist initialized with RFC1918/private networks"
}

cmd_whitelist_list() {
    echo "=========================================="
    echo "IPv4 Whitelist ($WHITELIST_IPSET)"
    echo "=========================================="
    ipset list "$WHITELIST_IPSET" 2>/dev/null | grep -E '^[0-9]' || echo "No entries"

    if [ "$ENABLE_IPV6" = true ]; then
        echo ""
        echo "=========================================="
        echo "IPv6 Whitelist ($WHITELIST_IPSET_V6)"
        echo "=========================================="
        ipset list "$WHITELIST_IPSET_V6" 2>/dev/null | grep -E '^[0-9a-fA-F:]' || echo "No entries"
    fi
}

cmd_clean_cache() {
    log_message "Cleaning cache for disabled feeds..."

    local removed=0
    local kept=0

    local enabled_feeds
    enabled_feeds=$(grep '^1|' "$FEEDS_CONFIG" 2>/dev/null | cut -d'|' -f2)

    for cache_file in "$CACHE_DIR"/*.raw "$CACHE_DIR"/*-v4.parsed "$CACHE_DIR"/*-v6.parsed; do
        [ -f "$cache_file" ] || continue

        local bn feed_name
        bn=$(basename "$cache_file")
        feed_name="${bn%%.raw}"
        feed_name="${feed_name%%-v4.parsed}"
        feed_name="${feed_name%%-v6.parsed}"

        if ! grep -q "^${feed_name}$" <<< "$enabled_feeds"; then
            rm -f "$cache_file"
            removed=$((removed + 1))
        else
            kept=$((kept + 1))
        fi
    done

    log_message "Removed $removed cache files, kept $kept active feeds"
}

cmd_test_rules() {
    log_message "Testing UFW rule generation (dry-run mode)..."

    if [ ! -f /usr/share/ufw/before.rules ]; then
        echo "ERROR: UFW default template /usr/share/ufw/before.rules not found"
        return 1
    fi

    local test_dir
    test_dir=$(mktemp -d)
    trap 'rm -rf "$test_dir"' RETURN

    local test_v4="$test_dir/before.rules.test"
    cp /usr/share/ufw/before.rules "$test_v4"
    [ "$ENABLE_IPV6" = true ] && cp /usr/share/ufw/before6.rules "$test_dir/before6.rules.test"

    local v4_rules="$test_dir/v4_rules"
    local v4_output="$test_dir/v4_output"

    _build_rules_block "v4" "$v4_rules"

    local feed_count
    feed_count=$(grep -c '^1|' "$FEEDS_CONFIG" 2>/dev/null || true)
    echo "Generated rules for $feed_count enabled feeds"

    if ! _insert_and_validate_rules "$test_v4" "$v4_rules" "$v4_output"; then
        echo "VALIDATION FAILED"
        return 1
    fi

    echo "Validation passed: exactly 1 *filter block found"

    local total_lines rule_lines
    total_lines=$(wc -l < "$v4_output")
    rule_lines=$(grep -c "^-A " "$v4_output" 2>/dev/null || true)

    echo "Generated $rule_lines iptables rules in $total_lines total lines"
    echo ""
    echo "=========================================="
    echo "Sample of generated rules:"
    echo "=========================================="
    grep "# UFW THREAT FEEDS" -A 10 "$v4_output" | head -15
    echo "..."
    echo ""
    echo "=========================================="
    echo "Test passed - rules would be generated safely"
    echo "  To apply these rules, run: $0 apply-rules"
    echo "=========================================="
}

cmd_install() {
    log_message "Installing per-feed threat blocking..."
    check_requirements
    create_directory_structure
    initialize_feeds_config
    setup_ipsets
    update_feeds

    # Ensure SSH is allowed before applying threat rules (only during install)
    ufw limit "$SSH_PORT/tcp" 2>/dev/null || ufw allow "$SSH_PORT/tcp"

    apply_ufw_rules
    setup_auto_update
    create_management_commands

    echo ""
    echo "=========================================="
    echo "Per-Feed Installation Complete"
    echo "=========================================="
    echo "Mode: Per-feed ipsets (detailed tracking)"
    echo "Feeds: $(grep -c '^1|' "$FEEDS_CONFIG")"
    echo "IPv6: $ENABLE_IPV6"
    echo "Auto-update: $ENABLE_AUTO_UPDATE ($UPDATE_INTERVAL)"
    echo ""
    echo "Commands:"
    echo "  $0 show-stats     # View per-feed statistics"
    echo "  $0 update         # Update all feeds"
    echo "  ufw-whitelist IP  # Whitelist an IP"
    echo ""
    echo "Logs: grep 'THREAT:' /var/log/syslog"
    echo "=========================================="
}

main() {
    # Internal command used by ipset-persistent.service at boot
    if [ "${1:-}" = "_ensure-ipsets" ]; then
        ensure_ipsets_exist
        exit 0
    fi

    parse_args "$@"
    case "$COMMAND" in
        install) cmd_install ;;
        update)
            check_requirements false
            create_directory_structure
            update_feeds
            ;;
        apply-rules)
            check_requirements
            apply_ufw_rules
            ;;
        test-rules) cmd_test_rules ;;
        list-feeds) cmd_list_feeds ;;
        show-stats) cmd_show_stats ;;
        add-feed) cmd_add_feed ;;
        remove-feed) cmd_remove_feed ;;
        enable-feed) cmd_enable_feed ;;
        disable-feed) cmd_disable_feed ;;
        whitelist-add) cmd_whitelist_add ;;
        whitelist-init) cmd_whitelist_init ;;
        whitelist-list) cmd_whitelist_list ;;
        clean-cache) cmd_clean_cache ;;
        *) show_usage ;;
    esac
}

main "$@"