<#
.SYNOPSIS
    Windows Update Monitoring Script for Prometheus Windows Exporter (wsus)

.DESCRIPTION
    This script monitors Windows Updates and exports metrics in Prometheus format.
    It checks for available updates, their severity levels, and installation status.
    Can optionally install updates when specified.

.PARAMETER MetricNames
    Custom metric names hashtable for overriding default metric names
.PARAMETER AutoInstall
    Enable automatic installation of available updates
.PARAMETER ScheduleDaily
    Create a scheduled task to run this script daily at 3 AM
.PARAMETER MetricsFilePath
    Path where Prometheus metrics will be written (must be accessible by windows_exporter)

.NOTES
    Version: 1.6.2
    Author: Phil Connor contact@mylinux.work
    License: MIT
    
    Features:
    - Monitors Windows Updates status
    - Exports metrics in Prometheus format
    - Supports automatic update installation
    - Graceful shutdown support
    - Retry mechanisms and timeout handling
    - Memory management and garbage collection
    - Comprehensive logging
    - Error handling and metrics
#>

param(
    # Custom metric names hashtable for overriding default metric names
    [hashtable]$MetricNames,
    # Enable automatic installation of available updates
    [switch]$AutoInstall = $false,                                                                  
    # Create a scheduled task to run this script daily at 3 AM
    [switch]$ScheduleDaily,                                                                 
    # Validate that the parent directory exists for the metrics file
    [ValidateScript({Test-Path (Split-Path $_ -Parent) -PathType Container})]                       
    # Path where Prometheus metrics will be written (must be accessible by windows_exporter)
    [string]$MetricsFilePath = "C:\Program Files\windows_exporter\textfile_inputs\updates.prom"     
)

# Check if script is already running to prevent multiple instances
$scriptName = $MyInvocation.MyCommand.Name
$currentProcess = Get-Process -Id $PID
$runningInstances = Get-WmiObject Win32_Process | 
    Where-Object { $_.CommandLine -like "*$scriptName*" -and $_.ProcessId -ne $currentProcess.Id }

if ($runningInstances) {
    Write-Host "Script is already running (PID: $($runningInstances.ProcessId)). Exiting to prevent conflicts."
    exit 0
}

