#!/bin/bash
################################################################################
# Script Name: keepalived-exporter.sh
# Version: 1.0
# Description: Prometheus exporter for Keepalived VRRP — instance state,
#              priority, VIP assignment, advertisement and transition counters
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - Keepalived installed and running
#   - Root/sudo access for sending signals and reading state files
#   - netcat (nc) for HTTP mode
#
# Usage:
#   sudo ./keepalived-exporter.sh
#   sudo ./keepalived-exporter.sh --http -p 9588
#   sudo ./keepalived-exporter.sh --textfile
#
# Configuration:
#   Default HTTP port: 9588
#   Textfile directory: /var/lib/node_exporter
#
# Metrics exported:
#   keepalived_up                                 Exporter status (1=up, 0=down)
#   keepalived_process_running                    Keepalived process status
#   keepalived_vrrp_state{instance}               VRRP state (2=MASTER, 1=BACKUP, 0=FAULT)
#   keepalived_vrrp_vip_status{instance,vip}      VIP assigned (1) or not (0)
#   keepalived_vrrp_priority{instance}            Configured base priority
#   keepalived_vrrp_effective_priority{instance}  Current effective priority
#   keepalived_vrrp_advert_interval{instance}     Advertisement interval in seconds
#   keepalived_vrrp_virtual_router_id{instance}   VRID
#   keepalived_vrrp_became_master_total{instance} Transitions to MASTER
#   keepalived_vrrp_released_master_total{instance} Transitions from MASTER
#   keepalived_vrrp_adverts_received_total{instance} Advertisements received
#   keepalived_vrrp_adverts_sent_total{instance}  Advertisements sent
#   keepalived_vrrp_garp_sent_total{instance}     Gratuitous ARPs sent
#   keepalived_vrrp_last_transition_timestamp{instance} Last state change timestamp
#   keepalived_exporter_duration_seconds          Script execution time
#   keepalived_exporter_last_run_timestamp        Last run timestamp
#
################################################################################

# ──────────────────────────────────────────────────────────────────────────────
# CONFIGURATION VARIABLES
# ──────────────────────────────────────────────────────────────────────────────

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

DATA_FILE="/tmp/keepalived.data"
STATS_FILE="/tmp/keepalived.stats"
SIGNAL_WAIT=0.5

# ──────────────────────────────────────────────────────────────────────────────
# HELPER FUNCTIONS
# ──────────────────────────────────────────────────────────────────────────────

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

Export Keepalived VRRP 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/keepalived.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
}

get_keepalived_pid() {
    local pid=""
    if [ -f /run/keepalived.pid ]; then
        pid=$(cat /run/keepalived.pid 2>/dev/null)
    elif [ -f /var/run/keepalived.pid ]; then
        pid=$(cat /var/run/keepalived.pid 2>/dev/null)
    fi
    if [ -z "$pid" ] || ! kill -0 "$pid" 2>/dev/null; then
        pid=$(pgrep -f "keepalived" | head -1)
    fi
    echo "$pid"
}

check_keepalived() {
    if pgrep keepalived >/dev/null 2>&1; then
        return 0
    fi
    if systemctl is-active keepalived >/dev/null 2>&1; then
        return 0
    fi
    return 1
}

map_state() {
    case "$1" in
        MASTER) echo 2 ;;
        BACKUP) echo 1 ;;
        FAULT)  echo 0 ;;
        INIT)   echo 0 ;;
        *)      echo 0 ;;
    esac
}

dump_keepalived_data() {
    local pid
    pid=$(get_keepalived_pid)
    [ -z "$pid" ] && return 1
    kill -USR1 "$pid" 2>/dev/null || return 1
    sleep "$SIGNAL_WAIT"
    [ -f "$DATA_FILE" ] || return 1
    return 0
}

dump_keepalived_stats() {
    local pid
    pid=$(get_keepalived_pid)
    [ -z "$pid" ] && return 1
    kill -USR2 "$pid" 2>/dev/null || return 1
    sleep "$SIGNAL_WAIT"
    [ -f "$STATS_FILE" ] || return 1
    return 0
}

