#!/bin/bash
################################################################################
# Script Name: headscale-metrics-exporter.sh
# Version: 1.0
# Description: Prometheus textfile collector exporter for Headscale — node
#              status, user counts, route health, pre-auth key inventory,
#              and key expiry tracking
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - Headscale CLI installed and headscale service running
#   - jq for JSON parsing
#   - Root/sudo access (headscale CLI requires it)
#   - netcat (nc) for HTTP mode
#
# Usage:
#   sudo ./headscale-metrics-exporter.sh
#   sudo ./headscale-metrics-exporter.sh --http -p 9588
#   sudo ./headscale-metrics-exporter.sh --textfile
#
# Metrics exported:
#   headscale_nodes_connected              - Count of online nodes
#   headscale_nodes_registered             - Total registered nodes
#   headscale_nodes_online                 - Per-node online status (1/0)
#   headscale_users_total                  - Total users
#   headscale_preauth_keys_total           - Pre-auth keys by state
#   headscale_routes_total                 - Routes by status
#   headscale_routes_exit_nodes            - Exit nodes by status
#   headscale_node_key_expiry_seconds      - Unix timestamp of key expiry per node
#   headscale_exporter_duration_seconds    - Script execution time
#   headscale_exporter_last_run_timestamp  - Last run timestamp
#
# Configuration:
#   Default HTTP port: 9588
#   Textfile directory: /var/lib/node_exporter
#
################################################################################

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

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9588

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

Export Headscale control server 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: 9588)
    -o, --output    Output file path

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

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            --textfile) OUTPUT_FILE="$TEXTFILE_DIR/headscale.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
}

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

check_headscale() {
    if ! command -v headscale >/dev/null 2>&1; then
        echo "ERROR: headscale command not found" >&2
        return 1
    fi
    if ! command -v jq >/dev/null 2>&1; then
        echo "ERROR: jq command not found" >&2
        return 1
    fi
    if ! headscale nodes list -o json >/dev/null 2>&1; then
        echo "ERROR: headscale not responding" >&2
        return 1
    fi
    return 0
}

prom_escape() {
    local val="$1"
    val="${val//\\/\\\\}"
    val="${val//\"/\\\"}"
    val="${val//$'\n'/\\n}"
    echo "$val"
}