# Create scheduled task for daily execution at 3 AM
if ($ScheduleDaily -eq $true) {
    $taskName = "WindowsUpdateMonitoring"
    # Check if the scheduled task already exists to avoid duplicates
    $existingTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue
    
    if (-not $existingTask) {
        # Define the action: run PowerShell with this script
        $taskAction = New-ScheduledTaskAction -Execute "powershell.exe" -Argument "-NoProfile -ExecutionPolicy Bypass -File `"$($MyInvocation.MyCommand.Path)`""
        # Set the trigger to run daily at 3 AM
        $taskTrigger = New-ScheduledTaskTrigger -Daily -At 3AM
        # Run as SYSTEM account with highest privileges for update operations
        $taskPrincipal = New-ScheduledTaskPrincipal -UserId "SYSTEM" -LogonType ServiceAccount -RunLevel Highest

        try {
            Write-Host "Creating scheduled task: $taskName"
            # Register the scheduled task with Windows Task Scheduler
            Register-ScheduledTask -TaskName $taskName -Action $taskAction -Trigger $taskTrigger -Principal $taskPrincipal -Description "Monitors for Windows updates and optionally installs them automatically"

            # Verify the task was created successfully
            $createdTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue
            if (-not $createdTask) {
                throw "Failed to verify scheduled task creation"
            }
            Write-Host "Successfully created scheduled task: $taskName"
        } catch {
            Write-Error "Failed to create auto-start task: $($_.Exception.Message)"
            throw
        }
    } else {
        Write-Host "Scheduled task $taskName already exists. Skipping creation."
    }
}


# Define standard metric values used throughout the script
$script:MetricValues = @{
    Success = 1         # Indicates successful operation
    Error = 0           # Indicates error state
    NoError = 0         # Indicates no error occurred
    NoLastInstall = 0   # Indicates no previous installation timestamp
    NoInstallStatus = 0 # Indicates no installation status available
}

# Define Prometheus metric names following naming conventions
# These metrics will be exposed to Prometheus for monitoring
Set-Variable -Name MetricNames -Value @{
    # Script execution metrics - track the health and performance of this script
    Status = "windows_update_script_status"                         # Status of the script execution (1 for running, 0 for not running)
    Timestamp = "windows_update_script_timestamp_seconds"           # Timestamp of the last script execution
    Error = "windows_update_script_error"                           # Error message if script execution fails
    Runtime = "windows_update_script_runtime_seconds"               # Total runtime of the script
    
    # Update information metrics - track Windows Update status and availability
    Available = "windows_updates_available"                         # Total number of Windows updates available for installation
    Info = "windows_update_info"                                    # General update information
    Reboot = "windows_update_reboot_required"                       # Reboot requirement flag
    AutoUpdate = "windows_update_auto_install_enabled"              # Automatic update installation enabled flag (0 or 1)
    LastInstall = "windows_update_last_install_timestamp_seconds"   # Timestamp of the last successful update installation
    InstallStatus = "windows_update_install_status"                 # Status of the last update installation attempt (0 for success, 1 for failure)
    UpdateList = "windows_update_available_list"                    # List of available updates with details
    } -Option ReadOnly

# Validate all metric names conform to Prometheus naming standards
# Metric names must start with a letter or underscore, followed by letters, numbers, or underscores
$MetricNames.Values | ForEach-Object {
    if (-not ($_ -match '^[a-zA-Z_:][a-zA-Z0-9_:]*$')) {
        throw "Invalid Prometheus metric name: $_"
    }
}

# Function to write errors to log file and create error metrics for monitoring
function Write-ScriptError {
    param(
        [string]$Message,                   # Error message to log
        [string]$ErrorCode = "unknown",     # Error code for categorization
        [hashtable]$MetricNames,            # Metric names hashtable
        [string]$MetricsFilePath            # Path to write error metrics
    )
    
    # Log the error to PowerShell error stream
    Write-Error $Message
    
    # Create a Prometheus error metric with sanitized error code
    $errorMetric = "$($MetricNames.Error){error=`"$($ErrorCode -replace '"', '')`"} 1"
    try {
        # Write the error metric to the metrics file for Prometheus to collect
        $errorMetric | Out-File -FilePath $MetricsFilePath -Encoding UTF8 -Force
    }
    catch {
        Write-Warning "Failed to write error metric: $_"
    }
}

# Function to execute actions with standardized error handling
function Invoke-WithErrorHandling {
    param(
        [scriptblock]$Action,           # Script block to execute
        [string]$Operation,             # Description of the operation for error messages
        [hashtable]$MetricNames,        # Metric names hashtable
        [string]$MetricsFilePath        # Path to write error metrics
    )
    try {
        # Execute the provided script block
        & $Action
    } catch {
        # Handle any errors by writing to metrics and re-throwing
        Write-ScriptError -Message "Failed to $Operation`: $_" -ErrorCode $Operation.Replace(' ', '_') -MetricNames $MetricNames -MetricsFilePath $MetricsFilePath
        throw
    }
}

# Function to download and install Windows updates
function Install-WindowsUpdate {
    param(
        [Parameter(Mandatory)]
        [ValidateNotNull()]
        $Update,                        # Windows Update object to install
        
        [Parameter(Mandatory)]
        [ValidateNotNull()]
        $UpdateSession,                 # Windows Update session object
        
        [Parameter(Mandatory)]
        [ValidateNotNull()]
        [hashtable]$MetricNames,        # Metric names hashtable
        
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string]$SanitizedTitle         # Sanitized update title for metrics
    )
    # Track COM objects for proper cleanup
    $comObjects = @()
    try {
        Write-Information "Installing update: $($Update.Title)" -InformationAction Continue
        
        # Create and configure the update downloader
        $UpdateDownloader = $UpdateSession.CreateUpdateDownloader()
        $comObjects += $UpdateDownloader
        $UpdateDownloader.Updates = New-Object -ComObject Microsoft.Update.UpdateColl
        $comObjects += $UpdateDownloader.Updates
        $UpdateDownloader.Updates.Add($Update) # Add the update to the collection
        $UpdateDownloader.Download() # Download the update files

        # Create and configure the update installer
        $UpdateInstaller = $UpdateSession.CreateUpdateInstaller()
        $UpdateInstaller.Updates = $UpdateDownloader.Updates
        $InstallResult = $UpdateInstaller.Install() # Install the downloaded update
        
        # Generate Unix timestamp for the installation
        $installTimestamp = [Math]::Floor([decimal](Get-Date(Get-Date).ToUniversalTime()-uformat '%s'))
        
        # Create Prometheus metrics for the installation
        $additionalMetrics = @()
        $lastMetric = "# HELP $($MetricNames.LastInstall) Unix timestamp when updates were last installed`n# TYPE $($MetricNames.LastInstall) gauge"
        $additionalMetrics += $lastMetric
        $additionalMetrics += "$($MetricNames.LastInstall){update=`"$SanitizedTitle`"} $installTimestamp"
        $installMetric = "# HELP $($MetricNames.InstallStatus) Status of update installation`n# TYPE $($MetricNames.InstallStatus) gauge"
        $additionalMetrics += $installMetric
        $additionalMetrics += "$($MetricNames.InstallStatus){update=`"$SanitizedTitle`",result=`"$($InstallResult.ResultCode)`"} 1"
        
        Write-Verbose "Installation completed with result: $($InstallResult.ResultCode)"
        
        # Return success result object
        return [PSCustomObject]@{
            Success = $InstallResult.ResultCode -eq 2  # ResultCode 2 = successful installation
            ResultCode = $InstallResult.ResultCode
            UpdateTitle = $Update.Title
            InstallTimestamp = $installTimestamp
            ErrorMessage = $null
            Metrics = $additionalMetrics
        }
    }
    catch {
        $errorMessage = $_.Exception.Message
        Write-Host "Error installing update: $_"
        
        # Return error result object
        return [PSCustomObject]@{
            Success = $false
            ResultCode = -1
            UpdateTitle = $Update.Title
            InstallTimestamp = $null
            ErrorMessage = $errorMessage
            Metrics = @()
        }
    }
    finally {
        # Clean up COM objects to prevent memory leaks
        $comObjects | ForEach-Object {
            if ($_ -ne $null) {
                try {
                    [System.Runtime.InteropServices.Marshal]::ReleaseComObject($_) | Out-Null
                } catch {
                    Write-Warning "Failed to release COM object: $_"
                }
            }
        }
    }
}

