#!/bin/bash

################################################################################
# IP Reputation Management Library
################################################################################
# Purpose: Centralized IP reputation tracking across all toolkit scripts
# Features:
# - Fast lookups using indexed file structure
# - Tracks: hits, country, last seen, reputation score, attack types
# - Optimized for high-volume traffic (attacks with thousands of IPs)
# - Automatic cleanup of old entries
# - GeoIP integration
# - Shared across all monitoring/analysis scripts
################################################################################

# Database location
IP_REP_DB_DIR="${IP_REP_DB_DIR:-/var/lib/server-toolkit/ip-reputation}"
IP_REP_DB="$IP_REP_DB_DIR/ip_database.db"
IP_REP_INDEX="$IP_REP_DB_DIR/ip_index.idx"
IP_REP_LOCK="$IP_REP_DB_DIR/.db.lock"

# Reputation score thresholds
REP_SCORE_CRITICAL=80    # Definitely malicious
REP_SCORE_HIGH=60        # Likely malicious
REP_SCORE_MEDIUM=40      # Suspicious
REP_SCORE_LOW=20         # Borderline
REP_SCORE_SAFE=0         # Safe/legitimate

# Attack type flags (bitmask for efficient storage)
ATTACK_FLAG_SQL_INJECTION=1
ATTACK_FLAG_XSS=2
ATTACK_FLAG_PATH_TRAVERSAL=4
ATTACK_FLAG_RCE=8
ATTACK_FLAG_BRUTEFORCE=16
ATTACK_FLAG_DDOS=32
ATTACK_FLAG_BOT=64
ATTACK_FLAG_SCANNER=128
ATTACK_FLAG_EXPLOIT=256

# Initialize the IP reputation database
init_ip_reputation_db() {
    mkdir -p "$IP_REP_DB_DIR" 2>/dev/null

    # Create empty database if it doesn't exist
    if [ ! -f "$IP_REP_DB" ]; then
        touch "$IP_REP_DB"
        chmod 600 "$IP_REP_DB"
    fi

    if [ ! -f "$IP_REP_INDEX" ]; then
        touch "$IP_REP_INDEX"
        chmod 600 "$IP_REP_INDEX"
    fi

    return 0
}

# Database format (pipe-delimited for fast parsing):
# IP|HIT_COUNT|REPUTATION_SCORE|COUNTRY|ATTACK_FLAGS|FIRST_SEEN|LAST_SEEN|LAST_ACTIVITY|NOTES|BAN_COUNT|LAST_BAN
# Example:
# 192.168.1.100|523|75|US|193|1730000000|1730800000|SQL injection on /admin|Auto-flagged|3|1730900000

# Lock management for concurrent access
acquire_lock() {
    local timeout=10
    local elapsed=0

    while [ -f "$IP_REP_LOCK" ] && [ $elapsed -lt $timeout ]; do
        sleep 0.1
        elapsed=$((elapsed + 1))
    done

    if [ $elapsed -ge $timeout ]; then
        # Stale lock, remove it
        rm -f "$IP_REP_LOCK" 2>/dev/null
    fi

    touch "$IP_REP_LOCK"
}

release_lock() {
    rm -f "$IP_REP_LOCK" 2>/dev/null
}

# Fast IP lookup using hash-based index for O(1) lookups
# Returns: IP data if found, empty if not found
lookup_ip() {
    local ip="$1"

    [ -z "$ip" ] && return 1
    [ ! -f "$IP_REP_DB" ] && return 1

    # Calculate hash bucket (first octet for IPv4 distributes IPs across 256 buckets)
    local hash_bucket="${ip%%.*}"
    local hash_file="${IP_REP_DB_DIR}/hash_${hash_bucket}.idx"

    # Fast path: Check hash bucket first (much smaller file to grep)
    if [ -f "$hash_file" ]; then
        # Hash bucket contains line numbers for IPs in this bucket
        local line_num=$(grep -m 1 "^${ip}|" "$hash_file" 2>/dev/null | cut -d'|' -f2)
        if [ -n "$line_num" ]; then
            # Direct line access - O(1) lookup!
            sed -n "${line_num}p" "$IP_REP_DB" 2>/dev/null
            return 0
        fi
    fi

    # Fallback: Linear search (for IPs not yet indexed)
    # Use tac to read file backwards, then grep for first match
    # This ensures we get the LATEST entry for IPs with duplicates
    tac "$IP_REP_DB" 2>/dev/null | grep -m 1 "^${ip}|" 2>/dev/null
}

