#!/bin/bash
################################################################################
# Script Name: ufw-blocklist-metrics.sh
# Version: 2.3
# Description: Production Prometheus exporter for UFW Blocklists (OPTIMIZED)
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Optimizations in v2.1:
#   - Single journalctl call with cached output
#   - Cached feed config parsing
#   - Eliminated redundant file operations
#   - 4.5 minutes → ~30 seconds typical runtime
#
# Fixes in v2.2:
#   - Fixed typo in script name header (bocklist → blocklist)
#   - Fixed ipset member counting to use Members: section
#   - Fixed empty journal data producing false grep counts
#   - Fixed HTTP response headers missing trailing \r\n
#   - Fixed SC2155/SC2126/SC2295 shellcheck warnings
#   - Added scrape timestamp metric
#   - Used SCRIPT_VERSION variable for version strings
#
# Fixes in v2.3:
#   - Fixed get_ipset_size using grep -c (exit 1 on 0 matches) causing
#     duplicate "0" output lines and arithmetic errors; switched to wc -l
#   - Fixed same grep -c || echo 0 bug in ufw_blocklist_enabled and
#     ufw_blocklist_total_rules heredoc substitutions
#   - Fixed misplaced 2>/dev/null on [ ] test for conntrack and effectiveness
#   - Fixed hardcoded v2.1 in usage text; now uses SCRIPT_VERSION
################################################################################

CONFIG_DIR="/etc/ufw-threats"
CACHE_DIR="$CONFIG_DIR/cache"
FEEDS_CONFIG="$CONFIG_DIR/feeds.conf"
IPSET_PREFIX="ufw-feed"
WHITELIST_IPSET="ufw-whitelist"
WHITELIST_IPSET_V6="ufw-whitelist-v6"
SCRIPT_VERSION="2.3"

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9418
LOCK_FILE="/var/run/ufw-blocklist-metrics.lock"

# Global cache variables
JOURNAL_1H=""
JOURNAL_24H=""
FEEDS_ARRAY=()

show_usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Export per-feed UFW threat statistics as Prometheus metrics (v${SCRIPT_VERSION}).

MODES:
    --textfile      Write to node_exporter textfile collector
    --http          Run HTTP server on port $HTTP_PORT

OPTIONS:
    -p, --port      HTTP port
    -o, --output    Output file
    -h, --help      Show help

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            --textfile) OUTPUT_FILE="$TEXTFILE_DIR/ufw_blocklist_metrics.prom"; shift ;;
            --http) HTTP_MODE=true; shift ;;
            -p|--port) HTTP_PORT="$2"; shift 2 ;;
            -o|--output) OUTPUT_FILE="$2"; shift 2 ;;
            *) echo "Unknown: $1"; exit 1 ;;
        esac
    done
}

# Load all journal data ONCE at startup
cache_journal_data() {
    JOURNAL_1H=$(timeout 10 journalctl --since "1 hour ago" 2>/dev/null | grep '\[THREAT' || echo "")
    JOURNAL_24H=$(timeout 30 journalctl --since "24 hours ago" 2>/dev/null | grep '\[THREAT' || echo "")
}

# Parse feeds config ONCE into array
cache_feeds_config() {
    FEEDS_ARRAY=()
    if [ -f "$FEEDS_CONFIG" ]; then
        while IFS='|' read -r enabled name url type description; do
            [[ "$enabled" =~ ^#.*$ ]] && continue
            [[ -z "$enabled" ]] && continue
            FEEDS_ARRAY+=("$enabled|$name|$url|$type|$description")
        done < "$FEEDS_CONFIG"
    fi
}

get_ipset_size() {
    local ipset_name="$1"
    local count
    count=$(ipset list "$ipset_name" 2>/dev/null | sed -n '/^Members:$/,$p' | tail -n +2 | wc -l)
    echo "${count:-0}"
}

# Optimized: Use cached journal data
get_feed_blocks() {
    local feed="$1"
    local period="$2"
    local data
    
    case "$period" in
        "1 hour ago") data="$JOURNAL_1H" ;;
        "24 hours ago") data="$JOURNAL_24H" ;;
        *) echo 0; return ;;
    esac
    
    if [ -z "$data" ]; then echo 0; return; fi
    local count
    count=$(printf '%s' "$data" | grep -c "\[THREAT:${feed}\]" 2>/dev/null)
    echo "${count:-0}"
}

