#!/bin/bash

################################################################################
# IP Reputation Manager
################################################################################
# Purpose: View, query, and manage the centralized IP reputation database
# Features:
# - Query individual IPs
# - View top malicious IPs
# - View top active IPs
# - Export database
# - Database statistics
# - Cleanup old entries
# - Manual IP flagging/whitelisting
################################################################################

# Get script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
source "$SCRIPT_DIR/lib/common-functions.sh"
source "$SCRIPT_DIR/lib/system-detect.sh"
source "$SCRIPT_DIR/lib/ip-reputation.sh"

# Require root
if [ "$EUID" -ne 0 ]; then
    print_error "This script must be run as root"
    exit 1
fi

# Menu display
show_menu() {
    clear
    print_banner "IP Reputation Manager"

    # Show quick stats
    local total_ips=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo 0)
    local db_size=$(du -h "$IP_REP_DB" 2>/dev/null | awk '{print $1}' || echo "0B")

    echo ""
    echo -e "${BLUE}${BOLD}Database Status:${NC} $total_ips IPs tracked | Size: $db_size"
    echo ""
    echo -e "${BOLD}Query & View:${NC}"
    echo ""
    echo -e "  ${GREEN}1)${NC} Query IP Reputation       - Look up specific IP"
    echo -e "  ${GREEN}2)${NC} Top Malicious IPs         - Highest reputation scores"
    echo -e "  ${GREEN}3)${NC} Top Active IPs            - Most hits/requests"
    echo -e "  ${GREEN}4)${NC} Database Statistics       - Overview of tracked IPs"
    echo -e "  ${GREEN}5)${NC} Live Monitoring           - Real-time reputation updates"
    echo ""
    echo -e "${BOLD}Database Management:${NC}"
    echo ""
    echo -e "  ${BLUE}6)${NC} Export Database           - Export to readable text file"
    echo -e "  ${BLUE}7)${NC} Cleanup Old Entries       - Remove IPs not seen in X days"
    echo -e "  ${BLUE}8)${NC} Compact Database          - Remove duplicate entries (faster writes)"
    echo -e "  ${BLUE}9)${NC} Rebuild Index             - Optimize database for speed"
    echo ""
    echo -e "${BOLD}Manual Actions:${NC}"
    echo ""
    echo -e "  ${YELLOW}10)${NC} Flag IP as Malicious      - Manually mark IP as threat"
    echo -e "  ${YELLOW}11)${NC} Mark IP as Legitimate     - Whitelist/reduce score"
    echo -e "  ${YELLOW}12)${NC} Import IPs from Log       - Batch import from file"
    echo ""
    echo -e "  ${RED}0)${NC} Exit"
    echo ""
    echo -e "${CYAN}──────────────────────────────────────────────────────────────${NC}"
    echo -n "Select option: "
}

# Query individual IP
query_ip_interactive() {
    clear
    print_banner "Query IP Reputation"
    echo ""
    echo -n "Enter IP address to query: "
    read -r ip_address

    if [ -z "$ip_address" ]; then
        print_error "No IP address provided"
        press_enter
        return
    fi

    # Validate IP format (basic check)
    if ! [[ "$ip_address" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        print_error "Invalid IP address format"
        press_enter
        return
    fi

    echo ""
    query_ip_reputation "$ip_address"

    echo ""
    press_enter
}

# View top malicious IPs
view_top_malicious() {
    clear
    print_banner "Top Malicious IPs"
    echo ""
    echo -n "How many top IPs to show? [20]: "
    read -r limit
    limit="${limit:-20}"

    echo ""
    echo -e "${RED}${BOLD}Top $limit Most Malicious IPs (by Reputation Score)${NC}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    printf "%-15s | %-7s | %-4s | %-8s | %-8s | %-30s\n" \
        "IP ADDRESS" "HITS" "CTRY" "REP" "LEVEL" "ATTACK TYPES"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    get_top_malicious_ips "$limit" | 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")

        # Color code by reputation
        local color="$NC"
        case "$category" in
            CRITICAL) color="$RED$BOLD" ;;
            HIGH) color="$RED" ;;
            MEDIUM) color="$YELLOW" ;;
            LOW) color="$CYAN" ;;
        esac

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

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    press_enter
}