# Add or update IP in database
# Usage: update_ip_reputation IP [HIT_INCREMENT] [SCORE_DELTA] [ATTACK_FLAGS] [ACTIVITY_NOTE]
update_ip_reputation() {
    local ip="$1"
    local hit_increment="${2:-1}"
    local score_delta="${3:-0}"
    local new_attack_flags="${4:-0}"
    local activity_note="${5:-}"

    [ -z "$ip" ] && return 1

    init_ip_reputation_db
    acquire_lock

    local existing
    existing=$(lookup_ip "$ip")

    local current_time=$(date +%s)

    if [ -n "$existing" ]; then
        # Parse existing entry
        IFS='|' read -r old_ip hit_count rep_score country attack_flags first_seen last_seen last_activity notes <<< "$existing"

        # Update values
        hit_count=$((hit_count + hit_increment))
        rep_score=$((rep_score + score_delta))

        # Cap reputation score at 0-100
        [ $rep_score -lt 0 ] && rep_score=0
        [ $rep_score -gt 100 ] && rep_score=100

        # Merge attack flags (bitwise OR)
        attack_flags=$((attack_flags | new_attack_flags))

        last_seen="$current_time"

        # Update activity note if provided
        if [ -n "$activity_note" ]; then
            last_activity="$activity_note"
        fi

        # OPTIMIZATION: Append-only writes (much faster than sed -i delete)
        # Append updated entry to end of file
        echo "$ip|$hit_count|$rep_score|$country|$attack_flags|$first_seen|$last_seen|$last_activity|$notes" >> "$IP_REP_DB"

        # Mark for compaction (file will have duplicates until compact_database runs)
        touch "${IP_REP_DB}.needs_compact" 2>/dev/null
    else
        # New entry
        local country=$(get_ip_country "$ip")
        echo "$ip|$hit_increment|$score_delta|$country|$new_attack_flags|$current_time|$current_time|$activity_note|" >> "$IP_REP_DB"
    fi

    release_lock

    # Auto-compact if file has lots of duplicates (from append-only writes)
    # Check if compaction is needed (marked file exists)
    if [ -f "${IP_REP_DB}.needs_compact" ]; then
        local db_size=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo "0")

        # Compact if database >50k lines (likely has significant duplicates)
        # Use random check to avoid all processes compacting simultaneously
        if [ "$db_size" -gt 50000 ] && [ $((RANDOM % 200)) -eq 0 ]; then
            compact_database &  # Background process (includes rebuild_index)
            return 0
        fi
    fi

    # Rebuild index automatically when database grows significantly
    # Check if hash index exists and is fresh
    local db_size=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo "0")
    local hash_count=$(ls -1 "${IP_REP_DB_DIR}"/hash_*.idx 2>/dev/null | wc -l)

    # Rebuild if:
    # 1. Database has >10k IPs but no hash index exists
    # 2. Database has >100k IPs and 1% chance (frequent enough during attacks)
    if [ "$hash_count" -eq 0 ] && [ "$db_size" -gt 10000 ]; then
        rebuild_index &  # Background process
    elif [ "$db_size" -gt 100000 ] && [ $((RANDOM % 100)) -eq 0 ]; then
        rebuild_index &  # Background process
    fi

    return 0
}