get_feed_blocks_v6() {
    local feed="$1"
    local period="$2"
    local data
    
    case "$period" in
        "1 hour ago") data="$JOURNAL_1H" ;;
        "24 hours ago") data="$JOURNAL_24H" ;;
        *) echo 0; return ;;
    esac
    
    if [ -z "$data" ]; then echo 0; return; fi
    local count
    count=$(printf '%s' "$data" | grep -c "\[THREAT-v6:${feed}\]" 2>/dev/null)
    echo "${count:-0}"
}

get_file_timestamp() {
    [ -f "$1" ] && stat -c %Y "$1" 2>/dev/null || echo "0"
}

get_file_size() {
    [ -f "$1" ] && stat -c %s "$1" 2>/dev/null || echo "0"
}

get_cache_age() {
    if [ -f "$1" ]; then
        echo $(($(date +%s) - $(stat -c %Y "$1" 2>/dev/null || echo 0)))
    else
        echo "0"
    fi
}

get_conntrack_count() {
    if [ -f /proc/sys/net/netfilter/nf_conntrack_count ]; then
        cat /proc/sys/net/netfilter/nf_conntrack_count
    else
        echo "0"
    fi
}

get_conntrack_max() {
    if [ -f /proc/sys/net/netfilter/nf_conntrack_max ]; then
        cat /proc/sys/net/netfilter/nf_conntrack_max
    else
        echo "0"
    fi
}

get_ipset_memory() {
    local ipset_name="$1"
    local mem
    mem=$(ipset list "$ipset_name" -t 2>/dev/null | grep "Size in memory:" | awk '{print $4}')
    echo "${mem:-0}"
}

get_cache_disk_usage() {
    if [ -d "$CACHE_DIR" ]; then
        df -B1 "$CACHE_DIR" 2>/dev/null | tail -1 | awk '{print $3"|"$4"|"$5}'
    else
        echo "0|0|0%"
    fi
}

get_total_cache_size() {
    if [ -d "$CACHE_DIR" ]; then
        du -sb "$CACHE_DIR" 2>/dev/null | awk '{print $1}'
    else
        echo "0"
    fi
}

acquire_lock() {
    if [ -f "$LOCK_FILE" ]; then
        local pid
        pid=$(cat "$LOCK_FILE" 2>/dev/null)
        if [ -n "$pid" ] && kill -0 "$pid" 2>/dev/null; then
            echo "ERROR: Another instance is already running (PID: $pid)" >&2
            exit 1
        else
            echo "Removing stale lock file" >&2
            rm -f "$LOCK_FILE"
        fi
    fi
    echo $$ > "$LOCK_FILE"
    trap cleanup EXIT INT TERM
}

cleanup() {
    rm -f "$LOCK_FILE"
}

