#!/bin/bash
#
# Rate-Based Anomaly Detection
# Detects HTTP floods, brute force, and other rate-based attacks

# Temporary directory for rate tracking
RATE_TRACKING_DIR="${RATE_TRACKING_DIR:-/var/tmp/rate-tracking}"
mkdir -p "$RATE_TRACKING_DIR" 2>/dev/null

# Record a request timestamp for an IP
# Usage: record_request "192.168.1.100" [timestamp]
record_request() {
    local ip="$1"
    local timestamp="${2:-$(date +%s)}"

    local rate_file="$RATE_TRACKING_DIR/${ip//\./_}.dat"
    echo "$timestamp" >> "$rate_file"
}

# Detect rate anomalies for an IP
# Usage: detect_rate_anomaly "192.168.1.100" [current_time]
# Returns: anomaly_score||anomaly_type||req_per_sec||req_per_10sec||req_per_min
detect_rate_anomaly() {
    local ip="$1"
    local current_time="${2:-$(date +%s)}"

    local rate_file="$RATE_TRACKING_DIR/${ip//\./_}.dat"

    # No history = no anomaly
    if [ ! -f "$rate_file" ]; then
        echo "0||NORMAL||0||0||0"
        return 0
    fi

    # Count requests in different time windows
    local req_1sec=$(awk -v cutoff="$((current_time - 1))" '$1 > cutoff' -- "$rate_file" 2>/dev/null | wc -l)
    local req_10sec=$(awk -v cutoff="$((current_time - 10))" '$1 > cutoff' -- "$rate_file" 2>/dev/null | wc -l)
    local req_60sec=$(awk -v cutoff="$((current_time - 60))" '$1 > cutoff' -- "$rate_file" 2>/dev/null | wc -l)

    local anomaly_score=0
    local anomaly_type="NORMAL"

    # HTTP flood detection thresholds
    if [ "$req_1sec" -gt 100 ]; then
        # >100 requests per second = Critical flood
        anomaly_score=95
        anomaly_type="HTTP_FLOOD_CRITICAL"
    elif [ "$req_1sec" -gt 50 ]; then
        # >50 requests per second = High flood
        anomaly_score=85
        anomaly_type="HTTP_FLOOD_HIGH"
    elif [ "$req_10sec" -gt 200 ]; then
        # >200 in 10 sec (20/sec sustained) = Sustained flood
        anomaly_score=80
        anomaly_type="HTTP_FLOOD_SUSTAINED"
    elif [ "$req_10sec" -gt 100 ]; then
        # >100 in 10 sec (10/sec sustained) = Moderate flood
        anomaly_score=70
        anomaly_type="HTTP_FLOOD_MODERATE"
    elif [ "$req_60sec" -gt 300 ]; then
        # >300 in 60 sec (5/sec sustained) = High rate
        anomaly_score=60
        anomaly_type="HIGH_RATE"
    elif [ "$req_60sec" -gt 150 ]; then
        # >150 in 60 sec (2.5/sec sustained) = Elevated rate
        anomaly_score=40
        anomaly_type="ELEVATED_RATE"
    elif [ "$req_60sec" -gt 60 ]; then
        # >60 in 60 sec (1/sec sustained) = Suspicious rate
        anomaly_score=20
        anomaly_type="SUSPICIOUS_RATE"
    fi

    # Cleanup old entries (keep last 60 seconds only)
    if [ -f "$rate_file" ]; then
        awk -v cutoff="$((current_time - 60))" '$1 > cutoff' -- "$rate_file" > "${rate_file}.tmp" 2>/dev/null
        mv "${rate_file}.tmp" "$rate_file" 2>/dev/null
    fi

    echo "$anomaly_score||$anomaly_type||$req_1sec||$req_10sec||$req_60sec"
}