# View top active IPs
view_top_active() {
    clear
    print_banner "Top Active IPs"
    echo ""
    echo -n "How many top IPs to show? [20]: "
    read -r limit
    limit="${limit:-20}"

    echo ""
    echo -e "${YELLOW}${BOLD}Top $limit Most Active IPs (by Hit Count)${NC}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    printf "%-15s | %-7s | %-4s | %-8s | %-8s | %-30s\n" \
        "IP ADDRESS" "HITS" "CTRY" "REP" "LEVEL" "ATTACK TYPES"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    get_top_active_ips "$limit" | 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")

        # Color code by hit count
        local color="$NC"
        if [ $hit_count -gt 10000 ]; then
            color="$RED$BOLD"
        elif [ $hit_count -gt 1000 ]; then
            color="$YELLOW"
        fi

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

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    press_enter
}

# Show statistics
view_statistics() {
    clear
    print_banner "Database Statistics"
    echo ""
    show_ip_statistics
    echo ""

    # Additional stats
    echo "Recent Activity (Last 24 hours):"
    local cutoff=$(($(date +%s) - 86400))
    local recent_count=$(awk -F'|' -v cut="$cutoff" '$7 >= cut' "$IP_REP_DB" 2>/dev/null | wc -l)
    echo "  Active IPs:         $recent_count"
    echo ""

    # Top countries
    echo "Top Countries by IP Count:"
    awk -F'|' '{print $4}' "$IP_REP_DB" 2>/dev/null | grep -v '^$' | sort | uniq -c | sort -rn | head -5 | \
    while read count country; do
        printf "  %-4s: %s IPs\n" "$country" "$count"
    done
    echo ""

    press_enter
}

# Export database
export_database_interactive() {
    clear
    print_banner "Export IP Reputation Database"
    echo ""
    echo -n "Enter output file path [/tmp/ip_reputation_export.txt]: "
    read -r output_path
    output_path="${output_path:-/tmp/ip_reputation_export.txt}"

    echo ""
    echo "Exporting database to $output_path..."
    export_ip_reputation "$output_path"

    echo ""
    print_success "Database exported successfully!"
    echo ""
    echo "View with: cat $output_path"
    echo "Or:        less $output_path"
    echo ""
    press_enter
}

# Cleanup old entries
cleanup_database_interactive() {
    clear
    print_banner "Cleanup Old Entries"
    echo ""
    echo "Remove IPs that haven't been seen in how many days?"
    echo ""
    echo -n "Days [90]: "
    read -r days
    days="${days:-90}"

    echo ""
    echo "This will remove IPs not seen in the last $days days."
    echo -n "Continue? (yes/no): "
    read -r confirm

    if [ "$confirm" != "yes" ]; then
        echo "Cancelled"
        press_enter
        return
    fi

    echo ""
    cleanup_old_ips "$days"
    echo ""
    print_success "Cleanup complete!"
    echo ""
    press_enter
}

# Compact database
compact_database_interactive() {
    clear
    print_banner "Compact Database"
    echo ""
    local total_before=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo 0)
    echo "Current database size: $total_before entries"
    echo ""
    echo "This will remove duplicate IP entries created by fast append-only writes."
    echo "The database will be compacted and re-indexed."
    echo ""
    echo -n "Continue? (yes/no): "
    read -r confirm

    if [ "$confirm" != "yes" ]; then
        echo "Cancelled"
        press_enter
        return
    fi

    echo ""
    compact_database
    echo ""
    print_success "Database compacted successfully!"
    echo ""
    press_enter
}

# Rebuild index
rebuild_index_interactive() {
    clear
    print_banner "Rebuild Database Index"
    echo ""
    echo "Rebuilding index for optimized lookups..."
    rebuild_index
    echo ""
    print_success "Index rebuilt successfully!"
    echo ""
    press_enter
}

# Flag IP as malicious
flag_ip_interactive() {
    clear
    print_banner "Flag IP as Malicious"
    echo ""
    echo -n "Enter IP address: "
    read -r ip_address

    if [ -z "$ip_address" ]; then
        print_error "No IP address provided"
        press_enter
        return
    fi

    echo ""
    echo "Attack Type:"
    echo "  1) SQL Injection"
    echo "  2) XSS"
    echo "  3) Path Traversal"
    echo "  4) RCE/Shell Upload"
    echo "  5) Brute Force"
    echo "  6) DDoS"
    echo "  7) Bot/Scanner"
    echo "  8) Exploit"
    echo ""
    echo -n "Select [1]: "
    read -r attack_choice
    attack_choice="${attack_choice:-1}"

    case "$attack_choice" in
        1) attack_type="SQL_INJECTION" ;;
        2) attack_type="XSS" ;;
        3) attack_type="PATH_TRAVERSAL" ;;
        4) attack_type="RCE" ;;
        5) attack_type="BRUTEFORCE" ;;
        6) attack_type="DDOS" ;;
        7) attack_type="SCANNER" ;;
        8) attack_type="EXPLOIT" ;;
        *) attack_type="SUSPICIOUS" ;;
    esac

    echo ""
    echo -n "Notes/Description: "
    read -r notes

    echo ""
    echo "Flagging $ip_address for $attack_type..."
    flag_ip_attack "$ip_address" "$attack_type" 0 "$notes"

    echo ""
    print_success "IP flagged successfully!"
    echo ""
    query_ip_reputation "$ip_address"
    echo ""
    press_enter
}