generate_metrics() {
    local start_time
    start_time=$(date +%s)
    
    cat <<EOF
# HELP ufw_blocklist_info Per-feed UFW threat blocking info
# TYPE ufw_blocklist_info gauge
ufw_blocklist_info{mode="per-feed",version="${SCRIPT_VERSION}"} 1

# HELP ufw_blocklist_enabled Total enabled feeds
# TYPE ufw_blocklist_enabled gauge
ufw_blocklist_enabled $(printf '%s\n' "${FEEDS_ARRAY[@]}" | grep -c '^1|')

# HELP ufw_blocklist_ipset_size Number of IPs per feed ipset
# TYPE ufw_blocklist_ipset_size gauge
EOF

    # Only export metrics for ipsets that actually exist and are enabled
    for ipset_name in $(ipset list -n 2>/dev/null | grep "^${IPSET_PREFIX}-"); do
        # Extract feed name and IP version
        local feed_name="${ipset_name#"${IPSET_PREFIX}"-}"
        local ip_version="4"
        
        if [[ "$feed_name" =~ -v6$ ]]; then
            feed_name="${feed_name%-v6}"
            ip_version="6"
        fi
        
        # Only show enabled feeds
        if ! printf '%s\n' "${FEEDS_ARRAY[@]}" | grep -q "^1|${feed_name}|" 2>/dev/null; then
            continue
        fi
        
        local size
        size=$(get_ipset_size "$ipset_name")
        echo "ufw_blocklist_ipset_size{feed=\"$feed_name\",ip_version=\"$ip_version\",status=\"enabled\"} $size"
    done

    cat <<EOF

# HELP ufw_blocklist_cache_age_seconds Seconds since feed update
# TYPE ufw_blocklist_cache_age_seconds gauge
EOF

    for feed_line in "${FEEDS_ARRAY[@]}"; do
        IFS='|' read -r enabled name url type description <<< "$feed_line"
        [ "$enabled" != "1" ] && continue
        
        age_v4=$(get_cache_age "$CACHE_DIR/${name}-v4.parsed")
        age_v6=$(get_cache_age "$CACHE_DIR/${name}-v6.parsed")
        
        echo "ufw_blocklist_cache_age_seconds{feed=\"$name\",ip_version=\"4\"} $age_v4"
        echo "ufw_blocklist_cache_age_seconds{feed=\"$name\",ip_version=\"6\"} $age_v6"
    done

    cat <<EOF

# HELP ufw_blocklist_last_update_timestamp Unix timestamp of last update
# TYPE ufw_blocklist_last_update_timestamp gauge
EOF

    for feed_line in "${FEEDS_ARRAY[@]}"; do
        IFS='|' read -r enabled name url type description <<< "$feed_line"
        [ "$enabled" != "1" ] && continue
        
        ts_v4=$(get_file_timestamp "$CACHE_DIR/${name}-v4.parsed")
        ts_v6=$(get_file_timestamp "$CACHE_DIR/${name}-v6.parsed")
        
        echo "ufw_blocklist_last_update_timestamp{feed=\"$name\",ip_version=\"4\"} $ts_v4"
        echo "ufw_blocklist_last_update_timestamp{feed=\"$name\",ip_version=\"6\"} $ts_v6"
    done

    cat <<EOF

# HELP ufw_blocklist_file_size_bytes Feed file sizes
# TYPE ufw_blocklist_file_size_bytes gauge
EOF

    for feed_line in "${FEEDS_ARRAY[@]}"; do
        IFS='|' read -r enabled name url type description <<< "$feed_line"
        [ "$enabled" != "1" ] && continue
        
        parsed_v4=$(get_file_size "$CACHE_DIR/${name}-v4.parsed")
        parsed_v6=$(get_file_size "$CACHE_DIR/${name}-v6.parsed")
        
        echo "ufw_blocklist_file_size_bytes{feed=\"$name\",ip_version=\"4\",type=\"parsed\"} $parsed_v4"
        echo "ufw_blocklist_file_size_bytes{feed=\"$name\",ip_version=\"6\",type=\"parsed\"} $parsed_v6"
    done

    cat <<EOF

# HELP ufw_blocklist_ip_version_ratio IPv4 vs IPv6 ratio
# TYPE ufw_blocklist_ip_version_ratio gauge
EOF

    for feed_line in "${FEEDS_ARRAY[@]}"; do
        IFS='|' read -r enabled name url type description <<< "$feed_line"
        [ "$enabled" != "1" ] && continue
        
        v4_size=$(get_ipset_size "${IPSET_PREFIX}-${name}")
        v6_size=$(get_ipset_size "${IPSET_PREFIX}-${name}-v6")
        total=$((v4_size + v6_size))
        
        if [ "$total" -gt 0 ]; then
            ratio_v4=$(awk "BEGIN {printf \"%.4f\", $v4_size / $total}")
            ratio_v6=$(awk "BEGIN {printf \"%.4f\", $v6_size / $total}")
        else
            ratio_v4="0.0000"
            ratio_v6="0.0000"
        fi
        
        echo "ufw_blocklist_ip_version_ratio{feed=\"$name\",version=\"4\"} $ratio_v4"
        echo "ufw_blocklist_ip_version_ratio{feed=\"$name\",version=\"6\"} $ratio_v6"
    done

    cat <<EOF

# HELP ufw_blocklist_blocked_total Blocked attempts per feed (rolling window)
# TYPE ufw_blocklist_blocked_total gauge
EOF

    for feed_line in "${FEEDS_ARRAY[@]}"; do
        IFS='|' read -r enabled name url type description <<< "$feed_line"
        [ "$enabled" != "1" ] && continue
        
        blocks_1h=$(get_feed_blocks "$name" "1 hour ago")
        blocks_24h=$(get_feed_blocks "$name" "24 hours ago")
        blocks_1h_v6=$(get_feed_blocks_v6 "$name" "1 hour ago")
        blocks_24h_v6=$(get_feed_blocks_v6 "$name" "24 hours ago")
        
        echo "ufw_blocklist_blocked_total{feed=\"$name\",ip_version=\"4\",period=\"1h\"} $blocks_1h"
        echo "ufw_blocklist_blocked_total{feed=\"$name\",ip_version=\"4\",period=\"24h\"} $blocks_24h"
        echo "ufw_blocklist_blocked_total{feed=\"$name\",ip_version=\"6\",period=\"1h\"} $blocks_1h_v6"
        echo "ufw_blocklist_blocked_total{feed=\"$name\",ip_version=\"6\",period=\"24h\"} $blocks_24h_v6"
    done

    # Calculate total blocks once for hit rate
    local total_blocks_24h=0
    for feed_line in "${FEEDS_ARRAY[@]}"; do
        IFS='|' read -r enabled name url type description <<< "$feed_line"
        [ "$enabled" != "1" ] && continue
        local b
        b=$(get_feed_blocks "$name" "24 hours ago")
        total_blocks_24h=$((total_blocks_24h + ${b:-0}))
    done

    cat <<EOF

# HELP ufw_blocklist_effectiveness Feed hit rate - percentage of total blocks from this feed (24h)
# TYPE ufw_blocklist_effectiveness gauge
EOF

    for feed_line in "${FEEDS_ARRAY[@]}"; do
        IFS='|' read -r enabled name url type description <<< "$feed_line"
        [ "$enabled" != "1" ] && continue
        
        blocks=$(get_feed_blocks "$name" "24 hours ago")
        blocks=$(echo "$blocks" | tr -d '\n' | tr -d ' ')
        blocks=${blocks:-0}
        
        if [ "${total_blocks_24h:-0}" -gt 0 ]; then
            effectiveness=$(awk "BEGIN {printf \"%.2f\", ($blocks / $total_blocks_24h) * 100}" 2>/dev/null || echo "0")
        else
            effectiveness="0"
        fi
        
        echo "ufw_blocklist_effectiveness{feed=\"$name\"} $effectiveness"
    done

    cat <<EOF

# HELP ufw_blocklist_ipset_memory_bytes Memory used by each ipset
# TYPE ufw_blocklist_ipset_memory_bytes gauge
EOF

    for feed_line in "${FEEDS_ARRAY[@]}"; do
        IFS='|' read -r enabled name url type description <<< "$feed_line"
        [ "$enabled" != "1" ] && continue
        
        mem_v4=$(get_ipset_memory "${IPSET_PREFIX}-${name}")
        mem_v6=$(get_ipset_memory "${IPSET_PREFIX}-${name}-v6")
        
        echo "ufw_blocklist_ipset_memory_bytes{feed=\"$name\",ip_version=\"4\"} $mem_v4"
        echo "ufw_blocklist_ipset_memory_bytes{feed=\"$name\",ip_version=\"6\"} $mem_v6"
    done

    # Conntrack metrics (system-wide)
    local conntrack_count conntrack_max conntrack_usage
    conntrack_count=$(get_conntrack_count)
    conntrack_max=$(get_conntrack_max)
    
    if [ "${conntrack_max:-0}" -gt 0 ]; then
        conntrack_usage=$(awk "BEGIN {printf \"%.2f\", ($conntrack_count / $conntrack_max) * 100}" 2>/dev/null || echo "0")
    else
        conntrack_usage="0"
    fi

    # Cache disk metrics
    local disk_info cache_size disk_used disk_avail disk_pct
    disk_info=$(get_cache_disk_usage)
    cache_size=$(get_total_cache_size)
    disk_used=$(echo "$disk_info" | cut -d'|' -f1)
    disk_avail=$(echo "$disk_info" | cut -d'|' -f2)
    disk_pct=$(echo "$disk_info" | cut -d'|' -f3 | tr -d '%')

    cat <<EOF

# HELP ufw_blocklist_conntrack_entries Current conntrack entries
# TYPE ufw_blocklist_conntrack_entries gauge
ufw_blocklist_conntrack_entries $conntrack_count

# HELP ufw_blocklist_conntrack_max Maximum conntrack entries
# TYPE ufw_blocklist_conntrack_max gauge
ufw_blocklist_conntrack_max $conntrack_max

# HELP ufw_blocklist_conntrack_usage_percent Conntrack table usage percentage
# TYPE ufw_blocklist_conntrack_usage_percent gauge
ufw_blocklist_conntrack_usage_percent $conntrack_usage

# HELP ufw_blocklist_cache_disk_used_bytes Disk space used by cache partition
# TYPE ufw_blocklist_cache_disk_used_bytes gauge
ufw_blocklist_cache_disk_used_bytes $disk_used

# HELP ufw_blocklist_cache_disk_available_bytes Disk space available on cache partition
# TYPE ufw_blocklist_cache_disk_available_bytes gauge
ufw_blocklist_cache_disk_available_bytes $disk_avail

# HELP ufw_blocklist_cache_disk_usage_percent Cache partition disk usage percentage
# TYPE ufw_blocklist_cache_disk_usage_percent gauge
ufw_blocklist_cache_disk_usage_percent ${disk_pct:-0}

# HELP ufw_blocklist_cache_total_size_bytes Total size of cache directory
# TYPE ufw_blocklist_cache_total_size_bytes gauge
ufw_blocklist_cache_total_size_bytes $cache_size

# HELP ufw_blocklist_whitelist_size Whitelist ipset size
# TYPE ufw_blocklist_whitelist_size gauge
ufw_blocklist_whitelist_size{ip_version="4"} $(get_ipset_size "$WHITELIST_IPSET")
ufw_blocklist_whitelist_size{ip_version="6"} $(get_ipset_size "$WHITELIST_IPSET_V6")

# HELP ufw_blocklist_total_unique_ips Total deduplicated IPs across all feeds
# TYPE ufw_blocklist_total_unique_ips gauge
ufw_blocklist_total_unique_ips{ip_version="4"} $(cat "$CACHE_DIR"/*-v4.parsed 2>/dev/null | sort -u | wc -l)
ufw_blocklist_total_unique_ips{ip_version="6"} $(cat "$CACHE_DIR"/*-v6.parsed 2>/dev/null | sort -u | wc -l)

# HELP ufw_blocklist_total_rules Total UFW firewall rules
# TYPE ufw_blocklist_total_rules gauge
ufw_blocklist_total_rules $(ufw status numbered 2>/dev/null | grep -c '^\[')

# HELP ufw_blocklist_scrape_timestamp_seconds Unix timestamp of metric generation
# TYPE ufw_blocklist_scrape_timestamp_seconds gauge
ufw_blocklist_scrape_timestamp_seconds $(date +%s)

# HELP ufw_blocklist_exporter_duration_seconds Time to generate all metrics
# TYPE ufw_blocklist_exporter_duration_seconds gauge
ufw_blocklist_exporter_duration_seconds $(($(date +%s) - start_time))
EOF

    echo ""
}

run_http_server() {
    echo "Starting exporter on port $HTTP_PORT..." >&2
    
    while true; do
        {
            read -r request
            if [[ "$request" =~ ^GET\ /metrics ]]; then
                printf "HTTP/1.1 200 OK\r\nContent-Type: text/plain; version=0.0.4; charset=utf-8\r\n\r\n"
                cache_journal_data
                cache_feeds_config
                generate_metrics
            else
                printf "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\n\r\n"
                echo "<h1>UFW Blocklist Exporter v${SCRIPT_VERSION}</h1><a href='/metrics'>Metrics</a>"
            fi
        } | nc -l -p "$HTTP_PORT" -q 1 2>/dev/null
    done
}

main() {
    parse_args "$@"
    
    # Prevent multiple instances (skip for HTTP mode as it should run continuously)
    [ "$HTTP_MODE" != true ] && acquire_lock
    
    if [ "$HTTP_MODE" = true ]; then
        run_http_server
    elif [ -n "$OUTPUT_FILE" ]; then
        # Cache data before generating metrics
        cache_journal_data
        cache_feeds_config
        
        # Ensure output directory exists
        mkdir -p "$(dirname "$OUTPUT_FILE")"
        
        # Create temp file in /tmp (not in node_exporter directory!)
        local temp_file
        temp_file=$(mktemp /tmp/ufw_metrics.XXXXXX)
        
        # Generate metrics to temp file
        generate_metrics > "$temp_file"
        
        # FORCE NEW INODE: Delete old file first, then move
        rm -f "$OUTPUT_FILE"
        
        # Move temp file to final location
        mv "$temp_file" "$OUTPUT_FILE"
        
        # Ensure node_exporter user can read it
        chmod 644 "$OUTPUT_FILE"
        
        # Force filesystem sync
        sync
    else
        cache_journal_data
        cache_feeds_config
        generate_metrics
    fi
}

main "$@"