check_vip_assigned() {
    local vip="$1"
    local ip_only="${vip%%/*}"
    if ip addr show 2>/dev/null | grep -qw "$ip_only"; then
        echo 1
    else
        echo 0
    fi
}

# ──────────────────────────────────────────────────────────────────────────────
# METRIC GENERATION
# ──────────────────────────────────────────────────────────────────────────────

parse_data_file() {
    [ -f "$DATA_FILE" ] || return
    local current_instance=""
    local in_vip_block=false

    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]]*VRRP\ Instance:\ (.+)$ ]]; then
            current_instance="${BASH_REMATCH[1]}"
            current_instance=$(echo "$current_instance" | xargs)
            in_vip_block=false
            continue
        fi
        [ -z "$current_instance" ] && continue

        if [[ "$line" =~ ^[[:space:]]*State\ =\ (.+)$ ]]; then
            local state="${BASH_REMATCH[1]}"
            state=$(echo "$state" | xargs)
            local state_val
            state_val=$(map_state "$state")
            echo "keepalived_vrrp_state{instance=\"$current_instance\"} $state_val"
            in_vip_block=false
        elif [[ "$line" =~ ^[[:space:]]*Wantstate\ =\ (.+)$ ]]; then
            in_vip_block=false
        elif [[ "$line" =~ ^[[:space:]]*Last\ transition\ =\ ([0-9]+) ]]; then
            local ts="${BASH_REMATCH[1]}"
            echo "keepalived_vrrp_last_transition_timestamp{instance=\"$current_instance\"} $ts"
            in_vip_block=false
        elif [[ "$line" =~ ^[[:space:]]*Base\ priority\ =\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_priority{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
            in_vip_block=false
        elif [[ "$line" =~ ^[[:space:]]*Effective\ priority\ =\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_effective_priority{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
            in_vip_block=false
        elif [[ "$line" =~ ^[[:space:]]*Advert\ interval\ =\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_advert_interval{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
            in_vip_block=false
        elif [[ "$line" =~ ^[[:space:]]*Virtual\ Router\ ID\ =\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_virtual_router_id{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
            in_vip_block=false
        elif [[ "$line" =~ ^[[:space:]]*Virtual\ IP ]]; then
            in_vip_block=true
        elif $in_vip_block && [[ "$line" =~ ^[[:space:]]+([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?) ]]; then
            local vip="${BASH_REMATCH[1]}"
            vip=$(echo "$vip" | xargs)
            local assigned
            assigned=$(check_vip_assigned "$vip")
            echo "keepalived_vrrp_vip_status{instance=\"$current_instance\",vip=\"$vip\"} $assigned"
        elif $in_vip_block && [[ ! "$line" =~ ^[[:space:]] ]]; then
            in_vip_block=false
        fi
    done < "$DATA_FILE"
}

parse_stats_file() {
    [ -f "$STATS_FILE" ] || return
    local current_instance=""
    local in_adverts=false

    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]]*VRRP\ Instance:\ (.+)$ ]]; then
            current_instance="${BASH_REMATCH[1]}"
            current_instance=$(echo "$current_instance" | xargs)
            in_adverts=false
            continue
        fi
        [ -z "$current_instance" ] && continue

        if [[ "$line" =~ ^[[:space:]]*Advertisements: ]]; then
            in_adverts=true
        elif $in_adverts && [[ "$line" =~ ^[[:space:]]*Received:\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_adverts_received_total{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
        elif $in_adverts && [[ "$line" =~ ^[[:space:]]*Sent:\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_adverts_sent_total{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
            in_adverts=false
        elif [[ "$line" =~ ^[[:space:]]*Became\ master:\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_became_master_total{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
        elif [[ "$line" =~ ^[[:space:]]*Released\ master:\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_released_master_total{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
        elif [[ "$line" =~ ^[[:space:]]*Gratuitous\ ARP:\ ([0-9]+) ]]; then
            echo "keepalived_vrrp_garp_sent_total{instance=\"$current_instance\"} ${BASH_REMATCH[1]}"
        fi
    done < "$STATS_FILE"
}

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

    if ! check_keepalived; then
        echo "# HELP keepalived_up Keepalived exporter status"
        echo "# TYPE keepalived_up gauge"
        echo "keepalived_up 0"
        echo ""
        echo "# HELP keepalived_process_running Keepalived process status (1=running, 0=stopped)"
        echo "# TYPE keepalived_process_running gauge"
        echo "keepalived_process_running 0"
        return
    fi

    echo "# HELP keepalived_up Keepalived exporter status"
    echo "# TYPE keepalived_up gauge"
    echo "keepalived_up 1"
    echo ""
    echo "# HELP keepalived_process_running Keepalived process status (1=running, 0=stopped)"
    echo "# TYPE keepalived_process_running gauge"
    echo "keepalived_process_running 1"
    echo ""

    # Signal keepalived to dump data and stats
    local has_data=false
    local has_stats=false
    if dump_keepalived_data; then
        has_data=true
    fi
    if dump_keepalived_stats; then
        has_stats=true
    fi

    if ! $has_data && ! $has_stats; then
        echo "# No VRRP data available — could not signal keepalived or read dump files"
        echo ""
        local script_end script_duration
        script_end=$(date +%s)
        script_duration=$((script_end - script_start))
        cat <<EOF
# HELP keepalived_exporter_duration_seconds Time to generate all metrics
# TYPE keepalived_exporter_duration_seconds gauge
keepalived_exporter_duration_seconds $script_duration

# HELP keepalived_exporter_last_run_timestamp Unix timestamp of last successful run
# TYPE keepalived_exporter_last_run_timestamp gauge
keepalived_exporter_last_run_timestamp $script_end
EOF
        return
    fi

    # Emit HELP/TYPE headers for data-file metrics
    if $has_data; then
        echo "# HELP keepalived_vrrp_state VRRP instance state (2=MASTER, 1=BACKUP, 0=FAULT/INIT)"
        echo "# TYPE keepalived_vrrp_state gauge"
        echo "# HELP keepalived_vrrp_vip_status VIP assigned on local interface (1=yes, 0=no)"
        echo "# TYPE keepalived_vrrp_vip_status gauge"
        echo "# HELP keepalived_vrrp_priority Configured base priority"
        echo "# TYPE keepalived_vrrp_priority gauge"
        echo "# HELP keepalived_vrrp_effective_priority Current effective priority"
        echo "# TYPE keepalived_vrrp_effective_priority gauge"
        echo "# HELP keepalived_vrrp_advert_interval Advertisement interval in seconds"
        echo "# TYPE keepalived_vrrp_advert_interval gauge"
        echo "# HELP keepalived_vrrp_virtual_router_id Virtual Router ID"
        echo "# TYPE keepalived_vrrp_virtual_router_id gauge"
        echo "# HELP keepalived_vrrp_last_transition_timestamp Unix timestamp of last state transition"
        echo "# TYPE keepalived_vrrp_last_transition_timestamp gauge"
        parse_data_file
        echo ""
    fi

    # Emit HELP/TYPE headers for stats-file metrics
    if $has_stats; then
        echo "# HELP keepalived_vrrp_became_master_total Transitions to MASTER state"
        echo "# TYPE keepalived_vrrp_became_master_total counter"
        echo "# HELP keepalived_vrrp_released_master_total Transitions from MASTER state"
        echo "# TYPE keepalived_vrrp_released_master_total counter"
        echo "# HELP keepalived_vrrp_adverts_received_total VRRP advertisements received"
        echo "# TYPE keepalived_vrrp_adverts_received_total counter"
        echo "# HELP keepalived_vrrp_adverts_sent_total VRRP advertisements sent"
        echo "# TYPE keepalived_vrrp_adverts_sent_total counter"
        echo "# HELP keepalived_vrrp_garp_sent_total Gratuitous ARPs sent"
        echo "# TYPE keepalived_vrrp_garp_sent_total counter"
        parse_stats_file
        echo ""
    fi

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

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

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

# ──────────────────────────────────────────────────────────────────────────────
# HTTP SERVER
# ──────────────────────────────────────────────────────────────────────────────

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