#!/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
# Example:
# 192.168.1.100|523|75|US|193|1730000000|1730800000|SQL injection on /admin|Auto-flagged

# 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 grep with optimizations
# Returns: IP data if found, empty if not found
lookup_ip() {
    local ip="$1"

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

    # Use grep with fixed string for speed
    grep -m 1 "^${ip}|" "$IP_REP_DB" 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

        # Remove old entry and add updated one
        sed -i "/^${ip}|/d" "$IP_REP_DB"
        echo "$ip|$hit_count|$rep_score|$country|$attack_flags|$first_seen|$last_seen|$last_activity|$notes" >> "$IP_REP_DB"
    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

    # Rebuild index if database is getting large
    local db_size=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo "0")
    if [ $db_size -gt 10000 ] && [ $((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

    # Sort by reputation score (field 3), descending
    sort -t'|' -k3 -rn "$IP_REP_DB" | head -n "$limit"
}

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

    init_ip_reputation_db

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

    # Sort by hit count (field 2), descending
    sort -t'|' -k2 -rn "$IP_REP_DB" | head -n "$limit"
}

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

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

    # Create sorted index by IP for binary search (future optimization)
    sort -t'|' -k1 "$IP_REP_DB" > "$IP_REP_INDEX"

    release_lock
}

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

# Initialize on library load
init_ip_reputation_db