#!/bin/bash
################################################################################
# Script Name: wazuh-exporter.sh
# Version: 1.0
# Description: Prometheus exporter for Wazuh SIEM providing operational
#              metrics via the Wazuh API — agent status, manager daemon health,
#              active alerts, vulnerability counts, rule and decoder totals,
#              syscheck events, and API health
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - Wazuh Manager installed and running
#   - Wazuh API accessible (default: https://localhost:55000)
#   - curl for API calls
#   - jq for JSON parsing
#   - netcat (nc) for HTTP mode
#
# Usage:
#   # Output to stdout
#   ./wazuh-exporter.sh
#
#   # HTTP server mode
#   ./wazuh-exporter.sh --http -p 9193
#
#   # Textfile collector mode
#   ./wazuh-exporter.sh --textfile
#
#   # Custom API credentials
#   ./wazuh-exporter.sh --api-user admin --api-pass secret
#
# Metrics Exported:
#   - wazuh_up - API reachability (1=up, 0=down)
#   - wazuh_info{version,cluster_name} - Wazuh version info
#   - wazuh_agents_total - Total agent count
#   - wazuh_agents_by_status{status} - Agents by status
#   - wazuh_agents_by_os{os} - Agents by OS platform
#   - wazuh_manager_status{daemon} - Per-daemon status (1=running, 0=stopped)
#   - wazuh_manager_log_errors - Error log entry count
#   - wazuh_manager_log_warnings - Warning log entry count
#   - wazuh_rules_total - Total rules loaded
#   - wazuh_decoders_total - Total decoders loaded
#   - wazuh_alerts_total_today - Total alerts today
#   - wazuh_syscheck_events_today - Syscheck (FIM) events today
#   - wazuh_vulnerability_detector_critical - Critical vulnerabilities
#   - wazuh_vulnerability_detector_high - High severity vulnerabilities
#   - wazuh_exporter_duration_seconds - Script execution time
#   - wazuh_exporter_last_run_timestamp - Last run timestamp
#
# Configuration:
#   Default HTTP port: 9193
#   Default API URL: https://localhost:55000
#   Default API user: wazuh
#   Default API pass: wazuh
#   Textfile directory: /var/lib/node_exporter
#
################################################################################

set -euo pipefail

# ============================================================================
# CONFIGURATION VARIABLES
# ============================================================================

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9193
API_URL="https://localhost:55000"
API_USER="wazuh"
API_PASS="wazuh"

# JWT token cache (populated on first API call)
JWT_TOKEN=""

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

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

Export Wazuh SIEM statistics as Prometheus metrics via the Wazuh API.

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

OPTIONS:
    -p, --port        HTTP port (default: 9193)
    -o, --output      Output file path
    --api-url         Wazuh API URL (default: https://localhost:55000)
    --api-user        API username (default: wazuh)
    --api-pass        API password (default: wazuh)

EXAMPLES:
    $0 --textfile                          # Write to textfile collector
    $0 --http --port 9193                  # Run HTTP server
    $0 --api-user admin --api-pass secret  # Custom credentials
    $0 -o /tmp/wazuh.prom                  # Write to custom file

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            --textfile) OUTPUT_FILE="$TEXTFILE_DIR/wazuh.prom"; shift ;;
            --http) HTTP_MODE=true; shift ;;
            -p|--port) HTTP_PORT="$2"; shift 2 ;;
            -o|--output) OUTPUT_FILE="$2"; shift 2 ;;
            --api-url) API_URL="$2"; shift 2 ;;
            --api-user) API_USER="$2"; shift 2 ;;
            --api-pass) API_PASS="$2"; shift 2 ;;
            *) echo "Unknown option: $1" >&2; exit 1 ;;
        esac
    done
}

# Check prerequisites
# Returns: 0 if OK, 1 if error
check_prerequisites() {
    if ! command -v curl >/dev/null 2>&1; then
        echo "ERROR: curl not found" >&2
        return 1
    fi

    if ! command -v jq >/dev/null 2>&1; then
        echo "ERROR: jq not found (required for JSON parsing)" >&2
        return 1
    fi

    return 0
}

# Escape special characters in Prometheus label values
# Args: $1 - string to escape
# Returns: escaped string safe for Prometheus labels
prom_escape() {
    local val="$1"
    val="${val//\\/\\\\}"
    val="${val//\"/\\\"}"
    val="${val//$'\n'/}"
    echo "$val"
}

