#!/bin/bash
################################################################################
# Script Name: tailscale-exporter.sh
# Version: 1.0
# Description: Prometheus exporter for Tailscale VPN — node status, peer
#              connections, transfer bytes, relay tracking, and exit node status
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - Tailscale CLI installed and tailscaled running
#   - jq for JSON parsing
#   - Root/sudo access (some tailscale commands require it)
#   - netcat (nc) for HTTP mode
#
# Usage:
#   sudo ./tailscale-exporter.sh
#   sudo ./tailscale-exporter.sh --http -p 9587
#   sudo ./tailscale-exporter.sh --textfile
#
# Metrics exported:
#   tailscale_up                          - Tailscale daemon reachable (1/0)
#   tailscale_info                        - Self node info (version, hostname, tailnet)
#   tailscale_self_online                 - Whether this node is online (1/0)
#   tailscale_self_ip                     - This node's Tailscale IP addresses
#   tailscale_peers_configured            - Total peer count
#   tailscale_peers_online                - Currently online peers
#   tailscale_peer_online                 - Per-peer online status (1/0)
#   tailscale_peer_direct                 - Per-peer direct connection (1/0)
#   tailscale_peer_relay                  - Per-peer DERP relay usage (1 if relayed)
#   tailscale_peer_tx_bytes               - Per-peer transmitted bytes
#   tailscale_peer_rx_bytes               - Per-peer received bytes
#   tailscale_peer_last_seen_timestamp    - Per-peer last seen Unix timestamp
#   tailscale_exit_node_active            - Whether an exit node is in use (1/0)
#   tailscale_peers_by_os                 - Peer count by operating system
#   tailscale_exporter_duration_seconds   - Script execution time
#   tailscale_exporter_last_run_timestamp - Last run timestamp
#
# Configuration:
#   Default HTTP port: 9587
#   Textfile directory: /var/lib/node_exporter
#
################################################################################

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

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

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