# Mark IP as legitimate
whitelist_ip_interactive() {
    clear
    print_banner "Mark IP as Legitimate"
    echo ""
    echo -n "Enter IP address: "
    read -r ip_address

    if [ -z "$ip_address" ]; then
        print_error "No IP address provided"
        press_enter
        return
    fi

    echo ""
    echo -n "Reason/Notes: "
    read -r notes

    echo ""
    echo "Marking $ip_address as legitimate..."
    mark_ip_legitimate "$ip_address" "$notes"

    echo ""
    print_success "IP marked as legitimate!"
    echo ""
    query_ip_reputation "$ip_address"
    echo ""
    press_enter
}

# Import from log file
import_log_interactive() {
    clear
    print_banner "Import IPs from Log File"
    echo ""
    echo -n "Enter log file path: "
    read -r log_path

    if [ ! -f "$log_path" ]; then
        print_error "File not found: $log_path"
        press_enter
        return
    fi

    echo ""
    echo "Attack Type (will be applied to all IPs):"
    echo "  1) Suspicious (default)"
    echo "  2) SQL Injection"
    echo "  3) XSS"
    echo "  4) Bot/Scanner"
    echo "  5) DDoS"
    echo ""
    echo -n "Select [1]: "
    read -r attack_choice
    attack_choice="${attack_choice:-1}"

    case "$attack_choice" in
        2) attack_type="SQL_INJECTION" ;;
        3) attack_type="XSS" ;;
        4) attack_type="SCANNER" ;;
        5) attack_type="DDOS" ;;
        *) attack_type="SUSPICIOUS" ;;
    esac

    echo ""
    echo "Importing IPs from $log_path as $attack_type..."
    import_ips_from_log "$log_path" "$attack_type" 5

    echo ""
    print_success "Import complete!"
    echo ""
    press_enter
}

# Live monitoring (real-time updates)
live_monitoring() {
    clear
    print_banner "Live IP Reputation Monitoring"
    echo ""
    echo "Watching database for changes... (Press Ctrl+C to exit)"
    echo ""

    local last_count=0
    local last_update=0

    while true; do
        local current_count=$(wc -l < "$IP_REP_DB" 2>/dev/null || echo 0)
        local current_time=$(stat -c %Y "$IP_REP_DB" 2>/dev/null || echo 0)

        if [ $current_count -ne $last_count ] || [ $current_time -ne $last_update ]; then
            clear
            echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
            echo "LIVE IP REPUTATION MONITORING - $(date '+%H:%M:%S')"
            echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
            echo ""
            show_ip_statistics
            echo ""
            echo "Recent Top 10 Updates:"
            get_top_malicious_ips 10 | head -10 | while IFS='|' read -r ip hit_count rep_score country attack_flags _ _ _ _; do
                local category=$(get_ip_reputation_category "$rep_score")
                local attacks=$(decode_attack_flags "$attack_flags")
                printf "%-15s | %5s hits | %3s/100 | %-8s | %s\n" "$ip" "$hit_count" "$rep_score" "$category" "${attacks:0:40}"
            done
            echo ""
            echo "Press Ctrl+C to exit | Refreshing every 2 seconds..."

            last_count=$current_count
            last_update=$current_time
        fi

        sleep 2
    done
}

# Main loop
main() {
    while true; do
        show_menu
        read -r choice

        case $choice in
            1) query_ip_interactive ;;
            2) view_top_malicious ;;
            3) view_top_active ;;
            4) view_statistics ;;
            5) live_monitoring ;;
            6) export_database_interactive ;;
            7) cleanup_database_interactive ;;
            8) compact_database_interactive ;;
            9) rebuild_index_interactive ;;
            10) flag_ip_interactive ;;
            11) whitelist_ip_interactive ;;
            12) import_log_interactive ;;
            0)
                clear
                echo "Exiting..."
                exit 0
                ;;
            *)
                print_error "Invalid option"
                sleep 1
                ;;
        esac
    done
}

# Run main
main