# Authenticate against Wazuh API and cache JWT token
# Returns: 0 if authenticated, 1 if error
authenticate() {
    if [ -n "$JWT_TOKEN" ]; then
        return 0
    fi

    local response
    response=$(curl -s -k -X POST \
        -u "${API_USER}:${API_PASS}" \
        "${API_URL}/security/user/authenticate" 2>/dev/null) || return 1

    JWT_TOKEN=$(echo "$response" | jq -r '.data.token // empty' 2>/dev/null)

    if [ -z "$JWT_TOKEN" ]; then
        echo "ERROR: Failed to authenticate with Wazuh API" >&2
        return 1
    fi

    return 0
}

# Make an authenticated API call
# Args: $1 - API endpoint path (e.g., /agents/summary/status)
# Returns: JSON response on stdout
api_call() {
    local endpoint="$1"
    curl -s -k -X GET \
        -H "Authorization: Bearer ${JWT_TOKEN}" \
        "${API_URL}${endpoint}" 2>/dev/null
}

# ============================================================================
# METRIC GENERATION
# ============================================================================

# Generate all Prometheus metrics
# Returns: Prometheus text format metrics on stdout
generate_metrics() {
    local script_start
    script_start=$(date +%s)

    # Check prerequisites
    if ! check_prerequisites; then
        cat <<EOF
# HELP wazuh_up Wazuh API reachability (1=up, 0=down)
# TYPE wazuh_up gauge
wazuh_up 0
EOF
        return
    fi

    # Authenticate
    JWT_TOKEN=""
    if ! authenticate; then
        cat <<EOF
# HELP wazuh_up Wazuh API reachability (1=up, 0=down)
# TYPE wazuh_up gauge
wazuh_up 0
EOF
        return
    fi

    # Verify API is responding
    local root_response
    root_response=$(api_call "/" 2>/dev/null)

    if [ -z "$root_response" ]; then
        cat <<EOF
# HELP wazuh_up Wazuh API reachability (1=up, 0=down)
# TYPE wazuh_up gauge
wazuh_up 0
EOF
        return
    fi

    cat <<EOF
# HELP wazuh_up Wazuh API reachability (1=up, 0=down)
# TYPE wazuh_up gauge
wazuh_up 1
EOF

    echo ""

    # ========================================================================
    # VERSION INFO
    # ========================================================================

    local wazuh_version cluster_name
    wazuh_version=$(echo "$root_response" | jq -r '.data.api_version // "unknown"' 2>/dev/null)
    cluster_name=$(echo "$root_response" | jq -r '.data.cluster_name // "unknown"' 2>/dev/null)

    cat <<EOF
# HELP wazuh_info Wazuh version information
# TYPE wazuh_info gauge
wazuh_info{version="$(prom_escape "$wazuh_version")",cluster_name="$(prom_escape "$cluster_name")"} 1
EOF

    echo ""

    # ========================================================================
    # AGENT METRICS
    # ========================================================================

    local agents_summary
    agents_summary=$(api_call "/agents/summary/status")

    if [ -n "$agents_summary" ] && [ "$agents_summary" != "null" ]; then
        local active disconnected never_connected pending total
        active=$(echo "$agents_summary" | jq -r '.data.connection.active // 0' 2>/dev/null)
        disconnected=$(echo "$agents_summary" | jq -r '.data.connection.disconnected // 0' 2>/dev/null)
        never_connected=$(echo "$agents_summary" | jq -r '.data.connection.never_connected // 0' 2>/dev/null)
        pending=$(echo "$agents_summary" | jq -r '.data.connection.pending // 0' 2>/dev/null)
        total=$(echo "$agents_summary" | jq -r '.data.connection.total // 0' 2>/dev/null)

        cat <<EOF
# HELP wazuh_agents_total Total number of agents
# TYPE wazuh_agents_total gauge
wazuh_agents_total $total

# HELP wazuh_agents_by_status Agents by connection status
# TYPE wazuh_agents_by_status gauge
wazuh_agents_by_status{status="active"} $active
wazuh_agents_by_status{status="disconnected"} $disconnected
wazuh_agents_by_status{status="never_connected"} $never_connected
wazuh_agents_by_status{status="pending"} $pending
EOF
    else
        cat <<EOF
# HELP wazuh_agents_total Total number of agents
# TYPE wazuh_agents_total gauge
wazuh_agents_total 0

# HELP wazuh_agents_by_status Agents by connection status
# TYPE wazuh_agents_by_status gauge
wazuh_agents_by_status{status="active"} 0
wazuh_agents_by_status{status="disconnected"} 0
wazuh_agents_by_status{status="never_connected"} 0
wazuh_agents_by_status{status="pending"} 0
EOF
    fi

    echo ""

    # Agents by OS
    local agents_os
    agents_os=$(api_call "/agents/summary/os")

    cat <<EOF
# HELP wazuh_agents_by_os Agents by OS platform
# TYPE wazuh_agents_by_os gauge
EOF

    if [ -n "$agents_os" ] && [ "$agents_os" != "null" ]; then
        echo "$agents_os" | jq -r '
            .data.affected_items // [] | .[] |
            "\(.os // "unknown") \(.count // 0)"
        ' 2>/dev/null | while read -r os count; do
            [ -z "$os" ] && continue
            echo "wazuh_agents_by_os{os=\"$(prom_escape "$os")\"} $count"
        done
    fi

    echo ""

    # ========================================================================
    # MANAGER METRICS
    # ========================================================================

    local manager_status
    manager_status=$(api_call "/manager/status")

    cat <<EOF
# HELP wazuh_manager_status Manager daemon status (1=running, 0=stopped)
# TYPE wazuh_manager_status gauge
EOF

    if [ -n "$manager_status" ] && [ "$manager_status" != "null" ]; then
        echo "$manager_status" | jq -r '
            .data.affected_items[0] // {} | to_entries[] |
            "\(.key) \(.value)"
        ' 2>/dev/null | while read -r daemon status; do
            [ -z "$daemon" ] && continue
            local val=0
            if [ "$status" = "running" ]; then
                val=1
            fi
            echo "wazuh_manager_status{daemon=\"$(prom_escape "$daemon")\"} $val"
        done
    fi

    echo ""

    # Manager logs summary
    local logs_summary
    logs_summary=$(api_call "/manager/logs/summary")

    local total_errors=0
    local total_warnings=0

    if [ -n "$logs_summary" ] && [ "$logs_summary" != "null" ]; then
        total_errors=$(echo "$logs_summary" | jq '[.data.affected_items[0] // {} | to_entries[] | .value.error // 0] | add // 0' 2>/dev/null)
        total_warnings=$(echo "$logs_summary" | jq '[.data.affected_items[0] // {} | to_entries[] | .value.warning // 0] | add // 0' 2>/dev/null)
        total_errors=${total_errors:-0}
        total_warnings=${total_warnings:-0}
    fi

    cat <<EOF
# HELP wazuh_manager_log_errors Total error log entries
# TYPE wazuh_manager_log_errors gauge
wazuh_manager_log_errors $total_errors

# HELP wazuh_manager_log_warnings Total warning log entries
# TYPE wazuh_manager_log_warnings gauge
wazuh_manager_log_warnings $total_warnings
EOF

    echo ""

    # ========================================================================
    # RULES METRICS
    # ========================================================================

    local rules_response
    rules_response=$(api_call "/rules?limit=1")

    local total_rules=0
    if [ -n "$rules_response" ] && [ "$rules_response" != "null" ]; then
        total_rules=$(echo "$rules_response" | jq '.data.total_affected_items // 0' 2>/dev/null)
        total_rules=${total_rules:-0}
    fi

    cat <<EOF
# HELP wazuh_rules_total Total rules loaded
# TYPE wazuh_rules_total gauge
wazuh_rules_total $total_rules
EOF

    echo ""

    # ========================================================================
    # DECODER METRICS
    # ========================================================================

    local decoders_response
    decoders_response=$(api_call "/decoders?limit=1")

    local total_decoders=0
    if [ -n "$decoders_response" ] && [ "$decoders_response" != "null" ]; then
        total_decoders=$(echo "$decoders_response" | jq '.data.total_affected_items // 0' 2>/dev/null)
        total_decoders=${total_decoders:-0}
    fi

    cat <<EOF
# HELP wazuh_decoders_total Total decoders loaded
# TYPE wazuh_decoders_total gauge
wazuh_decoders_total $total_decoders
EOF

    echo ""

    # ========================================================================
    # ALERT / STATS METRICS
    # ========================================================================

    local stats_response
    stats_response=$(api_call "/manager/stats")

    local total_alerts_today=0
    local total_syscheck_today=0

    if [ -n "$stats_response" ] && [ "$stats_response" != "null" ]; then
        total_alerts_today=$(echo "$stats_response" | jq '[.data.affected_items[] | .alerts[] | .sigid // 0] | length // 0' 2>/dev/null)
        total_syscheck_today=$(echo "$stats_response" | jq '[.data.affected_items[] | .syscheck // 0] | add // 0' 2>/dev/null)
        total_alerts_today=${total_alerts_today:-0}
        total_syscheck_today=${total_syscheck_today:-0}
    fi

    cat <<EOF
# HELP wazuh_alerts_total_today Total alerts today
# TYPE wazuh_alerts_total_today gauge
wazuh_alerts_total_today $total_alerts_today

# HELP wazuh_syscheck_events_today Total syscheck (FIM) events today
# TYPE wazuh_syscheck_events_today gauge
wazuh_syscheck_events_today $total_syscheck_today
EOF

    echo ""

    # ========================================================================
    # VULNERABILITY METRICS
    # ========================================================================

    local vuln_critical_response vuln_high_response
    vuln_critical_response=$(api_call "/vulnerability?severity=Critical&limit=1")
    vuln_high_response=$(api_call "/vulnerability?severity=High&limit=1")

    local critical_vulns=0
    local high_vulns=0

    if [ -n "$vuln_critical_response" ] && [ "$vuln_critical_response" != "null" ]; then
        critical_vulns=$(echo "$vuln_critical_response" | jq '.data.total_affected_items // 0' 2>/dev/null)
        critical_vulns=${critical_vulns:-0}
    fi

    if [ -n "$vuln_high_response" ] && [ "$vuln_high_response" != "null" ]; then
        high_vulns=$(echo "$vuln_high_response" | jq '.data.total_affected_items // 0' 2>/dev/null)
        high_vulns=${high_vulns:-0}
    fi

    cat <<EOF
# HELP wazuh_vulnerability_detector_critical Critical severity vulnerabilities
# TYPE wazuh_vulnerability_detector_critical gauge
wazuh_vulnerability_detector_critical $critical_vulns

# HELP wazuh_vulnerability_detector_high High severity vulnerabilities
# TYPE wazuh_vulnerability_detector_high gauge
wazuh_vulnerability_detector_high $high_vulns
EOF

    echo ""

    # ========================================================================
    # EXPORTER RUNTIME
    # ========================================================================

    local script_end script_duration
    script_end=$(date +%s)
    script_duration=$((script_end - script_start))

    cat <<EOF
# HELP wazuh_exporter_duration_seconds Time to generate all metrics
# TYPE wazuh_exporter_duration_seconds gauge
wazuh_exporter_duration_seconds $script_duration

# HELP wazuh_exporter_last_run_timestamp Unix timestamp of last successful run
# TYPE wazuh_exporter_last_run_timestamp gauge
wazuh_exporter_last_run_timestamp $script_end
EOF

    echo ""
}

