#!/bin/bash
################################################################################
# Script Name: iptables-blocklists.sh
# Version: 1.04
# Description: Per-feed iptables threat intelligence blocking with ipset
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
################################################################################
# Changelog:
# v1.04 - Fix IPv6 parser: reject timestamps (e.g. 14:34:21) misidentified
#          as IPv6 addresses by requiring hex letters, 3+ groups, or ::
# v1.03 - Fix remove-feed/disable-feed: delete iptables rules before destroying
#          ipsets. Remove dangerous ipset destruction of disabled feeds during
#          update. Always swap ipsets even when feed is empty (clears stale
#          blocks). Use -I (insert) for IPv6 rules instead of -A (append).
#          Fix show-stats to count both IPv4+IPv6 blocks with cached journal.
#          Add curl requirement check. Fix ipset member counting.
################################################################################
# Don't use 'set -e' - causes issues with ipset error handling

CONFIG_DIR="/etc/iptables-threats"
FEEDS_CONFIG="$CONFIG_DIR/feeds.conf"
CACHE_DIR="$CONFIG_DIR/cache"
BACKUP_DIR="$CONFIG_DIR/backups"
IPSET_PREFIX="iptables-feed"
WHITELIST_IPSET="iptables-whitelist"
WHITELIST_IPSET_V6="iptables-whitelist-v6"
LOG_FILE="/var/log/iptables-threats.log"
SSH_PORT="22"
ENABLE_AUTO_UPDATE=true
UPDATE_INTERVAL="daily"
ENABLE_IPV6=true
MAX_BACKUPS=5

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

PER-FEED VERSION for iptables: Each threat feed gets its own ipset.
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 rules reload)
    apply-rules             Regenerate and apply iptables rules (use with caution!)
    test-rules              Test rule generation without applying (dry-run)
    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 updates
    --no-ipv6               Disable IPv6
    --update-interval TIME  hourly, daily, weekly (default: daily)

EXAMPLES:
    # Install with default feeds
    sudo $0 install

    # Update feeds manually (safe - only updates ipsets)
    sudo $0 update

    # Test rule generation (safe - no changes)
    sudo $0 test-rules

    # Apply rules after testing (regenerates iptables)
    sudo $0 apply-rules

    # Add custom feed
    sudo $0 add-feed "my-blocklist" "https://example.com/blocklist.txt"

    # View statistics
    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"
}

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
}

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: $1"; exit 1 ;;
        esac
    done
    [ -z "$COMMAND" ] && COMMAND="install"
}

check_requirements() {
    [ "$EUID" -ne 0 ] && { echo "Run as root"; exit 1; }
    
    # Install iptables, ipset, curl if needed
    if ! command -v iptables >/dev/null 2>&1 || ! command -v ipset >/dev/null 2>&1; then
        if command -v dnf >/dev/null 2>&1; then
            dnf install -y iptables ipset curl iptables-services
        elif command -v yum >/dev/null 2>&1; then
            yum install -y iptables ipset curl iptables-services
        elif command -v apt-get >/dev/null 2>&1; then
            apt-get update && apt-get install -y iptables ipset curl iptables-persistent
        else
            echo "Cannot install requirements automatically"
            exit 1
        fi
    fi

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

    ensure_ipsets_exist
}

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

cleanup_old_backups() {
    local backup_count
    backup_count=$(find "$BACKUP_DIR" -name 'iptables-save-*.txt' | wc -l)
    
    if [ "$backup_count" -gt "$MAX_BACKUPS" ]; then
        local to_delete=$((backup_count - MAX_BACKUPS))
        find "$BACKUP_DIR" -name 'iptables-save-*.txt' -type f | \
            sort | head -n "$to_delete" | xargs rm -f
        log_message "Cleaned up $to_delete old backups (keeping last $MAX_BACKUPS)"
    fi
}

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 confidence score 100 1 day
0|abuseipd-3d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-3d.ipv4|commented|AbuseIPDB confidence score 100 3 day
0|abuseipdb-7d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-7d.ipv4|commented|AbuseIPDB confidence score 100 7 day
1|abuseipdb-14d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-14d.ipv4|commented|AbuseIPDB confidence score 100 14 day
0|abuseipdb-30d|https://raw.githubusercontent.com/borestad/blocklist-abuseipdb/main/abuseipdb-s100-30d.ipv4|commented|AbuseIPDB confidence score 100 30 day
# Add custom feeds below this line
EOF
    chmod 600 "$FEEDS_CONFIG"
}

