#!/bin/bash
################################################################################
# Script Name: wireguard-exporter.sh
# Version: 1.0
# Description: Prometheus exporter for WireGuard VPN — interface status, peer
#              connections, transfer bytes, handshake health, and endpoint tracking
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - WireGuard installed and configured
#   - Root/sudo access for `wg show`
#   - netcat (nc) for HTTP mode
#
# Usage:
#   sudo ./wireguard-exporter.sh
#   sudo ./wireguard-exporter.sh --http -p 9586
#   sudo ./wireguard-exporter.sh --textfile
#
# Configuration:
#   Default HTTP port: 9586
#   Textfile directory: /var/lib/node_exporter
#
################################################################################

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9586
HANDSHAKE_TIMEOUT=180

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

Export WireGuard VPN statistics as Prometheus metrics.

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

OPTIONS:
    -p, --port      HTTP port (default: 9586)
    -o, --output    Output file path

EXAMPLES:
    sudo $0 --textfile
    sudo $0 --http --port 9586

EOF
    exit 0
}

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

check_wireguard() {
    if ! command -v wg >/dev/null 2>&1; then
        echo "ERROR: wg command not found" >&2
        return 1
    fi
    return 0
}

generate_metrics() {
    local script_start
    script_start=$(date +%s)

    if ! check_wireguard; then
        echo "# HELP wireguard_up WireGuard exporter status"
        echo "# TYPE wireguard_up gauge"
        echo "wireguard_up 0"
        return
    fi

    echo "# HELP wireguard_up WireGuard exporter status"
    echo "# TYPE wireguard_up gauge"
    echo "wireguard_up 1"
    echo ""

    local interfaces
    interfaces=$(wg show interfaces 2>/dev/null)

    if [ -z "$interfaces" ]; then
        echo "# HELP wireguard_interface_up Interface status"
        echo "# TYPE wireguard_interface_up gauge"
        echo "wireguard_interface_up 0"
        return
    fi

    echo "# HELP wireguard_interface_up Interface status (1=up, 0=down)"
    echo "# TYPE wireguard_interface_up gauge"
    echo "# HELP wireguard_interface_listening_port Listening port"
    echo "# TYPE wireguard_interface_listening_port gauge"
    echo "# HELP wireguard_peers_configured Total configured peers"
    echo "# TYPE wireguard_peers_configured gauge"
    echo "# HELP wireguard_peers_connected Connected peers (handshake within ${HANDSHAKE_TIMEOUT}s)"
    echo "# TYPE wireguard_peers_connected gauge"
    echo "# HELP wireguard_peer_connected Peer connection status"
    echo "# TYPE wireguard_peer_connected gauge"
    echo "# HELP wireguard_peer_transfer_received_bytes Total bytes received from peer"
    echo "# TYPE wireguard_peer_transfer_received_bytes counter"
    echo "# HELP wireguard_peer_transfer_transmitted_bytes Total bytes transmitted to peer"
    echo "# TYPE wireguard_peer_transfer_transmitted_bytes counter"
    echo "# HELP wireguard_peer_latest_handshake_seconds Unix timestamp of latest handshake"
    echo "# TYPE wireguard_peer_latest_handshake_seconds gauge"
    echo "# HELP wireguard_peer_handshake_age_seconds Seconds since last handshake"
    echo "# TYPE wireguard_peer_handshake_age_seconds gauge"
    echo "# HELP wireguard_peer_persistent_keepalive Persistent keepalive interval"
    echo "# TYPE wireguard_peer_persistent_keepalive gauge"

    local now
    now=$(date +%s)

    for iface in $interfaces; do
        local listening_port
        listening_port=$(wg show "$iface" listen-port 2>/dev/null)
        echo "wireguard_interface_up{interface=\"$iface\"} 1"
        echo "wireguard_interface_listening_port{interface=\"$iface\"} ${listening_port:-0}"

        local total_peers=0
        local connected_peers=0

        while IFS=$'\t' read -r public_key endpoint allowed_ips latest_handshake transfer_rx transfer_tx persistent_keepalive; do
            [ -z "$public_key" ] && continue
            [ "$public_key" = "(none)" ] && continue
            total_peers=$((total_peers + 1))

            # Short key for labels
            local short_key="${public_key:0:8}..."

            # Handshake age
            local handshake_age=0
            local is_connected=0
            if [ "$latest_handshake" -gt 0 ] 2>/dev/null; then
                handshake_age=$((now - latest_handshake))
                if [ "$handshake_age" -le "$HANDSHAKE_TIMEOUT" ]; then
                    is_connected=1
                    connected_peers=$((connected_peers + 1))
                fi
            fi

            echo "wireguard_peer_connected{interface=\"$iface\",public_key=\"$short_key\"} $is_connected"
            echo "wireguard_peer_transfer_received_bytes{interface=\"$iface\",public_key=\"$short_key\"} ${transfer_rx:-0}"
            echo "wireguard_peer_transfer_transmitted_bytes{interface=\"$iface\",public_key=\"$short_key\"} ${transfer_tx:-0}"
            echo "wireguard_peer_latest_handshake_seconds{interface=\"$iface\",public_key=\"$short_key\"} ${latest_handshake:-0}"
            echo "wireguard_peer_handshake_age_seconds{interface=\"$iface\",public_key=\"$short_key\"} $handshake_age"

            local keepalive=0
            [ "$persistent_keepalive" != "off" ] && keepalive="${persistent_keepalive:-0}"
            echo "wireguard_peer_persistent_keepalive{interface=\"$iface\",public_key=\"$short_key\"} $keepalive"

        done < <(wg show "$iface" dump 2>/dev/null | tail -n +2)

        echo "wireguard_peers_configured{interface=\"$iface\"} $total_peers"
        echo "wireguard_peers_connected{interface=\"$iface\"} $connected_peers"
    done

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

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

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

run_http_server() {
    echo "Starting WireGuard 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
    while true; do
        {
            read -r request
            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
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r"
                echo "<html><head><title>WireGuard Exporter</title></head><body><h1>WireGuard Prometheus Exporter</h1><p><a href=\"/metrics\">Metrics</a></p></body></html>"
            fi
        } | nc -l -p "$HTTP_PORT" -q 1 2>/dev/null
    done
}

main() {
    parse_args "$@"
    if [ "$HTTP_MODE" = true ]; then
        run_http_server
    elif [ -n "$OUTPUT_FILE" ]; then
        local output_dir
        output_dir="$(dirname "$OUTPUT_FILE")"
        mkdir -p "$output_dir"
        local temp_file
        temp_file=$(mktemp "${output_dir}/.wireguard_metrics.XXXXXX")
        if ! generate_metrics > "$temp_file" 2>/dev/null; then
            rm -f "$temp_file"
            echo "ERROR: Failed to generate metrics" >&2
            exit 1
        fi
        chmod 644 "$temp_file"
        mv -f "$temp_file" "$OUTPUT_FILE"
        echo "Metrics written to $OUTPUT_FILE" >&2
    else
        generate_metrics
    fi
}

main "$@"