#!/bin/bash
################################################################################
# Script Name: samba-exporter.sh
# Version: 1.0
# Description: Prometheus exporter for Samba metrics -- active sessions, share
#              connections, locked files, process status, version info, and
#              configured share counts
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - Samba installed (smbstatus, testparm, smbd)
#   - Root or appropriate permissions for smbstatus
#   - netcat (nc) for HTTP mode
#
# Usage:
#   ./samba-exporter.sh
#   ./samba-exporter.sh --http -p 9588
#   ./samba-exporter.sh --textfile
#
# Metrics Exported:
#   - samba_up - Exporter status (1=up, 0=down)
#   - samba_server_info{version} - Samba version info metric
#   - samba_sessions_active - Number of active SMB sessions
#   - samba_session_info{user,machine,protocol_version} - Per-session info
#   - samba_shares_active - Number of active share connections
#   - samba_share_connections{share,machine,user} - Connections per share
#   - samba_share_configured_total - Total configured shares from testparm
#   - samba_locked_files_total - Total number of locked files
#   - samba_smbd_running - Whether smbd process is running (1/0)
#   - samba_nmbd_running - Whether nmbd process is running (1/0)
#   - samba_winbindd_running - Whether winbindd process is running (1/0)
#   - samba_smbd_pid_count - Number of smbd processes
#   - samba_exporter_duration_seconds - Script execution time
#   - samba_exporter_last_run_timestamp - Last successful run time
#
# Configuration:
#   Default HTTP port: 9588
#   Textfile directory: /var/lib/node_exporter
#
################################################################################

set -o pipefail

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

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9588
LOCK_FILE="/tmp/samba-exporter.lock"

# ============================================================================
# LOGGING FUNCTIONS
# ============================================================================

RED='\033[0;31m'
YELLOW='\033[1;33m'
GREEN='\033[0;32m'
NC='\033[0m'

log()  { echo -e "${GREEN}[INFO]${NC} $*" >&2; }
warn() { echo -e "${YELLOW}[WARN]${NC} $*" >&2; }
error(){ echo -e "${RED}[ERROR]${NC} $*" >&2; }

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

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

Export Samba 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:
    $0 --textfile
    $0 --http --port 9588
    $0 -o /tmp/samba.prom

EOF
    exit 0
}

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

acquire_lock() {
    if [ -f "$LOCK_FILE" ]; then
        local lock_pid
        lock_pid=$(cat "$LOCK_FILE" 2>/dev/null)
        if [ -n "$lock_pid" ] && kill -0 "$lock_pid" 2>/dev/null; then
            error "Another instance is running (PID $lock_pid)"
            return 1
        fi
        rm -f "$LOCK_FILE"
    fi
    echo $$ > "$LOCK_FILE"
    trap 'rm -f "$LOCK_FILE"' EXIT
    return 0
}

# Check if Samba tools are available
check_samba() {
    if ! command -v smbstatus >/dev/null 2>&1; then
        error "smbstatus not found"
        return 1
    fi
    if ! smbstatus -b >/dev/null 2>&1; then
        error "smbstatus failed (insufficient permissions or smbd not running)"
        return 1
    fi
    return 0
}

# Get Samba version string from smbd
# Returns: Version string (e.g., "4.18.6")
get_samba_version() {
    local version
    if command -v smbd >/dev/null 2>&1; then
        version=$(smbd --version 2>/dev/null | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1)
    fi
    if [ -z "$version" ]; then
        version=$(smbstatus --version 2>/dev/null | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1)
    fi
    echo "${version:-unknown}"
}

# Get active session data from smbstatus --processes
# Returns: Lines with PID, username, group, machine, protocol_version
get_session_data() {
    smbstatus --processes --numeric 2>/dev/null | awk '
        BEGIN { started=0 }
        /^----/ { started=1; next }
        started && NF >= 4 && $1 ~ /^[0-9]+/ {
            pid=$1
            user=$2
            group=$3
            machine=$4
            proto=""
            # Protocol version is typically the last field or near end
            for (i=5; i<=NF; i++) {
                if ($i ~ /^SMB[0-9]/ || $i ~ /^NT1/ || $i ~ /^SMB[23]_[0-9]/) {
                    proto=$i
                }
            }
            if (proto == "") proto="unknown"
            printf "%s|%s|%s|%s|%s\n", pid, user, group, machine, proto
        }
    '
}

