#!/bin/bash

#############################################################
#### APT Package Updates Exporter for Prometheus         ####
#### Expose pending apt updates as Prometheus metrics    ####
#### for Debian and Ubuntu servers                       ####
####                                                     ####
#### Author: Phil Connor                                 ####
#### Contact: contact@mylinux.work                       ####
#### License: MIT                                        ####
#### Version: 1.7                                        ####
####                                                     ####
#### Usage: ./apt-updates-exporter.sh                    ####
#############################################################

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

Export pending APT package updates as Prometheus metrics (v1.7).

MODES:
    --textfile      Write to node_exporter textfile collector
    (default)       Output metrics to stdout

OPTIONS:
    -o, --output    Output file path
    -h, --help      Show this help message

ENVIRONMENT VARIABLES:
    METRICS_DIR              Textfile collector directory (default: /var/lib/node_exporter)
    AUTO_UPDATE_ENABLED      Enable automatic package updates (default: false)
    AUTO_REMOVE_ENABLED      Enable automatic removal of orphaned packages (default: false)
    WAIT_PERIOD_ENABLED      Wait 3 days before auto-updating (default: true)
    CRON_INTERVAL            Cron schedule expression (default: 0 0 * * *)

METRICS:
    - node_upgrades_pending       Pending updates by origin and architecture
    - node_upgradelist             Per-package upgrade details (name, versions, origin)
    - node_autoremove_packages     Packages available for auto-removal
    - node_auto_updates            Packages auto-updated (when enabled)
    - node_auto_remove             Packages auto-removed (when enabled)
    - node_updated_packages        Packages updated in last run
    - node_removed_packages        Packages removed in last auto-removal
    - node_reboot_required         Whether a reboot is required

EXAMPLES:
    $0                                          # Output metrics to stdout
    $0 --textfile                               # Write to textfile collector
    $0 -o /tmp/apt_updates.prom                 # Write to custom file
    AUTO_UPDATE_ENABLED=true $0                 # Enable auto-updates

EOF
    exit 0
}

OUTPUT_FILE=""

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

parse_args "$@"

# Configuration variables with default values
AUTO_UPDATE_ENABLED="${AUTO_UPDATE_ENABLED:-false}"         # Enable automatic package updates
AUTO_REMOVE_ENABLED="${AUTO_REMOVE_ENABLED:-false}"         # Enable automatic removal of orphaned packages
APT_GET_CMD="${APT_GET_CMD:-/usr/bin/apt-get}"              # Path to apt-get command
AWK_CMD="${AWK_CMD:-/usr/bin/awk}"                          # Path to awk command
CRON_INTERVAL="${CRON_INTERVAL:-0 0 * * *}"                 # Cron schedule (daily at midnight)
GREP_CMD="${GREP_CMD:-/usr/bin/grep}"                       # Path to grep command
METRICS_DIR="${METRICS_DIR:-/var/lib/node_exporter}"        # Directory for Prometheus metrics files
SORT_CMD="${SORT_CMD:-/usr/bin/sort}"                       # Path to sort command
UNIQ_CMD="${UNIQ_CMD:-/usr/bin/uniq}"                       # Path to uniq command

# File paths for tracking update state
UPDATES_TIMESTAMP_FILE="$METRICS_DIR/updates_detected"      # Tracks when updates were first detected
WAIT_PERIOD_ENABLED="${WAIT_PERIOD_ENABLED:-true}"          # Enable waiting period before auto-updates
UPDATED_PACKAGES_FILE="$METRICS_DIR/updated_packages"       # List of packages updated in last run
AUTO_REMOVE_FILE="$METRICS_DIR/auto_remove_packages"        # List of packages removed in last auto-removal
WAIT_PERIOD_SECONDS=$((3 * 24 * 60 * 60))                   # Wait period: 3 days in seconds

# Safety check: prevent concurrent execution that could cause lock conflicts
if pidof apt apt-get >/dev/null || fuser /var/lib/dpkg/lock /var/lib/apt/lists/lock /var/cache/apt/archives/lock >/dev/null 2>&1; then
    echo "node_upgrades_pending{origin=\"error\",arch=\"unknown\"} -1"
    echo "node_upgradelist{pkgname=\"error\", update_version=\"\", current_version=\"\", origin=\"\"} -1"
    echo "node_auto_updates{status=\"error\"} -1"
    exit 1
