#!/bin/bash
################################################################################
# Script Name: iptables-blocklists.sh
# Version: 1.0
# Description: Per-feed iptables threat intelligence blocking with ipset
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
################################################################################
# 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

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() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

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

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() {
    [ -f "$FEEDS_CONFIG" ] && return
    
    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
# TYPE: Format type (plain, cidr, commented, custom)
# 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"
}

setup_ipsets() {
    log_message "Setting up per-feed ipsets..."
    
    # Whitelist
    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
    
    # Create ipset per feed
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" != "1" ] && continue
        
        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
        
        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
        fi
    done < "$FEEDS_CONFIG"
}

download_feed() {
    curl -f -s -m 30 -L "$1" -o "$2" 2>/dev/null
}

parse_feed() {
    local file="$1" type="$2" out_v4="$3" out_v6="$4"
    true > "$out_v4"
    true > "$out_v6"
    
    case "$type" in
        plain)
            grep -E '^[0-9.]+(/[0-9]+)?$' "$file" >> "$out_v4" 2>/dev/null || true
            [ "$ENABLE_IPV6" = true ] && grep -E '^[0-9a-fA-F:]+(/[0-9]+)?$' "$file" | grep ':' >> "$out_v6" 2>/dev/null || true
            ;;
        cidr)
            grep -E '^[0-9.]+' "$file" | cut -d' ' -f1 | cut -d'#' -f1 | grep -v '^$' >> "$out_v4" 2>/dev/null || true
            [ "$ENABLE_IPV6" = true ] && grep -E '^[0-9a-fA-F:]+' "$file" | grep ':' | cut -d' ' -f1 | cut -d'#' -f1 >> "$out_v6" 2>/dev/null || true
            ;;
        commented)
            grep -v -E '^[#;]|^$' "$file" | grep -oE '[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?' >> "$out_v4" 2>/dev/null || true
            [ "$ENABLE_IPV6" = true ] && grep -v -E '^[#;]|^$' "$file" | grep -oE '[0-9a-fA-F:]+(/[0-9]+)?' | grep -E '^[0-9a-fA-F]{1,4}:[0-9a-fA-F:]+' >> "$out_v6" 2>/dev/null || true
            ;;
    esac
}