# Get IP country using multiple methods
get_ip_country() {
    local ip="$1"
    local country="??"

    # Method 1: Check if geoiplookup is available
    if command -v geoiplookup >/dev/null 2>&1; then
        country=$(geoiplookup "$ip" 2>/dev/null | grep -oP 'Country Edition: \K[A-Z]{2}' | head -1)
    fi

    # Method 2: Check if geoiplookup6 for IPv6
    if [ -z "$country" ] || [ "$country" = "??" ]; then
        if command -v geoiplookup6 >/dev/null 2>&1 && [[ "$ip" =~ : ]]; then
            country=$(geoiplookup6 "$ip" 2>/dev/null | grep -oP 'Country Edition: \K[A-Z]{2}' | head -1)
        fi
    fi

    # Method 3: Check /usr/share/GeoIP databases directly
    if [ -z "$country" ] || [ "$country" = "??" ]; then
        if [ -f "/usr/share/GeoIP/GeoIP.dat" ] && command -v geoiplookup >/dev/null 2>&1; then
            country=$(geoiplookup "$ip" 2>/dev/null | awk -F': ' '{print $2}' | cut -d',' -f1 | head -1)
        fi
    fi

    # Method 4: Fallback - use whois (slower, only if critically needed)
    # Disabled by default for performance
    # if [ -z "$country" ] || [ "$country" = "??" ]; then
    #     country=$(whois "$ip" 2>/dev/null | grep -iE "^country:" | head -1 | awk '{print $2}')
    # fi

    # Default if all methods fail
    [ -z "$country" ] && country="??"

    echo "$country"
}

# Increment IP hit count (fast path for common case)
increment_ip_hits() {
    local ip="$1"
    local increment="${2:-1}"

    update_ip_reputation "$ip" "$increment" 0 0 ""
}

# Flag IP for specific attack type
flag_ip_attack() {
    local ip="$1"
    local attack_type="$2"
    local score_increase="${3:-5}"
    local note="${4:-$attack_type}"

    local attack_flag=0

    case "$attack_type" in
        SQL_INJECTION|sql) attack_flag=$ATTACK_FLAG_SQL_INJECTION; score_increase=15 ;;
        XSS|xss) attack_flag=$ATTACK_FLAG_XSS; score_increase=10 ;;
        PATH_TRAVERSAL|path) attack_flag=$ATTACK_FLAG_PATH_TRAVERSAL; score_increase=12 ;;
        RCE|rce|shell) attack_flag=$ATTACK_FLAG_RCE; score_increase=20 ;;
        BRUTEFORCE|brute) attack_flag=$ATTACK_FLAG_BRUTEFORCE; score_increase=8 ;;
        DDOS|ddos) attack_flag=$ATTACK_FLAG_DDOS; score_increase=10 ;;
        BOT|bot) attack_flag=$ATTACK_FLAG_BOT; score_increase=3 ;;
        SCANNER|scan) attack_flag=$ATTACK_FLAG_SCANNER; score_increase=5 ;;
        EXPLOIT|exploit) attack_flag=$ATTACK_FLAG_EXPLOIT; score_increase=15 ;;
        *) attack_flag=0; score_increase=5 ;;
    esac

    update_ip_reputation "$ip" 1 "$score_increase" "$attack_flag" "$note"
}

# Mark IP as legitimate (reduces reputation score)
mark_ip_legitimate() {
    local ip="$1"
    local note="${2:-Marked as legitimate}"

    update_ip_reputation "$ip" 0 -20 0 "$note"
}

# Get IP reputation category
get_ip_reputation_category() {
    local score="$1"

    if [ $score -ge $REP_SCORE_CRITICAL ]; then
        echo "CRITICAL"
    elif [ $score -ge $REP_SCORE_HIGH ]; then
        echo "HIGH"
    elif [ $score -ge $REP_SCORE_MEDIUM ]; then
        echo "MEDIUM"
    elif [ $score -ge $REP_SCORE_LOW ]; then
        echo "LOW"
    else
        echo "SAFE"
    fi
}

# Get attack types from flags
decode_attack_flags() {
    local flags="$1"
    local attacks=""

    [ $((flags & ATTACK_FLAG_SQL_INJECTION)) -ne 0 ] && attacks="${attacks}SQL,"
    [ $((flags & ATTACK_FLAG_XSS)) -ne 0 ] && attacks="${attacks}XSS,"
    [ $((flags & ATTACK_FLAG_PATH_TRAVERSAL)) -ne 0 ] && attacks="${attacks}PATH,"
    [ $((flags & ATTACK_FLAG_RCE)) -ne 0 ] && attacks="${attacks}RCE,"
    [ $((flags & ATTACK_FLAG_BRUTEFORCE)) -ne 0 ] && attacks="${attacks}BRUTE,"
    [ $((flags & ATTACK_FLAG_DDOS)) -ne 0 ] && attacks="${attacks}DDOS,"
    [ $((flags & ATTACK_FLAG_BOT)) -ne 0 ] && attacks="${attacks}BOT,"
    [ $((flags & ATTACK_FLAG_SCANNER)) -ne 0 ] && attacks="${attacks}SCAN,"
    [ $((flags & ATTACK_FLAG_EXPLOIT)) -ne 0 ] && attacks="${attacks}EXPLOIT,"

    # Remove trailing comma
    attacks="${attacks%,}"

    [ -z "$attacks" ] && attacks="NONE"

    echo "$attacks"
}