fi

#### Setup: Ensure metrics directory exists with proper permissions ####
if [[ ! -d "$METRICS_DIR" ]]; then
    # Create metrics directory
    mkdir -p "$METRICS_DIR" || {
        echo "Failed to create $METRICS_DIR"
        exit 1
    }
    # Set ownership to prometheus user (try different formats for compatibility)
    chown prometheus:prometheus "$METRICS_DIR" 2>/dev/null || chown prometheus. "$METRICS_DIR" || {
        echo "Failed to set ownership of $METRICS_DIR"
        exit 1
    }
    # Set appropriate permissions for metrics directory
    chmod 755 "$METRICS_DIR" || {
        echo "Failed to set permissions on $METRICS_DIR"
        exit 1
    }
fi

#### Setup: Ensure cron job exists for automated execution ####
if ! crontab -l | grep -q "updates.sh"; then
    # Add cron job to run this script automatically
    echo -e "$(crontab -u root -l)\n$CRON_INTERVAL /usr/local/bin/updates.sh > $METRICS_DIR/updates.prom 2>&1" | crontab -u root -
    # Verify the cron job was added successfully
    crontab -l | grep -q "updates.sh" || {
        echo "Failed to add cron job"
        exit 1
    }
fi

#### Setup: Ensure logrotate configuration exists for state files ####
LOGROTATE_CONFIG="/etc/logrotate.d/node-exporter-metrics"
if [[ ! -f "$LOGROTATE_CONFIG" ]]; then
    # Create logrotate configuration for monthly rotation of state files
    cat >"$LOGROTATE_CONFIG" <<EOF
# Logrotate configuration for node_exporter metrics files
# Created to handle updates.sh state files

$UPDATES_TIMESTAMP_FILE
$UPDATED_PACKAGES_FILE
$AUTO_REMOVE_FILE {
    monthly
    rotate 12
    compress
    delaycompress
    missingok
    notifempty
    create 0644 root root
    postrotate
        # Signal node_exporter to reload if running
        systemctl reload-or-restart node_exporter 2>/dev/null || true
    endscript
}
EOF
    # Verify the logrotate configuration is valid
    logrotate -d "$LOGROTATE_CONFIG" >/dev/null 2>&1 || {
        echo "Failed to create valid logrotate configuration"
        rm -f "$LOGROTATE_CONFIG"
        exit 1
    }
fi

#### Function: Count pending package upgrades grouped by origin and architecture ####
get_upgrades() {
    # Test apt-get upgrade command and exit on failure
    if ! $APT_GET_CMD -qq --just-print upgrade >/dev/null 2>&1; then
        echo "node_upgrades_pending{origin=\"error\",arch=\"unknown\"} -1"
        return 1
    fi

    # Parse apt-get output to extract package info to create Prometheus metrics
    $APT_GET_CMD -qq --just-print upgrade |
        $AWK_CMD -F '[()]' '/^Inst/ {
        sub("^[^ ]+ ", "", $2)           # Remove package name from origin field
        gsub(" ","",$2)                  # Remove spaces from origin
        sub(/\[|\]/, " ", $2)            # Replace brackets with space
        print $2
    }' |
        $SORT_CMD |                      # Sort the output
        $UNIQ_CMD -c |                   # Count unique entries
        $AWK_CMD '{
        gsub(/\\\\/, "\\\\", $2)         # Escape backslashes for Prometheus labels
        gsub(/\\/, "\\\\", $2)
        gsub(/"/, "\\\"", $2)            # Escape quotes for Prometheus labels
        gsub(/\[|\]/, "", $3)            # Remove brackets from architecture
        printf "node_upgrades_pending{origin=\"%s\",arch=\"%s\"} %d\n", $2, $3, $1
    }'
}