Export Tailscale 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: 9587)
    -o, --output    Output file path

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

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            --textfile) OUTPUT_FILE="$TEXTFILE_DIR/tailscale.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_tailscale() {
    if ! command -v tailscale >/dev/null 2>&1; then
        echo "ERROR: tailscale 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 ! tailscale status >/dev/null 2>&1; then
        echo "ERROR: tailscaled 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_tailscale; then
        echo "# HELP tailscale_up Tailscale daemon reachable"
        echo "# TYPE tailscale_up gauge"
        echo "tailscale_up 0"
        return
    fi

    # Capture tailscale status JSON once
    local tmp_json
    tmp_json=$(mktemp /tmp/.tailscale_status.XXXXXX)
    if ! tailscale status --json > "$tmp_json" 2>/dev/null; then
        echo "# HELP tailscale_up Tailscale daemon reachable"
        echo "# TYPE tailscale_up gauge"
        echo "tailscale_up 0"
        rm -f "$tmp_json"
        return
    fi

    echo "# HELP tailscale_up Tailscale daemon reachable"
    echo "# TYPE tailscale_up gauge"
    echo "tailscale_up 1"
    echo ""

    # --- Self node info ---
    local self_hostname self_version tailnet_name self_online
    self_hostname=$(jq -r '.Self.HostName // ""' "$tmp_json")
    self_version=$(jq -r '.Version // ""' "$tmp_json")
    tailnet_name=$(jq -r '.MagicDNSSuffix // ""' "$tmp_json")
    self_online=$(jq -r '.Self.Online // false' "$tmp_json")

    echo "# HELP tailscale_info Tailscale node information"
    echo "# TYPE tailscale_info gauge"
    echo "tailscale_info{version=\"$(prom_escape "$self_version")\",hostname=\"$(prom_escape "$self_hostname")\",tailnet_name=\"$(prom_escape "$tailnet_name")\"} 1"
    echo ""

    echo "# HELP tailscale_self_online Whether this node is online"
    echo "# TYPE tailscale_self_online gauge"
    if [ "$self_online" = "true" ]; then
        echo "tailscale_self_online 1"
    else
        echo "tailscale_self_online 0"
    fi
    echo ""

    # --- Self IPs ---
    echo "# HELP tailscale_self_ip Tailscale IP address of this node"
    echo "# TYPE tailscale_self_ip gauge"
    local self_ips
    self_ips=$(jq -r '.Self.TailscaleIPs[]? // empty' "$tmp_json")
    while IFS= read -r ip; do
        [ -z "$ip" ] && continue
        echo "tailscale_self_ip{ip=\"$(prom_escape "$ip")\"} 1"
    done <<< "$self_ips"
    echo ""

    # --- Peer metrics ---
    local total_peers=0
    local online_peers=0
    local exit_node_active=0

    # Collect per-peer data via jq
    local peer_data
    peer_data=$(jq -r '
        .Peer // {} | to_entries[] |
        [
            .value.HostName // "",
            (.value.TailscaleIPs[0]? // ""),
            .value.OS // "",
            (.value.Online // false | tostring),
            .value.CurAddr // "",
            .value.Relay // "",
            (.value.TxBytes // 0 | tostring),
            (.value.RxBytes // 0 | tostring),
            .value.LastSeen // "",
            (.value.ExitNode // false | tostring)
        ] | @tsv
    ' "$tmp_json")

    # Declare associative array for OS counts
    declare -A os_counts

    echo "# HELP tailscale_peer_online Per-peer online status (1=online, 0=offline)"
    echo "# TYPE tailscale_peer_online gauge"
    echo "# HELP tailscale_peer_direct Per-peer direct connection (1=direct, 0=relayed)"
    echo "# TYPE tailscale_peer_direct gauge"
    echo "# HELP tailscale_peer_relay Per-peer DERP relay usage (1 if relayed)"
    echo "# TYPE tailscale_peer_relay gauge"
    echo "# HELP tailscale_peer_tx_bytes Per-peer transmitted bytes"
    echo "# TYPE tailscale_peer_tx_bytes counter"
    echo "# HELP tailscale_peer_rx_bytes Per-peer received bytes"
    echo "# TYPE tailscale_peer_rx_bytes counter"
    echo "# HELP tailscale_peer_last_seen_timestamp Per-peer last seen Unix timestamp"
    echo "# TYPE tailscale_peer_last_seen_timestamp gauge"

    while IFS=$'\t' read -r hostname tailscale_ip os online cur_addr relay tx_bytes rx_bytes last_seen exit_node; do
        [ -z "$hostname" ] && continue
        total_peers=$((total_peers + 1))

        local escaped_hostname escaped_ip escaped_os
        escaped_hostname=$(prom_escape "$hostname")
        escaped_ip=$(prom_escape "$tailscale_ip")
        escaped_os=$(prom_escape "$os")

        # Online status
        local online_val=0
        if [ "$online" = "true" ]; then
            online_val=1
            online_peers=$((online_peers + 1))
        fi
        echo "tailscale_peer_online{hostname=\"$escaped_hostname\",tailscale_ip=\"$escaped_ip\",os=\"$escaped_os\"} $online_val"

        # Direct connection
        local direct_val=0
        if [ -n "$cur_addr" ] && [[ ! "$cur_addr" =~ ^DERP ]]; then
            direct_val=1
        fi
        echo "tailscale_peer_direct{hostname=\"$escaped_hostname\"} $direct_val"

        # Relay usage
        if [ -n "$relay" ]; then
            echo "tailscale_peer_relay{hostname=\"$escaped_hostname\",relay=\"$(prom_escape "$relay")\"} 1"
        else
            echo "tailscale_peer_relay{hostname=\"$escaped_hostname\",relay=\"\"} 0"
        fi

        # Transfer bytes
        echo "tailscale_peer_tx_bytes{hostname=\"$escaped_hostname\"} ${tx_bytes:-0}"
        echo "tailscale_peer_rx_bytes{hostname=\"$escaped_hostname\"} ${rx_bytes:-0}"

        # Last seen timestamp
        local last_seen_unix
        last_seen_unix=$(iso_to_unix "$last_seen")
        echo "tailscale_peer_last_seen_timestamp{hostname=\"$escaped_hostname\"} $last_seen_unix"

        # Exit node check
        if [ "$exit_node" = "true" ]; then
            exit_node_active=1
        fi

        # OS aggregation
        if [ -n "$os" ]; then
            os_counts["$os"]=$(( ${os_counts["$os"]:-0} + 1 ))
        fi

    done <<< "$peer_data"

    echo ""
    echo "# HELP tailscale_peers_configured Total number of configured peers"
    echo "# TYPE tailscale_peers_configured gauge"
    echo "tailscale_peers_configured $total_peers"
    echo ""

    echo "# HELP tailscale_peers_online Number of currently online peers"
    echo "# TYPE tailscale_peers_online gauge"
    echo "tailscale_peers_online $online_peers"
    echo ""

    echo "# HELP tailscale_exit_node_active Whether an exit node is currently in use"
    echo "# TYPE tailscale_exit_node_active gauge"
    echo "tailscale_exit_node_active $exit_node_active"
    echo ""

    # --- Peers by OS ---
    echo "# HELP tailscale_peers_by_os Number of peers by operating system"
    echo "# TYPE tailscale_peers_by_os gauge"
    for os in "${!os_counts[@]}"; do
        echo "tailscale_peers_by_os{os=\"$(prom_escape "$os")\"} ${os_counts[$os]}"
    done
    echo ""

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

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

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

    rm -f "$tmp_json"
}

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

run_http_server() {
    echo "Starting Tailscale 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>Tailscale Exporter</title></head><body><h1>Tailscale 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}/.tailscale_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 "$@"