_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/sysconfig/ipset ]; then
        ipset restore -f /etc/sysconfig/ipset 2>/dev/null || true
    elif [ -f /etc/iptables/ipsets ]; then
        ipset restore -f /etc/iptables/ipsets 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
}

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

    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 2>/dev/null || true
    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 2>/dev/null || true
    fi

    for_each_enabled_feed _ensure_feed_ipset
}

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 "iptables-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" out_v4="$3" out_v6="$4"

    : > "$out_v4"
    : > "$out_v6"

    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" >> "$out_v4" 2>/dev/null || true
            if [ "$ENABLE_IPV6" = true ]; then
                grep -E '^[0-9a-fA-F:]+(/[0-9]+)?$' "$cleaned" | grep ':' \
                    | grep -E "$v6_filter" >> "$out_v6" 2>/dev/null || true
            fi
            ;;
        cidr)
            grep -oE '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?' "$cleaned" \
                | grep -v '^$' >> "$out_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 '^$' >> "$out_v6" 2>/dev/null || true
            fi
            ;;
        commented)
            grep -v -E '^[#;]|^$' "$cleaned" \
                | grep -oE '[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?' >> "$out_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" >> "$out_v6" 2>/dev/null || true
            fi
            ;;
    esac

    rm -f "$cleaned"
}