iso_to_unix() {
    local ts="$1"
    if [ -z "$ts" ] || [ "$ts" = "null" ]; then
        echo "0"
        return
    fi
    local unix_ts
    unix_ts=$(date -d "$ts" +%s 2>/dev/null)
    echo "${unix_ts:-0}"
}

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

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

    if ! check_headscale; then
        echo "# HELP headscale_nodes_registered Total number of registered nodes"
        echo "# TYPE headscale_nodes_registered gauge"
        echo "headscale_nodes_registered 0"
        return
    fi

    # ------------------------------------------------------------------
    # Collect data from headscale CLI
    # ------------------------------------------------------------------
    local nodes_json users_json routes_json
    nodes_json=$(headscale nodes list -o json 2>/dev/null)
    users_json=$(headscale users list -o json 2>/dev/null)
    routes_json=$(headscale routes list -o json 2>/dev/null)

    if [ -z "$nodes_json" ] || [ "$nodes_json" = "null" ]; then
        nodes_json="[]"
    fi
    if [ -z "$users_json" ] || [ "$users_json" = "null" ]; then
        users_json="[]"
    fi
    if [ -z "$routes_json" ] || [ "$routes_json" = "null" ]; then
        routes_json="[]"
    fi

    # ------------------------------------------------------------------
    # Node metrics
    # ------------------------------------------------------------------
    local total_nodes online_nodes
    total_nodes=$(echo "$nodes_json" | jq 'length')
    online_nodes=$(echo "$nodes_json" | jq '[.[] | select(.online == true)] | length')

    echo "# HELP headscale_nodes_connected Number of currently online nodes"
    echo "# TYPE headscale_nodes_connected gauge"
    echo "headscale_nodes_connected ${online_nodes:-0}"
    echo ""

    echo "# HELP headscale_nodes_registered Total number of registered nodes"
    echo "# TYPE headscale_nodes_registered gauge"
    echo "headscale_nodes_registered ${total_nodes:-0}"
    echo ""

    echo "# HELP headscale_nodes_online Per-node online status (1=online, 0=offline)"
    echo "# TYPE headscale_nodes_online gauge"

    local node_data
    node_data=$(echo "$nodes_json" | jq -r '
        .[] |
        [
            .givenName // .name // "",
            .user.name // "",
            (.ipAddresses[0]? // ""),
            (.online // false | tostring),
            .expiry // ""
        ] | @tsv
    ')

    declare -a node_names=()
    declare -a node_users=()
    declare -a node_expiries=()

    while IFS=$'\t' read -r name user ip online expiry; do
        [ -z "$name" ] && continue
        local escaped_name escaped_user escaped_ip
        escaped_name=$(prom_escape "$name")
        escaped_user=$(prom_escape "$user")
        escaped_ip=$(prom_escape "$ip")

        local online_val=0
        if [ "$online" = "true" ]; then
            online_val=1
        fi
        echo "headscale_nodes_online{node=\"$escaped_name\",user=\"$escaped_user\",ip=\"$escaped_ip\"} $online_val"

        node_names+=("$escaped_name")
        node_users+=("$escaped_user")
        node_expiries+=("$expiry")
    done <<< "$node_data"
    echo ""

    # ------------------------------------------------------------------
    # Key expiry metrics
    # ------------------------------------------------------------------
    echo "# HELP headscale_node_key_expiry_seconds Unix timestamp of node key expiry"
    echo "# TYPE headscale_node_key_expiry_seconds gauge"

    for i in "${!node_names[@]}"; do
        local expiry_unix
        expiry_unix=$(iso_to_unix "${node_expiries[$i]}")
        echo "headscale_node_key_expiry_seconds{node=\"${node_names[$i]}\",user=\"${node_users[$i]}\"} $expiry_unix"
    done
    echo ""

    # ------------------------------------------------------------------
    # User metrics
    # ------------------------------------------------------------------
    local total_users
    total_users=$(echo "$users_json" | jq 'length')

    echo "# HELP headscale_users_total Total number of users"
    echo "# TYPE headscale_users_total gauge"
    echo "headscale_users_total ${total_users:-0}"
    echo ""

    # ------------------------------------------------------------------
    # Pre-auth key metrics
    # ------------------------------------------------------------------
    local usable_keys=0
    local expired_keys=0
    local used_keys=0

    local user_list
    user_list=$(echo "$users_json" | jq -r '.[].name // empty')

    while IFS= read -r username; do
        [ -z "$username" ] && continue
        local keys_json
        keys_json=$(headscale preauthkeys list -u "$username" -o json 2>/dev/null)
        if [ -z "$keys_json" ] || [ "$keys_json" = "null" ]; then
            continue
        fi

        local u e s
        u=$(echo "$keys_json" | jq '[.[] | select(.used == false and .expiration > now)] | length' 2>/dev/null)
        e=$(echo "$keys_json" | jq '[.[] | select(.used == false and .expiration <= now)] | length' 2>/dev/null)
        s=$(echo "$keys_json" | jq '[.[] | select(.used == true)] | length' 2>/dev/null)

        usable_keys=$((usable_keys + ${u:-0}))
        expired_keys=$((expired_keys + ${e:-0}))
        used_keys=$((used_keys + ${s:-0}))
    done <<< "$user_list"

    echo "# HELP headscale_preauth_keys_total Pre-auth keys by state"
    echo "# TYPE headscale_preauth_keys_total gauge"
    echo "headscale_preauth_keys_total{state=\"usable\"} $usable_keys"
    echo "headscale_preauth_keys_total{state=\"expired\"} $expired_keys"
    echo "headscale_preauth_keys_total{state=\"used\"} $used_keys"
    echo ""

    # ------------------------------------------------------------------
    # Route metrics
    # ------------------------------------------------------------------
    local approved_routes pending_routes disabled_routes
    local approved_exit pending_exit disabled_exit

    approved_routes=$(echo "$routes_json" | jq '[.[] | select(.advertised == true and .enabled == true and (.prefix | test("^0\\.0\\.0\\.0/0$|^::/0$") | not))] | length')
    pending_routes=$(echo "$routes_json" | jq '[.[] | select(.advertised == true and .enabled == false and (.prefix | test("^0\\.0\\.0\\.0/0$|^::/0$") | not))] | length')
    disabled_routes=$(echo "$routes_json" | jq '[.[] | select(.advertised == false and (.prefix | test("^0\\.0\\.0\\.0/0$|^::/0$") | not))] | length')

    approved_exit=$(echo "$routes_json" | jq '[.[] | select(.advertised == true and .enabled == true and (.prefix | test("^0\\.0\\.0\\.0/0$|^::/0$")))] | length')
    pending_exit=$(echo "$routes_json" | jq '[.[] | select(.advertised == true and .enabled == false and (.prefix | test("^0\\.0\\.0\\.0/0$|^::/0$")))] | length')
    disabled_exit=$(echo "$routes_json" | jq '[.[] | select(.advertised == false and (.prefix | test("^0\\.0\\.0\\.0/0$|^::/0$")))] | length')

    echo "# HELP headscale_routes_total Routes by status"
    echo "# TYPE headscale_routes_total gauge"
    echo "headscale_routes_total{status=\"approved\"} ${approved_routes:-0}"
    echo "headscale_routes_total{status=\"pending\"} ${pending_routes:-0}"
    echo "headscale_routes_total{status=\"disabled\"} ${disabled_routes:-0}"
    echo ""

    echo "# HELP headscale_routes_exit_nodes Exit nodes by status"
    echo "# TYPE headscale_routes_exit_nodes gauge"
    echo "headscale_routes_exit_nodes{status=\"approved\"} ${approved_exit:-0}"
    echo "headscale_routes_exit_nodes{status=\"pending\"} ${pending_exit:-0}"
    echo "headscale_routes_exit_nodes{status=\"disabled\"} ${disabled_exit:-0}"
    echo ""

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

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

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

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

run_http_server() {
    echo "Starting Headscale 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>Headscale Exporter</title></head><body><h1>Headscale 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
# ==============================================================================

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}/.headscale_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 "$@"