# Query and display IP information
query_ip_reputation() {
    local ip="$1"

    init_ip_reputation_db

    local data
    data=$(lookup_ip "$ip")

    if [ -z "$data" ]; then
        echo "IP $ip not found in reputation database"
        return 1
    fi

    IFS='|' read -r ip hit_count rep_score country attack_flags first_seen last_seen last_activity notes <<< "$data"

    local category=$(get_ip_reputation_category "$rep_score")
    local attacks=$(decode_attack_flags "$attack_flags")
    local first_seen_date=$(date -d "@$first_seen" '+%Y-%m-%d %H:%M:%S' 2>/dev/null || echo "$first_seen")
    local last_seen_date=$(date -d "@$last_seen" '+%Y-%m-%d %H:%M:%S' 2>/dev/null || echo "$last_seen")

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "IP Address:       $ip"
    echo "Country:          $country"
    echo "Reputation:       $rep_score/100 [$category]"
    echo "Total Hits:       $hit_count"
    echo "Attack Types:     $attacks"
    echo "First Seen:       $first_seen_date"
    echo "Last Seen:        $last_seen_date"
    echo "Last Activity:    ${last_activity:-None recorded}"
    echo "Notes:            ${notes:-None}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    return 0
}

# Get top IPs by reputation score
get_top_malicious_ips() {
    local limit="${1:-20}"

    init_ip_reputation_db

    [ ! -f "$IP_REP_DB" ] && return 1

    # OPTIMIZATION: For large files, use partial sort (much faster)
    # Only sort enough to find top N instead of sorting entire file
    local db_size=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo "0")

    if [ "$db_size" -gt 100000 ]; then
        # For very large databases, use awk to find high-scoring IPs first
        # then sort only those (much faster than sorting 500k lines)
        awk -F'|' '$3 >= 50' "$IP_REP_DB" | sort -t'|' -k3 -rn | head -n "$limit"
    else
        # For smaller databases, regular sort is fine
        sort -t'|' -k3 -rn "$IP_REP_DB" | head -n "$limit"
    fi
}

# Get top IPs by hit count
get_top_active_ips() {
    local limit="${1:-20}"

    init_ip_reputation_db

    [ ! -f "$IP_REP_DB" ] && return 1

    # OPTIMIZATION: For large files, filter first then sort
    local db_size=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo "0")

    if [ "$db_size" -gt 100000 ]; then
        # Filter to IPs with >100 hits, then sort (much faster)
        awk -F'|' '$2 >= 100' "$IP_REP_DB" | sort -t'|' -k2 -rn | head -n "$limit"
    else
        # For smaller databases, regular sort is fine
        sort -t'|' -k2 -rn "$IP_REP_DB" | head -n "$limit"
    fi
}

# Clean up old entries (not seen in X days)
cleanup_old_ips() {
    local days_old="${1:-90}"

    init_ip_reputation_db
    acquire_lock

    local cutoff_time=$(($(date +%s) - (days_old * 86400)))
    local temp_file="${IP_REP_DB}.tmp"

    # Keep only IPs seen within the cutoff time
    awk -F'|' -v cutoff="$cutoff_time" '$7 >= cutoff' "$IP_REP_DB" > "$temp_file"

    mv "$temp_file" "$IP_REP_DB"

    release_lock

    echo "Cleaned up IPs not seen in $days_old days"
}