update_feeds() {
    log_message "Starting per-feed update (FAST ipset restore mode)..."
    
    # Auto-cleanup cache and ipsets for disabled feeds
    local enabled_feeds
    enabled_feeds=$(grep '^1|' "$FEEDS_CONFIG" 2>/dev/null | cut -d'|' -f2)
    local cleaned_cache=0
    
    # Clean cache files
    for cache_file in "$CACHE_DIR"/*.raw "$CACHE_DIR"/*-v4.parsed "$CACHE_DIR"/*-v6.parsed "$CACHE_DIR"/*-v4.restore "$CACHE_DIR"/*-v6.restore; do
        [ -f "$cache_file" ] || continue
        local bn
        bn=$(basename "$cache_file")
        local fn="${bn%%.raw}"; fn="${fn%%-v4.parsed}"; fn="${fn%%-v6.parsed}"; fn="${fn%%-v4.restore}"; fn="${fn%%-v6.restore}"
        if ! echo "$enabled_feeds" | grep -q "^${fn}$"; then
            rm -f "$cache_file" && cleaned_cache=$((cleaned_cache + 1))
        fi
    done
    
    [ "$cleaned_cache" -gt 0 ] && log_message "  Cleaned $cleaned_cache stale cache files"
    
    local total=0 failed=0
    
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue
        
        total=$((total + 1))
        log_message "Updating: $name"
        
        local raw="$CACHE_DIR/${name}.raw"
        local v4="$CACHE_DIR/${name}-v4.parsed"
        local v6="$CACHE_DIR/${name}-v6.parsed"
        
        if download_feed "$url" "$raw" && parse_feed "$raw" "$type" "$v4" "$v6"; then
            local c4 c6=0
            c4=$(wc -l < "$v4" 2>/dev/null || echo 0)
            [ "$ENABLE_IPV6" = true ] && c6=$(wc -l < "$v6" 2>/dev/null || echo 0)
            
            # FAST IPv4: Use ipset restore (always swap, even if empty, to clear stale entries)
            # Ensure target ipset exists for swap
            if ! ipset list "${IPSET_PREFIX}-${name}" >/dev/null 2>&1; then
                ipset create "${IPSET_PREFIX}-${name}" hash:net family inet hashsize 4096 maxelem 200000
            fi
            
            {
                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"
                echo "swap ${IPSET_PREFIX}-${name} ${IPSET_PREFIX}-${name}-tmp"
                echo "destroy ${IPSET_PREFIX}-${name}-tmp"
            } > "$CACHE_DIR/${name}-v4.restore"
            
            ipset restore < "$CACHE_DIR/${name}-v4.restore" 2>/dev/null || {
                log_message "  ⚠ Batch load failed for $name IPv4, using fallback"
                ipset flush "${IPSET_PREFIX}-${name}" 2>/dev/null || true
                while IFS= read -r ip; do
                    [ -z "$ip" ] && continue
                    ipset add "${IPSET_PREFIX}-${name}" "$ip" 2>/dev/null || true
                done < "$v4"
            }
            
            # FAST IPv6: Use ipset restore (always swap, even if empty, to clear stale entries)
            if [ "$ENABLE_IPV6" = true ]; then
                # Ensure target ipset exists for swap
                if ! 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
                fi
                
                {
                    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"
                    echo "swap ${IPSET_PREFIX}-${name}-v6 ${IPSET_PREFIX}-${name}-v6-tmp"
                    echo "destroy ${IPSET_PREFIX}-${name}-v6-tmp"
                } > "$CACHE_DIR/${name}-v6.restore"
                
                ipset restore < "$CACHE_DIR/${name}-v6.restore" 2>/dev/null || {
                    log_message "  ⚠ Batch load failed for $name IPv6, using fallback"
                    ipset flush "${IPSET_PREFIX}-${name}-v6" 2>/dev/null || true
                    while IFS= read -r ip; do
                        [ -z "$ip" ] && continue
                        ipset add "${IPSET_PREFIX}-${name}-v6" "$ip" 2>/dev/null || true
                    done < "$v6"
                }
            fi
            
            log_message "  ✓ $name: $c4 IPv4, $c6 IPv6"
        else
            log_message "  ✗ Failed: $name"
            failed=$((failed + 1))
        fi
    done < "$FEEDS_CONFIG"
    
    # Save ipsets
    ipset save > /etc/sysconfig/ipset 2>/dev/null || ipset save > /etc/iptables/ipsets 2>/dev/null || true
    
    log_message "✓ Updated $total feeds ($failed failed) - FAST IPSET RESTORE MODE"
}

apply_iptables_rules() {
    log_message "Applying per-feed iptables rules..."
    
    # Backup current rules
    iptables-save > "$BACKUP_DIR/iptables-save-$(date +%Y%m%d-%H%M%S).txt" 2>/dev/null || true
    cleanup_old_backups
    
    # Remove old threat feed rules
    iptables -D INPUT -m set --match-set "$WHITELIST_IPSET" src -j ACCEPT 2>/dev/null || true
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        iptables -D INPUT -m set --match-set "${IPSET_PREFIX}-${name}" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT:${name}] " 2>/dev/null || true
        iptables -D INPUT -m set --match-set "${IPSET_PREFIX}-${name}" src -j DROP 2>/dev/null || true
    done < "$FEEDS_CONFIG" 2>/dev/null || true
    
    if [ "$ENABLE_IPV6" = true ]; then
        ip6tables -D INPUT -m set --match-set "$WHITELIST_IPSET_V6" src -j ACCEPT 2>/dev/null || true
        while IFS='|' read -r enabled name url type description; do
            [[ "$enabled" =~ ^#.*$ ]] && continue
            [[ -z "$enabled" ]] && continue
            ip6tables -D INPUT -m set --match-set "${IPSET_PREFIX}-${name}-v6" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT-v6:${name}] " 2>/dev/null || true
            ip6tables -D INPUT -m set --match-set "${IPSET_PREFIX}-${name}-v6" src -j DROP 2>/dev/null || true
        done < "$FEEDS_CONFIG" 2>/dev/null || true
    fi
    
    # Add whitelist rules (highest priority)
    iptables -I INPUT 1 -m set --match-set "$WHITELIST_IPSET" src -j ACCEPT
    [ "$ENABLE_IPV6" = true ] && ip6tables -I INPUT 1 -m set --match-set "$WHITELIST_IPSET_V6" src -j ACCEPT
    
    # Add per-feed rules
    local line=2
    local v6_line=2
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue
        
        # IPv4
        iptables -I INPUT $line -m set --match-set "${IPSET_PREFIX}-${name}" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT:${name}] "
        line=$((line + 1))
        iptables -I INPUT $line -m set --match-set "${IPSET_PREFIX}-${name}" src -j DROP
        line=$((line + 1))
        
        # IPv6
        if [ "$ENABLE_IPV6" = true ]; then
            ip6tables -I INPUT $v6_line -m set --match-set "${IPSET_PREFIX}-${name}-v6" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT-v6:${name}] "
            v6_line=$((v6_line + 1))
            ip6tables -I INPUT $v6_line -m set --match-set "${IPSET_PREFIX}-${name}-v6" src -j DROP
            v6_line=$((v6_line + 1))
        fi
    done < "$FEEDS_CONFIG"
    
    # SSH rate limiting
    if ! iptables -C INPUT -p tcp --dport "$SSH_PORT" -m conntrack --ctstate NEW -m recent --set 2>/dev/null; then
        iptables -I INPUT -p tcp --dport "$SSH_PORT" -m conntrack --ctstate NEW -m recent --set
        iptables -I INPUT -p tcp --dport "$SSH_PORT" -m conntrack --ctstate NEW -m recent --update --seconds 60 --hitcount 4 -j DROP
    fi
    
    # Save rules
    if [ -d /etc/sysconfig ]; then
        iptables-save > /etc/sysconfig/iptables
        [ "$ENABLE_IPV6" = true ] && ip6tables-save > /etc/sysconfig/ip6tables
    elif [ -d /etc/iptables ]; then
        iptables-save > /etc/iptables/rules.v4
        [ "$ENABLE_IPV6" = true ] && ip6tables-save > /etc/iptables/rules.v6
    fi
    
    log_message "✓ iptables rules applied (per-feed)"
}

setup_iptables_persistence() {
    log_message "Setting up iptables persistence..."

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

    cat > /etc/systemd/system/iptables-restore.service <<EOF
[Unit]
Description=Restore iptables rules
Before=network-pre.target
Wants=network-pre.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/bin/bash -c 'ipset restore -f /etc/sysconfig/ipset 2>/dev/null || ipset restore -f /etc/iptables/ipsets 2>/dev/null || true; ${script_path} _ensure-ipsets 2>/dev/null || true'
ExecStart=/bin/bash -c 'iptables-restore /etc/sysconfig/iptables 2>/dev/null || iptables-restore /etc/iptables/rules.v4 2>/dev/null || true'
ExecStart=/bin/bash -c 'ip6tables-restore /etc/sysconfig/ip6tables 2>/dev/null || ip6tables-restore /etc/iptables/rules.v6 2>/dev/null || true'

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    systemctl enable iptables-restore.service 2>/dev/null || true
    log_message "✓ iptables persistence configured"
}

setup_auto_update() {
    [ "$ENABLE_AUTO_UPDATE" = false ] && return
    
    local script
    script=$(readlink -f "$0")
    
    cat > /etc/systemd/system/iptables-threat-feeds-update.service <<EOF
[Unit]
Description=Update iptables threat feeds (per-feed)
After=network-online.target

[Service]
Type=oneshot
ExecStart=$script update
EOF
    
    cat > /etc/systemd/system/iptables-threat-feeds-update.timer <<EOF
[Unit]
Description=Update threat feeds $UPDATE_INTERVAL

[Timer]
OnCalendar=$UPDATE_INTERVAL
Persistent=true

[Install]
WantedBy=timers.target
EOF
    
    systemctl daemon-reload
    systemctl enable --now iptables-threat-feeds-update.timer
}

cmd_show_stats() {
    local journal_1h
    journal_1h=$(journalctl -k --since "1 hour ago" 2>/dev/null | grep '\[THREAT' || true)
    
    echo "Per-Feed Blocking Statistics"
    printf "%-25s %10s %10s %12s\n" "FEED" "IPv4" "IPv6" "BLOCKS(1h)"
    echo "-------------------------------------------------------------------"
    
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue
        
        local v4 v6=0 blocks_v4=0 blocks_v6=0
        v4=$(ipset list "${IPSET_PREFIX}-${name}" 2>/dev/null | sed -n '/^Members:$/,$p' | tail -n +2 | wc -l)
        [ "$ENABLE_IPV6" = true ] && v6=$(ipset list "${IPSET_PREFIX}-${name}-v6" 2>/dev/null | sed -n '/^Members:$/,$p' | tail -n +2 | wc -l)
        
        if [ -n "$journal_1h" ]; then
            blocks_v4=$(printf '%s' "$journal_1h" | grep -c "\[THREAT:${name}\]" 2>/dev/null || true)
            blocks_v6=$(printf '%s' "$journal_1h" | grep -c "\[THREAT-v6:${name}\]" 2>/dev/null || true)
        fi
        local blocks=$(( ${blocks_v4:-0} + ${blocks_v6:-0} ))
        
        printf "%-25s %10s %10s %12s\n" "$name" "${v4:-0}" "${v6:-0}" "$blocks"
    done < "$FEEDS_CONFIG"
}

_print_feed() {
    local enabled="$1" name="$2" _url="$3" _type="$4" description="$5"
    printf "%-10s %-25s %s\n" "$([ "$enabled" = "1" ] && echo "ENABLED" || echo "DISABLED")" "$name" "$description"
}

cmd_list_feeds() {
    printf "%-10s %-25s %s\n" "STATUS" "NAME" "DESC"
    for_each_feed _print_feed
}

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/sysconfig/ipset 2>/dev/null || ipset save > /etc/iptables/ipsets 2>/dev/null || true
}

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/sysconfig/ipset 2>/dev/null || ipset save > /etc/iptables/ipsets 2>/dev/null || true
    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 '^[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 '^[0-9a-fA-F:]' || echo "No entries"
    fi
}

cmd_add_feed() {
    [ -z "$FEED_NAME" ] || [ -z "$FEED_URL" ] && { echo "Usage: $0 add-feed <NAME> <URL>"; exit 1; }
    validate_feed_name "$FEED_NAME" || 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() {
    [ -z "$FEED_NAME" ] && { echo "Usage: $0 remove-feed <NAME>"; exit 1; }

    # Remove rules for this feed first (while ipsets still exist)
    iptables -D INPUT -m set --match-set "${IPSET_PREFIX}-${FEED_NAME}" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT:${FEED_NAME}] " 2>/dev/null || true
    iptables -D INPUT -m set --match-set "${IPSET_PREFIX}-${FEED_NAME}" src -j DROP 2>/dev/null || true
    if [ "$ENABLE_IPV6" = true ]; then
        ip6tables -D INPUT -m set --match-set "${IPSET_PREFIX}-${FEED_NAME}-v6" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT-v6:${FEED_NAME}] " 2>/dev/null || true
        ip6tables -D INPUT -m set --match-set "${IPSET_PREFIX}-${FEED_NAME}-v6" src -j DROP 2>/dev/null || true
    fi

    # Now safe to destroy ipsets
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}" 2>/dev/null || true
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}-v6" 2>/dev/null || true

    # Remove from config last
    sed -i "/|${FEED_NAME}|/d" "$FEEDS_CONFIG"

    log_message "✓ Removed feed: $FEED_NAME"
}

cmd_enable_feed() {
    [ -z "$FEED_NAME" ] && { echo "Usage: $0 enable-feed <NAME>"; exit 1; }
    sed -i "s/^0|${FEED_NAME}|/1|${FEED_NAME}|/" "$FEEDS_CONFIG"
    log_message "✓ Enabled: $FEED_NAME"
    
    # Create ipsets if they don't exist
    if ! ipset list "${IPSET_PREFIX}-${FEED_NAME}" >/dev/null 2>&1; then
        ipset create "${IPSET_PREFIX}-${FEED_NAME}" hash:net family inet hashsize 4096 maxelem 200000
    fi
    if [ "$ENABLE_IPV6" = true ] && ! ipset list "${IPSET_PREFIX}-${FEED_NAME}-v6" >/dev/null 2>&1; then
        ipset create "${IPSET_PREFIX}-${FEED_NAME}-v6" hash:net family inet6 hashsize 4096 maxelem 200000
    fi
    
    log_message "Run 'update' to download IPs, then 'apply-rules' to add firewall rules"
}

cmd_disable_feed() {
    [ -z "$FEED_NAME" ] && { echo "Usage: $0 disable-feed <NAME>"; exit 1; }

    # Remove rules for this feed first (while ipsets still exist)
    iptables -D INPUT -m set --match-set "${IPSET_PREFIX}-${FEED_NAME}" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT:${FEED_NAME}] " 2>/dev/null || true
    iptables -D INPUT -m set --match-set "${IPSET_PREFIX}-${FEED_NAME}" src -j DROP 2>/dev/null || true
    if [ "$ENABLE_IPV6" = true ]; then
        ip6tables -D INPUT -m set --match-set "${IPSET_PREFIX}-${FEED_NAME}-v6" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT-v6:${FEED_NAME}] " 2>/dev/null || true
        ip6tables -D INPUT -m set --match-set "${IPSET_PREFIX}-${FEED_NAME}-v6" src -j DROP 2>/dev/null || true
    fi

    # Now safe to destroy ipsets
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}" 2>/dev/null || true
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}-v6" 2>/dev/null || true

    # Mark disabled in config
    sed -i "s/^1|${FEED_NAME}|/0|${FEED_NAME}|/" "$FEEDS_CONFIG"

    log_message "✓ Disabled: $FEED_NAME"
}

cmd_install() {
    log_message "Installing per-feed mode..."
    check_requirements
    create_directory_structure
    initialize_feeds_config
    setup_ipsets
    update_feeds
    apply_iptables_rules
    setup_iptables_persistence
    setup_auto_update
    
    echo ""
    echo "=========================================="
    echo "✓ Per-feed installation complete"
    echo "=========================================="
    echo "Feeds: $(grep -c '^1|' "$FEEDS_CONFIG")"
    echo "Config: $FEEDS_CONFIG"
    echo "Log: $LOG_FILE"
    echo ""
    echo "Commands:"
    echo "  $0 show-stats"
    echo "  $0 list-feeds"
    echo "  $0 update"
    echo "  $0 whitelist-add <IP>"
    echo "=========================================="
}

cmd_test_rules() {
    log_message "Testing iptables rule generation (dry-run mode)..."
    
    echo "=========================================="
    echo "Rule Generation Test"
    echo "=========================================="
    echo ""
    
    # Count enabled feeds
    local enabled_count=0
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue
        enabled_count=$((enabled_count + 1))
    done < "$FEEDS_CONFIG"
    
    echo "✓ Found $enabled_count enabled feeds"
    echo ""
    
    # Show what would be generated
    echo "IPv4 rules that would be created:"
    echo "  1. Whitelist bypass: -I INPUT 1 -m set --match-set $WHITELIST_IPSET src -j ACCEPT"
    
    local line=2
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue
        
        echo "  $line. [${name}] LOG: -I INPUT $line -m set --match-set ${IPSET_PREFIX}-${name} src -m limit --limit 5/min -j LOG"
        line=$((line + 1))
        echo "  $line. [${name}] DROP: -I INPUT $line -m set --match-set ${IPSET_PREFIX}-${name} src -j DROP"
        line=$((line + 1))
    done < "$FEEDS_CONFIG"
    
    echo ""
    echo "Total IPv4 rules: $((line - 1))"
    
    if [ "$ENABLE_IPV6" = true ]; then
        echo ""
        echo "IPv6 rules that would be created:"
        echo "  1. Whitelist bypass: -I INPUT 1 -m set --match-set $WHITELIST_IPSET_V6 src -j ACCEPT"
        
        local v6_line=2
        while IFS='|' read -r enabled name url type description; do
            [[ "$enabled" =~ ^#.*$ ]] && continue
            [[ -z "$enabled" ]] && continue
            [ "$enabled" != "1" ] && continue
            
            echo "  $v6_line. [${name}] LOG: -I INPUT $v6_line -m set --match-set ${IPSET_PREFIX}-${name}-v6 src -j LOG"
            v6_line=$((v6_line + 1))
            echo "  $v6_line. [${name}] DROP: -I INPUT $v6_line -m set --match-set ${IPSET_PREFIX}-${name}-v6 src -j DROP"
            v6_line=$((v6_line + 1))
        done < "$FEEDS_CONFIG"
        
        echo ""
        echo "Total IPv6 rules: $((v6_line - 1))"
    fi
    
    echo ""
    echo "=========================================="
    echo "✓ Test passed - rules would be generated successfully"
    echo "  To apply these rules, run: $0 apply-rules"
    echo "=========================================="
}

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 "$CACHE_DIR"/*-v4.restore "$CACHE_DIR"/*-v6.restore; 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}"
        feed_name="${feed_name%%-v4.restore}"
        feed_name="${feed_name%%-v6.restore}"

        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"
}

main() {
    # Internal command used by iptables-restore.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
            create_directory_structure
            update_feeds
            # DO NOT apply rules here - only update ipsets
            # To regenerate rules, use: apply-rules, enable-feed, disable-feed, or remove-feed
            ;;
        apply-rules)
            check_requirements
            apply_iptables_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 ;;
    esac
}

main "$@"