#### Function: Handle automatic package updates with optional wait period ####
handle_auto_updates() {
    # Skip if auto-updates are disabled
    [[ "$AUTO_UPDATE_ENABLED" != "true" ]] && return

    local should_update=false

    # Check if we should wait before updating (prevents immediate updates on detection)
    if [[ "$WAIT_PERIOD_ENABLED" == "true" ]]; then
        local current_time detected_time
        current_time=$(date +%s)
        detected_time=$(cat "$UPDATES_TIMESTAMP_FILE" 2>/dev/null || echo "0")
        # Only update if wait period has elapsed
        ((current_time - detected_time >= WAIT_PERIOD_SECONDS)) && should_update=true
    else
        # Update immediately if wait period is disabled
        should_update=true
    fi

    if [[ "$should_update" == "true" ]]; then
        perform_auto_update
        # Clear timestamp file after updating (reset wait period)
        [[ "$WAIT_PERIOD_ENABLED" == "true" ]] && rm -f "$UPDATES_TIMESTAMP_FILE"
    fi
}

#### Function: Execute automatic package updates and record metrics ####
perform_auto_update() {
    # Output Prometheus metric headers
    echo '# HELP node_auto_updates Number of packages auto-updated.'
    echo '# TYPE node_auto_updates gauge'

    local update_output update_count
    # Run apt update and upgrade non-interactively with timeout, capture output
    update_output=$(timeout 300 bash -c "DEBIAN_FRONTEND=noninteractive $APT_GET_CMD update >/dev/null 2>&1 && DEBIAN_FRONTEND=noninteractive $APT_GET_CMD -y -o Dpkg::Options::='--force-confdef' -o Dpkg::Options::='--force-confold' upgrade 2>&1" || echo "TIMEOUT_ERROR")
    # Check for timeout or other errors
    if [[ "$update_output" == *"TIMEOUT_ERROR"* ]]; then
        echo "node_auto_updates{status=\"timeout\"} -1"
        return 1
    fi
    
    # Count number of packages that were actually unpacked/updated
    update_count=$(echo "$update_output" | grep -c "^Unpacking ")

    echo "node_auto_updates{status=\"success\"} $update_count"
    # Save list of updated packages with versions for reporting
    echo "$update_output" | grep "^Unpacking " | awk '{gsub(/[()]/, "", $3); print $2 " " $3}' >"$UPDATED_PACKAGES_FILE"
}

#### Function: Generate detailed list of individual packages available for upgrade ####
get_upgrade_list() {
    # Test apt-get upgrade command and handle failures
    if ! $APT_GET_CMD --just-print upgrade >/dev/null 2>&1; then
        echo 'node_upgradelist{pkgname="error", update_version="", current_version="", origin=""} -1'
        return 1
    fi

    # Parse each package installation line to extract detailed package information
    $APT_GET_CMD --just-print upgrade |
        $GREP_CMD Inst |                # Filter for installation lines
        $AWK_CMD '{
        gsub(/\(|\)/, "", $4)            # Remove parentheses from version
        gsub(/:/, ".", $4)               # Replace colons with dots in version
        gsub(/\[|\]/, "", $3)            # Remove brackets from current version
        gsub(/:/, " ", $5)               # Replace colons with spaces in origin
        
        # Escape special characters for Prometheus label values
        gsub(/\\/, "\\\\", $2)           # Escape backslashes in package name
        gsub(/"/, "\\\"", $2)            # Escape quotes in package name
        gsub(/\\/, "\\\\", $4)           # Escape backslashes in version
        gsub(/"/, "\\\"", $4)            # Escape quotes in version
        
        # Output Prometheus metric with package details
        printf "node_upgradelist{pkgname=\"%s\",update_version=\"%s\", current_version=\"%s\", origin=\"%s\"} 1\n", $2, $4, $3, $5
    }'
}

#### Function: Get list of packages that can be automatically removed (orphaned) ####
get_auto_remove_list() {
    # Test autoremove command with dry-run to see what would be removed
    if ! $APT_GET_CMD --dry-run autoremove >/dev/null 2>&1; then
        echo 'node_autoremove_packages{pkgname="error"} -1'
        return 1
    fi

    # Parse dry-run output to find packages that would be removed
    $APT_GET_CMD --dry-run autoremove 2>/dev/null |
        $GREP_CMD "Remv" |              # Filter for removal lines
        $AWK_CMD '{
        # Escape special characters for Prometheus labels
        gsub(/\\/, "\\\\", $2)           # Escape backslashes in package name
        gsub(/"/, "\\\"", $2)            # Escape quotes in package name
        printf "node_autoremove_packages{pkgname=\"%s\"} 1\n", $2
    }'
}