# Compact database to remove duplicate IP entries (from append-only writes)
compact_database() {
    init_ip_reputation_db
    acquire_lock

    echo "Compacting database (removing duplicate IP entries)..."

    local temp_db="${IP_REP_DB}.compact_tmp"
    local original_size=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo "0")

    # Use awk to keep only the LAST occurrence of each IP (most recent data)
    # Read file backwards, keep first occurrence of each IP, then reverse again
    tac "$IP_REP_DB" | awk -F'|' '!seen[$1]++' | tac > "$temp_db"

    # Replace original with compacted version
    mv "$temp_db" "$IP_REP_DB"

    local new_size=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo "0")
    local removed=$((original_size - new_size))

    # Remove compaction marker
    rm -f "${IP_REP_DB}.needs_compact" 2>/dev/null

    release_lock

    echo "Compaction complete: Removed $removed duplicate entries ($original_size → $new_size IPs)"

    # Rebuild index after compaction
    rebuild_index
}

# Rebuild index for faster lookups (for very large databases)
rebuild_index() {
    init_ip_reputation_db
    acquire_lock

    echo "Rebuilding hash-based index for fast lookups..."

    # Remove old hash files
    rm -f "${IP_REP_DB_DIR}"/hash_*.idx 2>/dev/null

    # Build hash buckets (256 buckets based on first octet)
    # This distributes 500k IPs into ~2k IPs per bucket = MUCH faster
    local line_num=0
    while IFS='|' read -r ip rest; do
        ((line_num++))

        # Calculate hash bucket from first octet
        local hash_bucket="${ip%%.*}"
        local hash_file="${IP_REP_DB_DIR}/hash_${hash_bucket}.idx"

        # Store IP and its line number in the hash bucket file
        echo "${ip}|${line_num}" >> "$hash_file"
    done < "$IP_REP_DB"

    # Sort each hash bucket file for faster grep
    for hash_file in "${IP_REP_DB_DIR}"/hash_*.idx; do
        [ -f "$hash_file" ] && sort -t'|' -k1 -o "$hash_file" "$hash_file"
    done

    # Also create main sorted index for compatibility
    sort -t'|' -k1 "$IP_REP_DB" > "$IP_REP_INDEX"

    release_lock

    echo "Index rebuilt: $(ls -1 "${IP_REP_DB_DIR}"/hash_*.idx 2>/dev/null | wc -l) hash buckets created"
}

# Export reputation database to readable format
export_ip_reputation() {
    local output_file="${1:-/tmp/ip_reputation_export_$(date +%Y%m%d_%H%M%S).txt}"

    init_ip_reputation_db

    {
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        echo "SERVER TOOLKIT - IP REPUTATION DATABASE EXPORT"
        echo "Generated: $(date '+%Y-%m-%d %H:%M:%S')"
        echo "Total IPs: $(wc -l < "$IP_REP_DB" 2>/dev/null || echo 0)"
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        echo ""

        printf "%-15s | %-7s | %-4s | %-8s | %-6s | %-30s | %-19s\n" \
            "IP ADDRESS" "HITS" "CTRY" "REP" "LEVEL" "ATTACKS" "LAST SEEN"
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

        # Sort by reputation score, descending
        sort -t'|' -k3 -rn "$IP_REP_DB" | while IFS='|' read -r ip hit_count rep_score country attack_flags first_seen last_seen last_activity notes; do
            local category=$(get_ip_reputation_category "$rep_score")
            local attacks=$(decode_attack_flags "$attack_flags")
            local last_seen_date=$(date -d "@$last_seen" '+%Y-%m-%d %H:%M' 2>/dev/null || echo "$last_seen")

            printf "%-15s | %-7s | %-4s | %-3s/100 | %-8s | %-30s | %-19s\n" \
                "$ip" "$hit_count" "$country" "$rep_score" "$category" "${attacks:0:30}" "$last_seen_date"
        done

        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    } > "$output_file"

    echo "IP reputation database exported to: $output_file"
}

# Check if IP should be blocked (based on reputation)
should_block_ip() {
    local ip="$1"
    local threshold="${2:-$REP_SCORE_HIGH}"  # Default: block if reputation >= 60

    local data
    data=$(lookup_ip "$ip")

    [ -z "$data" ] && return 1  # Unknown IP, don't block

    IFS='|' read -r _ _ rep_score _ _ _ _ _ _ <<< "$data"

    [ $rep_score -ge $threshold ] && return 0  # Should block
    return 1  # Should not block
}