update_feeds() {
    log_message "Starting per-feed update (FAST ipset restore mode)..."
    
    # Auto-cleanup cache and ipsets for disabled feeds
    local enabled_feeds=$(grep '^1|' "$FEEDS_CONFIG" 2>/dev/null | cut -d'|' -f2)
    local cleaned_cache=0
    local cleaned_ipsets=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=$(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
    
    # Clean ipsets for disabled feeds
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        [ "$enabled" = "1" ] && continue
        
        if ipset list "${IPSET_PREFIX}-${name}" >/dev/null 2>&1; then
            ipset destroy "${IPSET_PREFIX}-${name}" 2>/dev/null && cleaned_ipsets=$((cleaned_ipsets + 1))
        fi
        if ipset list "${IPSET_PREFIX}-${name}-v6" >/dev/null 2>&1; then
            ipset destroy "${IPSET_PREFIX}-${name}-v6" 2>/dev/null && cleaned_ipsets=$((cleaned_ipsets + 1))
        fi
    done < "$FEEDS_CONFIG"
    
    [ "$cleaned_cache" -gt 0 ] && log_message "  Cleaned $cleaned_cache stale cache files"
    [ "$cleaned_ipsets" -gt 0 ] && log_message "  Destroyed $cleaned_ipsets stale ipsets"
    
    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
            if [ "$c4" -gt 0 ]; then
                # 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"
                }
            fi
            
            # FAST IPv6: Use ipset restore
            if [ "$ENABLE_IPV6" = true ] && [ "$c6" -gt 0 ]; 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
    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 -A INPUT -m set --match-set "${IPSET_PREFIX}-${name}-v6" src -m limit --limit 5/min -j LOG --log-prefix "[THREAT-v6:${name}] "
            ip6tables -A INPUT -m set --match-set "${IPSET_PREFIX}-${name}-v6" src -j DROP
        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..."
    
    # Create systemd service for iptables restore
    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'
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=$(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() {
    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=$(ipset list "${IPSET_PREFIX}-${name}" 2>/dev/null | grep -c '^[0-9.]' || echo 0)
        [ "$ENABLE_IPV6" = true ] && v6=$(ipset list "${IPSET_PREFIX}-${name}-v6" 2>/dev/null | grep -c '^[0-9a-fA-F:]' || echo 0)
        blocks=$(journalctl -k --since "1 hour ago" 2>/dev/null | grep -c "\[THREAT:${name}\]" || echo 0)
        
        printf "%-25s %10s %10s %12s\n" "$name" "$v4" "$v6" "$blocks"
    done < "$FEEDS_CONFIG"
}

cmd_list_feeds() {
    printf "%-10s %-25s %s\n" "STATUS" "NAME" "DESC"
    while IFS='|' read -r enabled name url type description; do
        [[ "$enabled" =~ ^#.*$ ]] && continue
        [[ -z "$enabled" ]] && continue
        printf "%-10s %-25s %s\n" "$([ "$enabled" = "1" ] && echo "ENABLED" || echo "DISABLED")" "$name" "$description"
    done < "$FEEDS_CONFIG"
}

cmd_whitelist_add() {
    [ -z "$WHITELIST_IP" ] && { echo "Usage: $0 whitelist-add <IP|CIDR>"; exit 1; }
    
    if echo "$WHITELIST_IP" | grep -q ':'; then
        ipset add "$WHITELIST_IPSET_V6" "$WHITELIST_IP" 2>/dev/null && \
            log_message "✓ Added to IPv6 whitelist: $WHITELIST_IP" || \
            { echo "Failed to add $WHITELIST_IP"; exit 1; }
    else
        ipset add "$WHITELIST_IPSET" "$WHITELIST_IP" 2>/dev/null && \
            log_message "✓ Added to IPv4 whitelist: $WHITELIST_IP" || \
            { echo "Failed to add $WHITELIST_IP"; exit 1; }
    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; }
    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; }
    sed -i "/|${FEED_NAME}|/d" "$FEEDS_CONFIG"
    
    # Remove ipsets and rules
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}" 2>/dev/null || true
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}-v6" 2>/dev/null || true
    
    log_message "✓ Removed feed: $FEED_NAME"
    log_message "Reapplying rules..."
    apply_iptables_rules
}

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; }
    sed -i "s/^1|${FEED_NAME}|/0|${FEED_NAME}|/" "$FEEDS_CONFIG"
    
    # Destroy ipsets to clear metrics
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}" 2>/dev/null || true
    ipset destroy "${IPSET_PREFIX}-${FEED_NAME}-v6" 2>/dev/null || true
    
    log_message "✓ Disabled: $FEED_NAME"
    log_message "Reapplying rules..."
    apply_iptables_rules
}

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_count=0
        while IFS='|' read -r enabled name url type description; do
            [[ "$enabled" =~ ^#.*$ ]] && continue
            [[ -z "$enabled" ]] && continue
            [ "$enabled" != "1" ] && continue
            
            v6_count=$((v6_count + 1))
            echo "  $((v6_count * 2)). [${name}] LOG: -A INPUT -m set --match-set ${IPSET_PREFIX}-${name}-v6 src -j LOG"
            echo "  $((v6_count * 2 + 1)). [${name}] DROP: -A INPUT -m set --match-set ${IPSET_PREFIX}-${name}-v6 src -j DROP"
        done < "$FEEDS_CONFIG"
        
        echo ""
        echo "Total IPv6 rules: $((v6_count * 2 + 1))"
    fi
    
    echo ""
    echo "=========================================="
    echo "✓ Test passed - rules would be generated successfully"
    echo "  To apply these rules, run: $0 apply-rules"
    echo "=========================================="
}

main() {
    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 ;;
    esac
}

main "$@"