#!/bin/bash
################################################################################
# Script Name: dovecot-metrics-exporter.sh
# Description: Prometheus exporter for Dovecot IMAP/POP3 server metrics
#
# Collects connection counts, authentication stats, mailbox operations,
# process info, and protocol-level metrics from doveadm and Dovecot stats.
#
# Author:  Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
# Version: 1.0
#
# Usage:
#   # Output to stdout
#   ./dovecot-metrics-exporter.sh
#
#   # Textfile collector mode (atomic write)
#   ./dovecot-metrics-exporter.sh --textfile
#
#   # Custom output file
#   ./dovecot-metrics-exporter.sh -o /path/to/metrics.prom
#
################################################################################

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

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HOSTNAME=$(hostname)

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

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

Export Dovecot IMAP/POP3 server metrics as Prometheus metrics.

MODES:
    --textfile      Write to node_exporter textfile collector
                    (writes to $TEXTFILE_DIR/dovecot-metrics.prom)

OPTIONS:
    -o, --output    Output file path (for custom locations)
    -h, --help      Show this help message

EXAMPLES:
    $0                               # Output to stdout
    $0 --textfile                    # Write to textfile collector
    $0 -o /tmp/dovecot.prom          # Write to custom file

EOF
    exit 0
}

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

# Safe integer extraction — returns 0 on failure
safe_int() {
    local val="$1"
    if [[ "$val" =~ ^[0-9]+$ ]]; then
        echo "$val"
    else
        echo 0
    fi
}

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