# Batch import IPs from various sources
import_ips_from_log() {
    local log_file="$1"
    local attack_type="${2:-SUSPICIOUS}"
    local score_per_hit="${3:-5}"

    [ ! -f "$log_file" ] && return 1

    # Extract IPs and count occurrences
    grep -oE '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}' "$log_file" | \
    sort | uniq -c | while read count ip; do
        update_ip_reputation "$ip" "$count" "$score_per_hit" 0 "Imported from $log_file"
    done

    echo "Imported IPs from $log_file"
}

# Statistics summary
show_ip_statistics() {
    init_ip_reputation_db

    local total_ips=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo 0)
    local critical=$(awk -F'|' -v thresh=$REP_SCORE_CRITICAL '$3 >= thresh' "$IP_REP_DB" 2>/dev/null | wc -l)
    local high=$(awk -F'|' -v low=$REP_SCORE_HIGH -v hi=$REP_SCORE_CRITICAL '$3 >= low && $3 < hi' "$IP_REP_DB" 2>/dev/null | wc -l)
    local medium=$(awk -F'|' -v low=$REP_SCORE_MEDIUM -v hi=$REP_SCORE_HIGH '$3 >= low && $3 < hi' "$IP_REP_DB" 2>/dev/null | wc -l)
    local low=$(awk -F'|' -v low=$REP_SCORE_LOW -v hi=$REP_SCORE_MEDIUM '$3 >= low && $3 < hi' "$IP_REP_DB" 2>/dev/null | wc -l)
    local safe=$(awk -F'|' -v thresh=$REP_SCORE_LOW '$3 < thresh' "$IP_REP_DB" 2>/dev/null | wc -l)

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "IP REPUTATION DATABASE STATISTICS"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "Total Tracked IPs:    $total_ips"
    echo ""
    echo "By Reputation Level:"
    echo "  CRITICAL (≥80):     $critical"
    echo "  HIGH (60-79):       $high"
    echo "  MEDIUM (40-59):     $medium"
    echo "  LOW (20-39):        $low"
    echo "  SAFE (<20):         $safe"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
}

################################################################################
# BAN MANAGEMENT & TRACKING
################################################################################

# Record that an IP was banned
# Usage: record_ip_ban IP DURATION_HOURS [REASON]
record_ip_ban() {
    local ip="$1"
    local duration="${2:-1}"
    local reason="${3:-Manual ban from live monitor}"

    [ -z "$ip" ] && return 1

    init_ip_reputation_db
    acquire_lock

    local existing
    existing=$(lookup_ip "$ip")

    local current_time=$(date +%s)

    if [ -n "$existing" ]; then
        # Parse existing entry (with new ban fields)
        IFS='|' read -r old_ip hit_count rep_score country attack_flags first_seen last_seen last_activity notes ban_count last_ban <<< "$existing"

        # Increment ban count
        ban_count=$((${ban_count:-0} + 1))
        last_ban="$current_time"

        # Increase reputation score for being banned
        rep_score=$((rep_score + 10))
        [ $rep_score -gt 100 ] && rep_score=100

        # Update notes
        notes="Banned ${ban_count}x (${duration}h): $reason"

        # Write updated entry (remove old, add new)
        local temp_file="${IP_REP_DB}.tmp.$$"
        grep -v "^${ip}|" "$IP_REP_DB" > "$temp_file" 2>/dev/null || touch "$temp_file"
        echo "$ip|$hit_count|$rep_score|$country|$attack_flags|$first_seen|$last_seen|$last_activity|$notes|$ban_count|$last_ban" >> "$temp_file"
        mv "$temp_file" "$IP_REP_DB"
    else
        # New IP - create entry with ban
        echo "$ip|0|70|unknown|0|$current_time|$current_time|Banned|Banned: $reason|1|$current_time" >> "$IP_REP_DB"
    fi

    release_lock
    return 0
}

# Get ban count for an IP
get_ip_ban_count() {
    local ip="$1"

    local data
    data=$(lookup_ip "$ip")

    [ -z "$data" ] && echo "0" && return 0

    # Extract ban_count (field 10)
    echo "$data" | awk -F'|' '{print $10}'
}

# Get last ban timestamp for an IP
get_ip_last_ban() {
    local ip="$1"

    local data
    data=$(lookup_ip "$ip")

    [ -z "$data" ] && echo "0" && return 0

    # Extract last_ban (field 11)
    echo "$data" | awk -F'|' '{print $11}'
}