# Get active share connection data from smbstatus --shares
# Returns: Lines with share name, machine, user, connection time
get_share_data() {
    smbstatus --shares --numeric 2>/dev/null | awk '
        BEGIN { started=0 }
        /^----/ { started=1; next }
        started && NF >= 4 && $1 ~ /^[0-9]+/ {
            # Fields: Service pid machine Connected_at ...
            share=$1
            pid=$2
            machine=$3
            # Look up from all remaining fields
            printf "%s|%s|%s\n", share, pid, machine
        }
    '
}

# Get number of locked files from smbstatus --locks
# Returns: Count of locked files
get_locked_files_count() {
    local count
    count=$(smbstatus --locks --numeric 2>/dev/null | awk '
        BEGIN { started=0; count=0 }
        /^----/ { started=1; next }
        started && NF >= 2 && $1 ~ /^[0-9]+/ { count++ }
        END { print count }
    ')
    echo "${count:-0}"
}

# Get total configured shares from testparm
# Returns: Number of configured shares (excluding global, printers meta-sections)
get_configured_shares() {
    local count
    if command -v testparm >/dev/null 2>&1; then
        count=$(testparm -s 2>/dev/null | grep -cE '^\[' | head -1)
        # Subtract 1 for [global] section
        if [ -n "$count" ] && [ "$count" -gt 0 ]; then
            # Count actual share sections, excluding [global]
            count=$(testparm -s 2>/dev/null | grep -E '^\[' | grep -vcE '^\[global\]')
        fi
    fi
    echo "${count:-0}"
}

# Check if a process is running
# Args: $1 - process name
# Returns: 1 if running, 0 if not
check_process_running() {
    local proc_name="$1"
    if pgrep -x "$proc_name" >/dev/null 2>&1; then
        echo "1"
    else
        echo "0"
    fi
}

# Get count of smbd processes
# Returns: Number of smbd processes
get_smbd_pid_count() {
    local count
    count=$(pgrep -cx smbd 2>/dev/null)
    echo "${count:-0}"
}

# Sanitize a label value for Prometheus exposition format
# Args: $1 - raw value
# Returns: Escaped string safe for label values
sanitize_label() {
    local val="$1"
    val="${val//\\/\\\\}"
    val="${val//\"/\\\"}"
    val="${val//$'\n'/}"
    echo "$val"
}

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

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

    if ! check_samba; then
        cat <<EOF
# HELP samba_up Samba exporter status
# TYPE samba_up gauge
samba_up 0
EOF
        return
    fi

    cat <<EOF
# HELP samba_up Samba exporter status
# TYPE samba_up gauge
samba_up 1

EOF

    # Server version info
    local version
    version=$(get_samba_version)

    cat <<EOF
# HELP samba_server_info Samba server version information
# TYPE samba_server_info gauge
samba_server_info{version="$(sanitize_label "$version")"} 1

EOF

    # Active sessions
    local session_data session_count
    session_data=$(get_session_data)
    if [ -n "$session_data" ]; then
        session_count=$(echo "$session_data" | wc -l)
    else
        session_count=0
    fi

    cat <<EOF
# HELP samba_sessions_active Number of active SMB sessions
# TYPE samba_sessions_active gauge
samba_sessions_active $session_count

EOF

    # Per-session info
    cat <<EOF
# HELP samba_session_info Per-session information
# TYPE samba_session_info gauge
EOF

    if [ -n "$session_data" ]; then
        while IFS='|' read -r _ user _ machine proto; do
            [ -z "$user" ] && continue
            user=$(sanitize_label "$user")
            machine=$(sanitize_label "$machine")
            proto=$(sanitize_label "$proto")
            echo "samba_session_info{user=\"$user\",machine=\"$machine\",protocol_version=\"$proto\"} 1"
        done <<< "$session_data"
    fi

    echo ""

    # Active share connections
    local share_data share_count
    share_data=$(get_share_data)
    if [ -n "$share_data" ]; then
        share_count=$(echo "$share_data" | wc -l)
    else
        share_count=0
    fi

    cat <<EOF
# HELP samba_shares_active Number of active share connections
# TYPE samba_shares_active gauge
samba_shares_active $share_count

EOF

    # Per-share connection info
    # We need user info for share connections; re-parse with user data
    cat <<EOF
# HELP samba_share_connections Per-share connection information
# TYPE samba_share_connections gauge
EOF

    if [ -n "$share_data" ]; then
        # Build a PID-to-user map from session data
        local -A pid_user_map
        if [ -n "$session_data" ]; then
            while IFS='|' read -r pid user _ _ _; do
                [ -z "$pid" ] && continue
                pid_user_map["$pid"]="$user"
            done <<< "$session_data"
        fi

        while IFS='|' read -r share pid machine; do
            [ -z "$share" ] && continue
            local user="${pid_user_map[$pid]:-unknown}"
            share=$(sanitize_label "$share")
            machine=$(sanitize_label "$machine")
            user=$(sanitize_label "$user")
            echo "samba_share_connections{share=\"$share\",machine=\"$machine\",user=\"$user\"} 1"
        done <<< "$share_data"
    fi

    echo ""

    # Configured shares
    local configured_shares
    configured_shares=$(get_configured_shares)

    cat <<EOF
# HELP samba_share_configured_total Total number of configured shares
# TYPE samba_share_configured_total gauge
samba_share_configured_total $configured_shares

EOF

    # Locked files
    local locked_files
    locked_files=$(get_locked_files_count)

    cat <<EOF
# HELP samba_locked_files_total Total number of locked files
# TYPE samba_locked_files_total gauge
samba_locked_files_total $locked_files

EOF

    # Process status
    local smbd_running nmbd_running winbindd_running smbd_pids
    smbd_running=$(check_process_running "smbd")
    nmbd_running=$(check_process_running "nmbd")
    winbindd_running=$(check_process_running "winbindd")
    smbd_pids=$(get_smbd_pid_count)

    cat <<EOF
# HELP samba_smbd_running Whether smbd process is running
# TYPE samba_smbd_running gauge
samba_smbd_running $smbd_running

# HELP samba_nmbd_running Whether nmbd process is running
# TYPE samba_nmbd_running gauge
samba_nmbd_running $nmbd_running

# HELP samba_winbindd_running Whether winbindd process is running
# TYPE samba_winbindd_running gauge
samba_winbindd_running $winbindd_running

# HELP samba_smbd_pid_count Number of smbd processes
# TYPE samba_smbd_pid_count gauge
samba_smbd_pid_count $smbd_pids

EOF

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

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

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

    echo ""
}

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

run_http_server() {
    log "Starting Samba exporter on port $HTTP_PORT..."

    if ! command -v nc >/dev/null 2>&1; then
        error "netcat (nc) required for HTTP mode"
        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>Samba Exporter</title></head><body><h1>Samba 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 EXECUTION
# ============================================================================

main() {
    parse_args "$@"

    if ! acquire_lock; then
        exit 1
    fi

    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}/.samba_metrics.XXXXXX")

        if ! generate_metrics > "$temp_file" 2>/dev/null; then
            rm -f "$temp_file"
            error "Failed to generate metrics"
            exit 1
        fi

        local file_lines
        file_lines=$(wc -l < "$temp_file" 2>/dev/null || echo 0)

        if [ "$file_lines" -lt 5 ]; then
            rm -f "$temp_file"
            error "Metrics file too small ($file_lines lines), keeping previous"
            exit 1
        fi

        chmod 644 "$temp_file"
        mv -f "$temp_file" "$OUTPUT_FILE"

        log "Metrics written to $OUTPUT_FILE ($file_lines lines)"
    else
        generate_metrics
    fi
}

main "$@"