generate_metrics() {
    local START_TIME
    START_TIME=$(date +%s.%N)

# --- Exporter info ---
echo "# HELP dovecot_up Exporter status (1=up, 0=down)"
echo "# TYPE dovecot_up gauge"

# Check if Dovecot is running
if systemctl is-active --quiet dovecot 2>/dev/null; then
    echo "dovecot_up 1"
else
    echo "dovecot_up 0"
fi

echo ""
echo "# HELP dovecot_exporter_info Exporter version information"
echo "# TYPE dovecot_exporter_info gauge"
echo 'dovecot_exporter_info{version="1.0"} 1'
echo ""

# --- Dovecot version ---
echo "# HELP dovecot_version_info Dovecot version information"
echo "# TYPE dovecot_version_info gauge"
local dovecot_version
dovecot_version=$(dovecot --version 2>/dev/null | awk '{print $1}') || dovecot_version="unknown"
echo "dovecot_version_info{version=\"${dovecot_version}\"} 1"
echo ""

# --- Process counts ---
echo "# HELP dovecot_processes Number of running Dovecot processes by type"
echo "# TYPE dovecot_processes gauge"
for proc_type in imap pop3 lmtp managesieve submission auth anvil; do
    count=$(pgrep -c "dovecot/${proc_type}" 2>/dev/null) || count=0
    echo "dovecot_processes{type=\"${proc_type}\"} ${count}"
done
local total_procs
total_procs=$(pgrep -c dovecot 2>/dev/null) || total_procs=0
echo "dovecot_processes{type=\"total\"} ${total_procs}"
echo ""

# --- Connected users (from doveadm) ---
echo "# HELP dovecot_connected_users Number of currently connected users by protocol"
echo "# TYPE dovecot_connected_users gauge"
local imap_users=0 pop3_users=0 lmtp_users=0 managesieve_users=0
if command -v doveadm >/dev/null 2>&1; then
    imap_users=$(doveadm who -1 2>/dev/null | grep -c 'imap' 2>/dev/null) || imap_users=0
    pop3_users=$(doveadm who -1 2>/dev/null | grep -c 'pop3' 2>/dev/null) || pop3_users=0
    lmtp_users=$(doveadm who -1 2>/dev/null | grep -c 'lmtp' 2>/dev/null) || lmtp_users=0
    managesieve_users=$(doveadm who -1 2>/dev/null | grep -c 'managesieve' 2>/dev/null) || managesieve_users=0
fi
echo "dovecot_connected_users{protocol=\"imap\"} ${imap_users}"
echo "dovecot_connected_users{protocol=\"pop3\"} ${pop3_users}"
echo "dovecot_connected_users{protocol=\"lmtp\"} ${lmtp_users}"
echo "dovecot_connected_users{protocol=\"managesieve\"} ${managesieve_users}"
echo ""

# --- Total connections (from doveadm who) ---
echo "# HELP dovecot_connections_total Total active connections by protocol"
echo "# TYPE dovecot_connections_total gauge"
local imap_conns=0 pop3_conns=0
if command -v doveadm >/dev/null 2>&1; then
    imap_conns=$(doveadm who -1 2>/dev/null | grep 'imap' | awk '{sum+=$3} END {print sum+0}' 2>/dev/null) || imap_conns=0
    pop3_conns=$(doveadm who -1 2>/dev/null | grep 'pop3' | awk '{sum+=$3} END {print sum+0}' 2>/dev/null) || pop3_conns=0
fi
echo "dovecot_connections_total{protocol=\"imap\"} ${imap_conns}"
echo "dovecot_connections_total{protocol=\"pop3\"} ${pop3_conns}"
echo ""

# --- Authentication stats from mail.log ---
local LOG_FILE="/var/log/mail.log"
if [[ ! -f "$LOG_FILE" ]]; then
    LOG_FILE="/var/log/maillog"
fi

echo "# HELP dovecot_auth_success_total Successful authentication attempts by protocol"
echo "# TYPE dovecot_auth_success_total counter"
local imap_auth_ok=0 pop3_auth_ok=0
if [[ -f "$LOG_FILE" ]]; then
    imap_auth_ok=$(grep -c 'imap-login: Info: Login:' "$LOG_FILE" 2>/dev/null) || imap_auth_ok=0
    pop3_auth_ok=$(grep -c 'pop3-login: Info: Login:' "$LOG_FILE" 2>/dev/null) || pop3_auth_ok=0
fi
echo "dovecot_auth_success_total{protocol=\"imap\"} ${imap_auth_ok}"
echo "dovecot_auth_success_total{protocol=\"pop3\"} ${pop3_auth_ok}"
echo ""

echo "# HELP dovecot_auth_failed_total Failed authentication attempts by protocol"
echo "# TYPE dovecot_auth_failed_total counter"
local imap_auth_fail=0 pop3_auth_fail=0
if [[ -f "$LOG_FILE" ]]; then
    imap_auth_fail=$(grep -c 'imap-login:.*auth failed\|imap-login: Info: Aborted login' "$LOG_FILE" 2>/dev/null) || imap_auth_fail=0
    pop3_auth_fail=$(grep -c 'pop3-login:.*auth failed\|pop3-login: Info: Aborted login' "$LOG_FILE" 2>/dev/null) || pop3_auth_fail=0
fi
echo "dovecot_auth_failed_total{protocol=\"imap\"} ${imap_auth_fail}"
echo "dovecot_auth_failed_total{protocol=\"pop3\"} ${pop3_auth_fail}"
echo ""

# --- TLS connections ---
echo "# HELP dovecot_tls_connections_total TLS connections by status"
echo "# TYPE dovecot_tls_connections_total counter"
local tls_yes=0 tls_no=0
if [[ -f "$LOG_FILE" ]]; then
    tls_yes=$(grep -c 'Login:.*TLS' "$LOG_FILE" 2>/dev/null) || tls_yes=0
    tls_no=$(grep 'Login:' "$LOG_FILE" 2>/dev/null | grep -cv 'TLS' 2>/dev/null) || tls_no=0
fi
echo "dovecot_tls_connections_total{tls=\"yes\"} ${tls_yes}"
echo "dovecot_tls_connections_total{tls=\"no\"} ${tls_no}"
echo ""

# --- Authentication methods ---
echo "# HELP dovecot_auth_method_total Logins by authentication method"
echo "# TYPE dovecot_auth_method_total counter"
if [[ -f "$LOG_FILE" ]]; then
    for method in PLAIN LOGIN CRAM-MD5 DIGEST-MD5; do
        count=$(grep -c "Login:.*method=${method}" "$LOG_FILE" 2>/dev/null) || count=0
        echo "dovecot_auth_method_total{method=\"${method}\"} ${count}"
    done
else
    for method in PLAIN LOGIN CRAM-MD5 DIGEST-MD5; do
        echo "dovecot_auth_method_total{method=\"${method}\"} 0"
    done
fi
echo ""

# --- Disconnections ---
echo "# HELP dovecot_disconnections_total Client disconnections by reason"
echo "# TYPE dovecot_disconnections_total counter"
local dc_logout=0 dc_timeout=0 dc_closed=0 dc_internal=0
if [[ -f "$LOG_FILE" ]]; then
    dc_logout=$(grep -c 'Logged out' "$LOG_FILE" 2>/dev/null) || dc_logout=0
    dc_timeout=$(grep -c 'Disconnected.*Timed out\|Connection timed out' "$LOG_FILE" 2>/dev/null) || dc_timeout=0
    dc_closed=$(grep -c 'Disconnected.*Connection closed' "$LOG_FILE" 2>/dev/null) || dc_closed=0
    dc_internal=$(grep -c 'Disconnected.*Internal error' "$LOG_FILE" 2>/dev/null) || dc_internal=0
fi
echo "dovecot_disconnections_total{reason=\"logout\"} ${dc_logout}"
echo "dovecot_disconnections_total{reason=\"timeout\"} ${dc_timeout}"
echo "dovecot_disconnections_total{reason=\"connection_closed\"} ${dc_closed}"
echo "dovecot_disconnections_total{reason=\"internal_error\"} ${dc_internal}"
echo ""

# --- LMTP delivery stats ---
echo "# HELP dovecot_lmtp_deliveries_total LMTP deliveries by status"
echo "# TYPE dovecot_lmtp_deliveries_total counter"
local lmtp_ok=0 lmtp_reject=0 lmtp_tempfail=0
if [[ -f "$LOG_FILE" ]]; then
    lmtp_ok=$(grep -c 'lmtp.*saved mail' "$LOG_FILE" 2>/dev/null) || lmtp_ok=0
    lmtp_reject=$(grep -c 'lmtp.*rejected' "$LOG_FILE" 2>/dev/null) || lmtp_reject=0
    lmtp_tempfail=$(grep -c 'lmtp.*temporary failure\|lmtp.*temp-fail' "$LOG_FILE" 2>/dev/null) || lmtp_tempfail=0
fi
echo "dovecot_lmtp_deliveries_total{status=\"delivered\"} ${lmtp_ok}"
echo "dovecot_lmtp_deliveries_total{status=\"rejected\"} ${lmtp_reject}"
echo "dovecot_lmtp_deliveries_total{status=\"tempfail\"} ${lmtp_tempfail}"
echo ""

# --- Sieve stats ---
echo "# HELP dovecot_sieve_actions_total Sieve filter actions"
echo "# TYPE dovecot_sieve_actions_total counter"
local sieve_filed=0 sieve_discard=0 sieve_redirect=0 sieve_reject=0
if [[ -f "$LOG_FILE" ]]; then
    sieve_filed=$(grep -c 'sieve:.*stored mail\|sieve:.*fileinto' "$LOG_FILE" 2>/dev/null) || sieve_filed=0
    sieve_discard=$(grep -c 'sieve:.*discard' "$LOG_FILE" 2>/dev/null) || sieve_discard=0
    sieve_redirect=$(grep -c 'sieve:.*redirect' "$LOG_FILE" 2>/dev/null) || sieve_redirect=0
    sieve_reject=$(grep -c 'sieve:.*reject' "$LOG_FILE" 2>/dev/null) || sieve_reject=0
fi
echo "dovecot_sieve_actions_total{action=\"filed\"} ${sieve_filed}"
echo "dovecot_sieve_actions_total{action=\"discard\"} ${sieve_discard}"
echo "dovecot_sieve_actions_total{action=\"redirect\"} ${sieve_redirect}"
echo "dovecot_sieve_actions_total{action=\"reject\"} ${sieve_reject}"
echo ""

# --- Dovecot stats (if old_stats or stats plugin enabled) ---
# Try doveadm stats dump for Dovecot 2.3+
echo "# HELP dovecot_mail_commands_total Mail commands executed"
echo "# TYPE dovecot_mail_commands_total counter"
local cmds_select=0 cmds_fetch=0 cmds_store=0 cmds_search=0 cmds_copy=0 cmds_expunge=0
if command -v doveadm >/dev/null 2>&1; then
    local stats_output
    stats_output=$(doveadm stats dump session 2>/dev/null | head -20)
    if [[ -n "$stats_output" ]]; then
        cmds_select=$(echo "$stats_output" | awk '{sum+=$4} END {print sum+0}') || cmds_select=0
        cmds_fetch=$(echo "$stats_output" | awk '{sum+=$5} END {print sum+0}') || cmds_fetch=0
    fi
fi
# Fallback: count from logs
if [[ -f "$LOG_FILE" ]]; then
    cmds_copy=$(grep -c 'Copy\|copy' "$LOG_FILE" 2>/dev/null | head -1) || cmds_copy=0
    cmds_expunge=$(grep -c 'Expunged' "$LOG_FILE" 2>/dev/null) || cmds_expunge=0
fi
echo "dovecot_mail_commands_total{command=\"copy\"} ${cmds_copy}"
echo "dovecot_mail_commands_total{command=\"expunge\"} ${cmds_expunge}"
echo ""

# --- Mail storage quota (top users if doveadm quota available) ---
echo "# HELP dovecot_quota_usage_bytes User quota usage in bytes (top users)"
echo "# TYPE dovecot_quota_usage_bytes gauge"
echo "# HELP dovecot_quota_limit_bytes User quota limit in bytes"
echo "# TYPE dovecot_quota_limit_bytes gauge"
if command -v doveadm >/dev/null 2>&1; then
    doveadm quota get -A 2>/dev/null | grep 'STORAGE' | head -20 | while IFS=$'\t' read -r user type value limit _; do
        local usage_bytes=$((value * 1024))
        local limit_bytes=$((limit * 1024))
        echo "dovecot_quota_usage_bytes{user=\"${user}\"} ${usage_bytes}"
        echo "dovecot_quota_limit_bytes{user=\"${user}\"} ${limit_bytes}"
    done 2>/dev/null
fi
echo ""

# --- Dovecot uptime ---
echo "# HELP dovecot_uptime_seconds Dovecot process uptime in seconds"
echo "# TYPE dovecot_uptime_seconds gauge"
local dovecot_pid uptime_seconds=0
dovecot_pid=$(pgrep -o dovecot 2>/dev/null) || dovecot_pid=""
if [[ -n "$dovecot_pid" ]] && [[ -d "/proc/${dovecot_pid}" ]]; then
    local start_time
    start_time=$(stat -c %Y "/proc/${dovecot_pid}" 2>/dev/null) || start_time=0
    if [[ "$start_time" -gt 0 ]]; then
        uptime_seconds=$(( $(date +%s) - start_time ))
    fi
fi
echo "dovecot_uptime_seconds ${uptime_seconds}"
echo ""

# --- Memory usage ---
echo "# HELP dovecot_memory_bytes Total memory usage of all Dovecot processes"
echo "# TYPE dovecot_memory_bytes gauge"
local total_mem=0
total_mem=$(pgrep dovecot 2>/dev/null | xargs -I {} cat /proc/{}/status 2>/dev/null | awk '/VmRSS/{sum+=$2} END {print sum*1024+0}') || total_mem=0
echo "dovecot_memory_bytes ${total_mem}"
echo ""

# --- Script execution time ---
local END_TIME
END_TIME=$(date +%s.%N)
local DURATION
DURATION=$(echo "$END_TIME - $START_TIME" | bc)

echo "# HELP dovecot_exporter_duration_seconds Time to generate all metrics"
echo "# TYPE dovecot_exporter_duration_seconds gauge"
echo "dovecot_exporter_duration_seconds ${DURATION}"
echo ""

echo "# HELP dovecot_exporter_last_run_timestamp Unix timestamp of last successful run"
echo "# TYPE dovecot_exporter_last_run_timestamp gauge"
echo "dovecot_exporter_last_run_timestamp $(date +%s)"
}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

main() {
    parse_args "$@"

    if [ -n "$OUTPUT_FILE" ]; then
        local output_dir
        output_dir="$(dirname "$OUTPUT_FILE")"
        mkdir -p "$output_dir"

        local temp_file
        temp_file=$(mktemp "${output_dir}/.dovecot_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

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

        if [ "$file_lines" -lt 10 ]; then
            rm -f "$temp_file"
            echo "ERROR: Metrics file too small ($file_lines lines), keeping previous" >&2
            exit 1
        fi

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

        echo "Metrics written to $OUTPUT_FILE ($file_lines lines)" >&2
    else
        generate_metrics
    fi
}

main "$@"