# Function to check if a system reboot is required after updates
function Get-RebootStatus {
    try {
        # Check if Windows OS indicates a reboot is pending
        $osRebootPending = (Get-CimInstance -ClassName Win32_OperatingSystem -ErrorAction Stop).RebootPending
        # Check Component Based Servicing registry key for pending operations
        $cbsRebootPending = Get-ItemProperty -Path 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Component Based Servicing' -Name 'RebootPending' -ErrorAction Stop
        # Return 1 if either check indicates reboot required, 0 otherwise
        return [int]($osRebootPending -or $cbsRebootPending)
    }
    catch {
        Write-Warning "Could not determine reboot status: $($_.Exception.Message)"
        return -1  # Return -1 to indicate indeterminate status
    }
}

# Check current reboot status for inclusion in metrics
$rebootRequired = Get-RebootStatus

# Function to sanitize strings for safe use as Prometheus labels
function ConvertTo-SafePrometheusLabel {
    param(
        [Parameter(Mandatory=$false)]
        [AllowEmptyString()]
        [string]$Value        # String value to sanitize
    )
    
    # Return empty string if input is null or empty
    if ([string]::IsNullOrEmpty($Value)) {
        return ""
    }
    
    # Sanitize special characters that could break Prometheus format
    # Replace backslashes, quotes, newlines, tabs, and non-printable characters
    $sanitized = $Value -replace '[\\"\r\n\t]|[^\x20-\x7E]', { 
        switch ($_.Value) {
            '\' { '\\' }     # Escape backslashes
            '"' { '\"' }     # Escape double quotes
            default { if ($_.Value -match '[\r\n\t]') { ' ' } else { '_' } }  # Replace whitespace/unprintable with underscore
        }
    }
    
    # Limit length to prevent excessively long labels (Prometheus best practice)
    return $sanitized.Substring(0, [Math]::Min($sanitized.Length, 256))
}