# Block IP using CSF (if available) or iptables
# Usage: block_ip_temporary IP DURATION_HOURS [REASON]
block_ip_temporary() {
    local ip="$1"
    local duration="${2:-1}"  # Default: 1 hour
    local reason="${3:-High threat activity detected}"

    [ -z "$ip" ] && return 1

    # Validate IP format
    if ! [[ "$ip" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        echo "ERROR: Invalid IP format: $ip"
        return 1
    fi

    # Check if CSF is available
    if command -v csf &>/dev/null; then
        # Use CSF temporary deny
        local duration_seconds=$((duration * 3600))
        csf -td "$ip" "$duration_seconds" "$reason" &>/dev/null

        if [ $? -eq 0 ]; then
            echo "✓ Blocked $ip using CSF for ${duration}h: $reason"
            record_ip_ban "$ip" "$duration" "$reason"
            return 0
        else
            echo "⚠ CSF block failed for $ip, trying iptables..."
        fi
    fi

    # Fallback to iptables
    if command -v iptables &>/dev/null; then
        # Check if already blocked
        if iptables -L INPUT -n | grep -q "$ip"; then
            echo "⚠ $ip already blocked in iptables"
            return 0
        fi

        # Add iptables rule
        iptables -I INPUT -s "$ip" -j DROP

        if [ $? -eq 0 ]; then
            echo "✓ Blocked $ip using iptables for ${duration}h: $reason"
            record_ip_ban "$ip" "$duration" "$reason"

            # Schedule removal using at (if available)
            if command -v at &>/dev/null; then
                echo "iptables -D INPUT -s $ip -j DROP 2>/dev/null" | at now + $duration hours 2>/dev/null
                echo "  (Scheduled auto-unblock in ${duration}h)"
            else
                echo "  (WARNING: Manual unblock required - 'at' command not available)"
            fi

            return 0
        else
            echo "✗ Failed to block $ip with iptables"
            return 1
        fi
    fi

    echo "✗ No firewall available (CSF or iptables required)"
    return 1
}

# Unblock IP
unblock_ip() {
    local ip="$1"

    [ -z "$ip" ] && return 1

    # Try CSF first
    if command -v csf &>/dev/null; then
        csf -tr "$ip" &>/dev/null
        if [ $? -eq 0 ]; then
            echo "✓ Unblocked $ip from CSF"
            return 0
        fi
    fi

    # Try iptables
    if command -v iptables &>/dev/null; then
        iptables -D INPUT -s "$ip" -j DROP 2>/dev/null
        if [ $? -eq 0 ]; then
            echo "✓ Unblocked $ip from iptables"
            return 0
        fi
    fi

    echo "⚠ $ip not found in firewall rules"
    return 1
}

# Check if IP is currently blocked
is_ip_blocked() {
    local ip="$1"

    [ -z "$ip" ] && return 1

    # Check CSF
    if command -v csf &>/dev/null; then
        if csf -g "$ip" 2>/dev/null | grep -q "DENY"; then
            return 0
        fi
    fi

    # Check iptables (use word boundaries to avoid partial matches)
    if command -v iptables &>/dev/null; then
        if iptables -L INPUT -n 2>/dev/null | grep -w "$ip" | grep -q "DROP\|REJECT"; then
            return 0
        fi
    fi

    return 1
}

# Get list of IPs that should be blocked based on reputation
# Usage: get_blockable_ips [MIN_SCORE]
get_blockable_ips() {
    local min_score="${1:-60}"  # Default: score >= 60

    [ ! -f "$IP_REP_DB" ] && return 1

    # Get IPs with score >= min_score, not already blocked
    while IFS='|' read -r ip hit_count rep_score rest; do
        # Skip if score too low
        [ "$rep_score" -lt "$min_score" ] 2>/dev/null && continue

        # Skip if already blocked
        is_ip_blocked "$ip" && continue

        # Output: IP|SCORE|HITS
        echo "$ip|$rep_score|$hit_count"
    done < "$IP_REP_DB" | sort -t'|' -k2 -rn
}

export -f record_ip_ban
export -f get_ip_ban_count
export -f get_ip_last_ban
export -f block_ip_temporary
export -f unblock_ip
export -f is_ip_blocked
export -f get_blockable_ips

# Initialize on library load
init_ip_reputation_db