#### Function: Handle automatic removal of orphaned packages ####
handle_auto_remove() {
    # Skip if auto-remove is disabled
    [[ "$AUTO_REMOVE_ENABLED" != "true" ]] && return

    perform_auto_remove
}

#### Function: Execute automatic package removal and record metrics ####
perform_auto_remove() {
    # Output Prometheus metric headers
    echo '# HELP node_auto_remove Number of packages auto-removed.'
    echo '# TYPE node_auto_remove gauge'

    local remove_output remove_count
    # Run autoremove non-interactively and capture output
    remove_output=$(DEBIAN_FRONTEND=noninteractive $APT_GET_CMD -y autoremove 2>&1)
    # Count packages that were actually removed
    remove_count=$(echo "$remove_output" | grep -c "^Removing ")

    echo "node_auto_remove{status=\"success\"} $remove_count"
    # Save list of removed packages for reporting
    echo "$remove_output" | grep "^Removing " | awk '{print $2}' >"$AUTO_REMOVE_FILE"
}

#### Generate all Prometheus metrics ####
generate_metrics() {
    #### Upgrade list metrics ####
    upgradelist=$(get_upgrade_list)
    echo '# HELP node_upgradelist List of packages for upgrade'
    echo '# TYPE node_upgradelist gauge'
    if [[ -n "${upgradelist}" ]]; then
        echo "${upgradelist}"
    else
        echo 'node_upgradelist{pkgname="", update_version="", current_version="", origin=""} 0'
    fi

    #### Pending upgrades metrics and auto-updates ####
    pending_upgrades=$(get_upgrades)
    echo '# HELP node_upgrades_pending Apt package pending updates by origin.'
    echo '# TYPE node_upgrades_pending gauge'

    if [[ -n "$pending_upgrades" ]]; then
        printf "%s\n" "$pending_upgrades"

        if [[ ! -f "$UPDATES_TIMESTAMP_FILE" ]]; then
            date +%s >"$UPDATES_TIMESTAMP_FILE"
        fi

        handle_auto_updates
    else
        echo 'node_upgrades_pending{origin="", arch=""} 0'
        echo '# HELP node_auto_updates Number of packages auto-updated.'
        echo '# TYPE node_auto_updates gauge'
        echo 'node_auto_updates{status="success"} 0'
        rm -f "$UPDATES_TIMESTAMP_FILE"
    fi

    #### Auto-removable packages metrics ####
    autoremovelist=$(get_auto_remove_list)
    echo '# HELP node_autoremove_packages List of packages available for auto-removal'
    echo '# TYPE node_autoremove_packages gauge'
    if [[ -n "${autoremovelist}" ]]; then
        echo "${autoremovelist}"
        handle_auto_remove
    else
        echo 'node_autoremove_packages{pkgname=""} 0'
    fi

    #### Packages updated in the last run ####
    if [[ -f "$UPDATED_PACKAGES_FILE" ]]; then
        echo '# HELP node_updated_packages List of packages updated in last update'
        echo '# TYPE node_updated_packages gauge'
        while IFS=' ' read -r package version; do
            echo "node_updated_packages{package=\"$package\",version=\"$version\"} 1"
        done <"$UPDATED_PACKAGES_FILE"
    fi

    #### Packages removed in the last auto-removal ####
    if [[ -f "$AUTO_REMOVE_FILE" ]]; then
        echo '# HELP node_removed_packages List of packages removed in last auto-removal'
        echo '# TYPE node_removed_packages gauge'
        while IFS= read -r package; do
            echo "node_removed_packages{package=\"$package\"} 1"
        done <"$AUTO_REMOVE_FILE"
    fi

    #### Reboot required check ####
    echo '# HELP node_reboot_required Node reboot is required for software updates.'
    echo '# TYPE node_reboot_required gauge'
    if [[ -f '/run/reboot-required' ]]; then
        echo 'node_reboot_required 1'
    else
        echo 'node_reboot_required 0'
    fi
}

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

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

    file_lines=$(wc -l < "$temp_file" 2>/dev/null || echo 0)
    if [[ "$file_lines" -lt 5 ]]; 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