# Function to build Prometheus metric strings and add them to StringBuilder
function Add-PrometheusMetric {
    param(
        [Parameter(Mandatory)]
        [System.Text.StringBuilder]$StringBuilder,    # StringBuilder to append metrics to
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string]$MetricName,                         # Name of the Prometheus metric
        [Parameter(Mandatory)]
        [string]$HelpText,                           # Help text describing the metric
        [Parameter(Mandatory)]
        $Value,                                      # Metric value
        [hashtable]$Labels = @{}                     # Optional labels for the metric
    )
    
    try {
        # Add Prometheus HELP comment explaining what the metric measures
        $StringBuilder.AppendLine("# HELP $MetricName $HelpText") | Out-Null
        # Add Prometheus TYPE comment (assuming gauge type for all metrics)
        $StringBuilder.AppendLine("# TYPE $MetricName gauge") | Out-Null
        
        if ($null -ne $Value) {
            # Build label string if labels are provided
            $labelString = ""
            if ($Labels.Count -gt 0) {
                $labelPairs = $Labels.GetEnumerator() | ForEach-Object { 
                    "$($_.Key)=`"$($_.Value)`"" 
                }
                $labelString = "{$($labelPairs -join ',')}"
            }
            # Add the actual metric line with name, labels, and value
            $StringBuilder.AppendLine("$MetricName$labelString $Value") | Out-Null
        }
    }
    catch {
        Write-Error "Failed to add Prometheus metric '$MetricName': $_"
        throw
    }
}

# Function to generate complete Prometheus metrics output
function Get-PrometheusMetrics {
    param(
        [int]$UpdateCount,              # Number of available updates
        [double]$ScriptRunTime,         # Time taken to run the script
        [bool]$rebootRequired,          # Whether a system reboot is required
        [bool]$AutoUpdateEnabled,       # Whether automatic updates are enabled
        [array]$AvailableUpdates = @()  # Array of available update details
    )
    
    # Validate input parameters
    if ($UpdateCount -lt 0) {
        throw "UpdateCount must be non-negative"
    }
    if ($ScriptRunTime -lt 0) {
        throw "ScriptRunTime must be non-negative"
    }

    try {
        # Get current Unix timestamp for the script execution time
        $currentTimestamp = [DateTimeOffset]::UtcNow.ToUnixTimeSeconds()
    } catch {
        Write-Error "Failed to initialize metrics generation: $_"
        return $null
    }

    # Convert boolean values to integers for Prometheus (0 or 1)
    $rebootValue = [int]$rebootRequired
    $autoUpdateValue = [int]$AutoUpdateEnabled

    # Define all metrics to be generated with their properties
    $metrics = @(
        @{ Name = $MetricNames.Available; Help = "Number of Windows updates available from WSUS"; Value = $UpdateCount }
        @{ Name = $MetricNames.Status; Help = "Indicates if the update script has run successfully"; Value = 1; Labels = @{status="success"} }
        @{ Name = $MetricNames.Timestamp; Help = "Unix timestamp when the script last ran"; Value = $currentTimestamp }
        @{ Name = $MetricNames.Error; Help = "Information about script errors"; Value = 0 }
        @{ Name = $MetricNames.Reboot; Help = "Indicates if a reboot is required after updates"; Value = $rebootValue }
        @{ Name = $MetricNames.Runtime; Help = "Time taken to execute the update script in seconds"; Value = $ScriptRunTime }
        @{ Name = $MetricNames.AutoUpdate; Help = "Indicates if automatic update installation is enabled"; Value = $autoUpdateValue }
        @{ Name = $MetricNames.LastInstall; Help = "Unix timestamp when updates were last installed"; Value = 0 }
        @{ Name = $MetricNames.InstallStatus; Help = "Status of update installation"; Value = 0 }
    )

    # Build the complete metrics string
    $metricsBuilder = [System.Text.StringBuilder]::new()
    foreach ($metric in $metrics) {
        try {
            Add-PrometheusMetric -StringBuilder $metricsBuilder -MetricName $metric.Name -HelpText $metric.Help -Value $metric.Value -Labels $metric.Labels
        } catch {
            Write-Warning "Failed to add metric '$($metric.Name)': $_"
        }
    }
    
    # Add individual update metrics - one per available update
    if ($AvailableUpdates.Count -gt 0) {
        $metricsBuilder.AppendLine("# HELP $($MetricNames.UpdateList) Individual Windows updates available for installation") | Out-Null
        $metricsBuilder.AppendLine("# TYPE $($MetricNames.UpdateList) gauge") | Out-Null
        
        foreach ($update in $AvailableUpdates) {
            $labels = @{
                title = $update.Title
                severity = $update.Severity
                downloaded = $update.IsDownloaded
                size_bytes = $update.Size
            }
            $labelString = ($labels.GetEnumerator() | ForEach-Object { "$($_.Key)=`"$($_.Value)`"" }) -join ','
            $metricsBuilder.AppendLine("$($MetricNames.UpdateList){$labelString} 1") | Out-Null
        }
    }
    
    return $metricsBuilder.ToString()
}

# Validate that metric names configuration is available
if (-not $MetricNames) {
    Write-Error "MetricNames configuration is not available"
    return $null
}

# Graceful shutdown flag
$global:shutdown = $false

# Register event handler for graceful shutdown
Register-EngineEvent -SourceIdentifier "PowerShell.Exiting" -Action {
    Write-Host "Initiating graceful shutdown..."
    $global:shutdown = $true
}

# Trap Ctrl+C and other termination signals
[Console]::TreatControlCAsInput = $false
Register-ObjectEvent -InputObject ([Console]) -EventName CancelKeyPress -Action {
    param($Sender, $CancelEventArgs)
    Write-Host "Shutdown signal received. Cleaning up..."
    $CancelEventArgs.Cancel = $true
    $global:shutdown = $true
}

# Function to check for shutdown signal during long operations
function Test-ShutdownSignal {
    if ($global:shutdown) {
        Write-Host "Shutdown requested. Exiting gracefully..."
        exit 0
    }
}

# Main execution: Check for Windows Updates from WSUS
Write-Host "Checking for Windows Updates from WSUS..."
$StartTime = Get-Date
$comObjects = @()
$availableUpdates = @()
try {
    # Check for shutdown before starting
    Test-ShutdownSignal
    
    # Create Windows Update session and searcher COM objects
    $UpdateSession = New-Object -ComObject Microsoft.Update.Session
    $comObjects += $UpdateSession
    $UpdateSearcher = $UpdateSession.CreateUpdateSearcher()
    $comObjects += $UpdateSearcher
    
    # Check for shutdown before search
    Test-ShutdownSignal
    
    # Search for updates that are not yet installed
    $SearchResult = $UpdateSearcher.Search("IsInstalled=0")
    
    # Process individual updates for detailed metrics
    foreach ($Update in $SearchResult.Updates) {
        $sanitizedTitle = ConvertTo-SafePrometheusLabel -Value $Update.Title
        $sanitizedDescription = ConvertTo-SafePrometheusLabel -Value $Update.Description
        $updateSize = if ($Update.MaxDownloadSize) { $Update.MaxDownloadSize } else { 0 }
        
        $availableUpdates += [PSCustomObject]@{
            Title = $sanitizedTitle
            Description = $sanitizedDescription
            Size = $updateSize
            Severity = if ($Update.MsrcSeverity) { $Update.MsrcSeverity } else { "Unknown" }
            IsDownloaded = [int]$Update.IsDownloaded
        }
    }
} catch {
    # Handle connection failures and write error metric
    Write-Error "Failed to connect to Windows Update service: $_"
    $errorMetric = "$($MetricNames.Error){error=`"wsus_connection_failed`"} 1"
    $errorMetric | Out-File -FilePath $MetricsFilePath -Encoding UTF8 -Force
    exit 1
} finally {
    # Clean up COM objects to prevent memory leaks
    $comObjects | ForEach-Object {
        if ($_ -and [System.Runtime.InteropServices.Marshal]::IsComObject($_)) {
            try {
                [System.Runtime.InteropServices.Marshal]::ReleaseComObject($_) | Out-Null
            } catch {
                Write-Warning "Failed to release COM object: $_"
            }
        }
    }
}

# Calculate script execution time and generate base metrics
$EndTime = Get-Date
$ScriptRunTime = ($EndTime - $StartTime).TotalSeconds
# Safely get update count even if SearchResult is null
$updateCount = if ($SearchResult -and $SearchResult.Updates) { $SearchResult.Updates.Count } else { 0 }
$prometheusMetric = Get-PrometheusMetrics -UpdateCount $updateCount -ScriptRunTime $ScriptRunTime -AutoUpdateEnabled $AutoInstall -AvailableUpdates $availableUpdates

# Initialize array for additional metrics from update installations
$additionalMetrics = @()

# Process the search results
if ($SearchResult.Updates.Count -eq 0) {
    Write-Host "No updates available from WSUS."
} else {
    Write-Host "Found $($SearchResult.Updates.Count) update(s) available:"
    $updateList = [System.Collections.ArrayList]::new()
    
    # Iterate through each available update
    foreach ($Update in $SearchResult.Updates) {
        # Check for shutdown signal between updates
        Test-ShutdownSignal
        
        Write-Host "- $($Update.Title)"
        $sanitizedTitle = ConvertTo-SafePrometheusLabel -Value $Update.Title
        [void]$updateList.Add("title=`"$sanitizedTitle`"")
        
        # Install the update if AutoInstall is enabled
        if ($AutoInstall) {
            # Check for shutdown before installing
            Test-ShutdownSignal
            
            $result = Install-WindowsUpdate -Update $Update -UpdateSession $UpdateSession -MetricNames $MetricNames -SanitizedTitle $sanitizedTitle
            if ($result.Metrics) {
                $additionalMetrics += $result.Metrics
            }
        }
    }
}

# Write all metrics to the output file for Prometheus collection
try {
    # Final shutdown check before writing metrics
    Test-ShutdownSignal
    
    $allMetrics = @($prometheusMetric) + $additionalMetrics
    ($allMetrics -join "`n") | Out-File -FilePath "$MetricsFilePath" -Encoding UTF8 -Force
    Write-Host "Metrics successfully written to $MetricsFilePath"
} catch {
    # Handle file write errors and create error metric
    Write-Error "Failed to write metrics to file: $_"
    $errorCode = $_.Exception.HResult -replace '"', '\"'
    "$($MetricNames.Error){error=`"file_write_failed`",error_code=`"$errorCode`"} 1" | Out-File -FilePath "$MetricsFilePath" -Encoding UTF8 -Force
}

Write-Host "Script completed successfully."