# ============================================================================
# HTTP SERVER MODE
# ============================================================================

# Run simple HTTP server using netcat
# Serves metrics on /metrics endpoint
run_http_server() {
    echo "Starting Wazuh exporter on port $HTTP_PORT..." >&2

    if ! command -v nc >/dev/null 2>&1; then
        echo "ERROR: netcat (nc) required for HTTP mode" >&2
        exit 1
    fi

    # Infinite loop accepting HTTP requests
    while true; do
        {
            read -r request
            # Check if request is for /metrics endpoint
            if [[ "$request" =~ ^GET\ /metrics ]]; then
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/plain; version=0.0.4\r\n\r"
                generate_metrics
            else  # Serve HTML landing page for other requests
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r"
                cat <<EOF
<!DOCTYPE html>
<html>
<head><title>Wazuh Exporter v1.0</title></head>
<body>
<h1>Wazuh Prometheus Exporter v1.0</h1>
<p><a href="/metrics">Metrics</a></p>
<p>Operational metrics from the Wazuh API.</p>
</body>
</html>
EOF
            fi
        } | nc -l -p "$HTTP_PORT" -q 1 2>/dev/null
    done
}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

# Main entry point - routes to appropriate output mode
main() {
    parse_args "$@"

    if [ "$HTTP_MODE" = true ]; then
        # Run HTTP server (blocks until killed)
        run_http_server
    elif [ -n "$OUTPUT_FILE" ]; then
        # Textfile collector mode: write atomically using temp file
        local output_dir
        output_dir="$(dirname "$OUTPUT_FILE")"
        mkdir -p "$output_dir"

        # Create temp file in SAME directory for atomic rename (same filesystem)
        local temp_file
        temp_file=$(mktemp "${output_dir}/.wazuh_metrics.XXXXXX")

        # Generate metrics to temp file
        if ! generate_metrics > "$temp_file" 2>/dev/null; then
            rm -f "$temp_file"
            echo "ERROR: Failed to generate metrics" >&2
            exit 1
        fi

        # Validate: file must exist, have content
        local file_lines
        file_lines=$(wc -l < "$temp_file" 2>/dev/null || echo 0)

        if [ "$file_lines" -lt 10 ]; then
            rm -f "$temp_file"
            echo "ERROR: Metrics file too small ($file_lines lines), keeping previous" >&2
            exit 1
        fi

        # Set permissions before move
        chmod 644 "$temp_file"

        # Atomic rename - no gap where file is missing
        mv -f "$temp_file" "$OUTPUT_FILE"

        echo "Metrics written to $OUTPUT_FILE ($file_lines lines)" >&2
    else
        # Default: output to stdout
        generate_metrics
    fi
}

# Execute main function with all script arguments
main "$@"