# Analyze request pattern (burst detection)
# Usage: analyze_request_pattern "192.168.1.100" [window_seconds]
# Returns: pattern_type||burst_count||distribution_score
analyze_request_pattern() {
    local ip="$1"
    local window="${2:-60}"  # Default 60 second window

    local rate_file="$RATE_TRACKING_DIR/${ip//\./_}.dat"

    if [ ! -f "$rate_file" ]; then
        echo "NONE||0||0"
        return 0
    fi

    local current_time=$(date +%s)
    local cutoff=$((current_time - window))

    # Get timestamps in window
    local timestamps=$(awk -v cutoff="$cutoff" '$1 > cutoff {print $1}' -- "$rate_file" 2>/dev/null | sort -n)
    local total_count=$(echo "$timestamps" | wc -l)

    if [ "$total_count" -lt 5 ]; then
        echo "NORMAL||0||0"
        return 0
    fi

    # Calculate time gaps between requests
    local prev_time=0
    local gaps=()
    local burst_count=0
    local regular_count=0

    while IFS= read -r ts; do
        if [ "$prev_time" -gt 0 ]; then
            local gap=$((ts - prev_time))
            if [ "$gap" -lt 1 ]; then
                # Burst: Multiple requests in same second
                burst_count=$((burst_count + 1))
            elif [ "$gap" -lt 5 ]; then
                # Rapid: Requests within 5 seconds
                burst_count=$((burst_count + 1))
            else
                # Regular spacing
                regular_count=$((regular_count + 1))
            fi
        fi
        prev_time=$ts
    done <<< "$timestamps"

    # Determine pattern type
    local pattern_type="NORMAL"
    local distribution_score=0

    if [ "$burst_count" -gt "$((total_count / 2))" ]; then
        # More than half are bursts
        pattern_type="BURST"
        distribution_score=70
    elif [ "$regular_count" -gt "$((total_count * 3 / 4))" ]; then
        # Regular intervals (bot-like behavior)
        pattern_type="AUTOMATED"
        distribution_score=50
    else
        # Mixed pattern
        pattern_type="MIXED"
        distribution_score=30
    fi

    echo "$pattern_type||$burst_count||$distribution_score"
}

# Cleanup old rate tracking files
# Usage: cleanup_rate_tracking [max_age_seconds]
cleanup_rate_tracking() {
    local max_age="${1:-300}"  # Default 5 minutes

    if [ ! -d "$RATE_TRACKING_DIR" ]; then
        return 0
    fi

    # Find and delete files older than max_age
    find "$RATE_TRACKING_DIR" -type f -name "*.dat" -mmin "+$((max_age / 60))" -delete 2>/dev/null

    # Also clean up empty files
    find "$RATE_TRACKING_DIR" -type f -name "*.dat" -empty -delete 2>/dev/null
}

# Get current request rate for an IP
# Usage: get_current_rate "192.168.1.100" [window_seconds]
# Returns: requests_per_second (as integer)
get_current_rate() {
    local ip="$1"
    local window="${2:-60}"  # Default 60 second window

    local rate_file="$RATE_TRACKING_DIR/${ip//\./_}.dat"

    if [ ! -f "$rate_file" ]; then
        echo "0"
        return 0
    fi

    local current_time=$(date +%s)
    local cutoff=$((current_time - window))
    local count=$(awk -v cutoff="$cutoff" '$1 > cutoff' -- "$rate_file" 2>/dev/null | wc -l)

    # Calculate requests per second
    local rate=$((count / window))
    echo "$rate"
}

# Check if IP is currently flooding
# Usage: is_flooding "192.168.1.100" [threshold]
# Returns: 0 if flooding, 1 if not
is_flooding() {
    local ip="$1"
    local threshold="${2:-10}"  # Default 10 req/sec

    local rate=$(get_current_rate "$ip" 10)  # Check 10 second window

    if [ "$rate" -ge "$threshold" ]; then
        return 0  # Is flooding
    else
        return 1  # Not flooding
    fi
}

# Format rate anomaly for display
# Usage: format_rate_anomaly "$anomaly_result"
format_rate_anomaly() {
    local result="$1"

    local score="${result%%||*}"
    local temp="${result#*||}"
    local type="${temp%%||*}"
    temp="${temp#*||}"
    local req_1s="${temp%%||*}"
    temp="${temp#*||}"
    local req_10s="${temp%%||*}"
    local req_60s="${temp#*||}"

    local color="\033[0;36m"  # Cyan
    if [ "$score" -ge 85 ]; then
        color="\033[0;31m"  # Red
    elif [ "$score" -ge 70 ]; then
        color="\033[1;33m"  # Yellow
    fi

    echo -e "${color}[$type:$score]${NC} Rate: $req_1s/sec | $req_10s/10s | $req_60s/min"
}

# Initialize rate tracking (create directory)
init_rate_tracking() {
    mkdir -p "$RATE_TRACKING_DIR" 2>/dev/null
    chmod 700 "$RATE_TRACKING_DIR" 2>/dev/null
}

# Auto-cleanup background task (run periodically)
start_rate_cleanup_task() {
    local interval="${1:-300}"  # Default 5 minutes

    while true; do
        sleep "$interval"
        cleanup_rate_tracking "$interval"
    done &

    echo $!  # Return PID of cleanup task
}

# Export functions for use in subshells
export -f record_request
export -f detect_rate_anomaly
export -f analyze_request_pattern
export -f cleanup_rate_tracking
export -f get_current_rate
export -f is_flooding
export -f format_rate_anomaly
export -f init_rate_tracking
export -f start_rate_cleanup_task
