PSModule_SpeculationControl.psm1

#
# PSModule_SpeculationControl.psm1
# Module implementation for speculative execution mitigation automation.
# Author: Krishnaramanan.S
# License: No restriction in use or modify
# Tested in Dev Environment only. Test this module in your Dev Environment first before using in Production. User must make changes if required on their own risk.
#

#region Public Functions

function Start-SpeculativeMitigation {
    <#
    .SYNOPSIS
        Applies registry-based mitigations and optionally installs required KBs.
    
    .DESCRIPTION
        Consolidates Fix-SpeculativeMitigations.ps1 functionality into a module function.
        Supports Level-based presets, granular CVE selection, KB installation,
        XML scan consumption, backup/restore, and user approval flows.
    
    .PARAMETER Level
        Mitigation level: 1, 2, or 3. Default is 1.
    
    .PARAMETER CVE
        Granular list of CVEs to fix. Overrides Level if specified.
    
    .PARAMETER DisableHyperThreading
        Disables Hyper-Threading in the registry for L1TF mitigation.
    
    .PARAMETER Backup
        Creates a registry backup before making changes.
    
    .PARAMETER Restore
        Restores the most recent backup.
    
    .PARAMETER InstallKBs
        Check for and install required Microsoft KBs.
    
    .PARAMETER ScanOnly
        Scan for missing KBs without installing.
    
    .PARAMETER XmlReportPath
        Path to XML report from Check-SpeculativeMitigations.ps1.
    
    .PARAMETER Force
        Skip confirmation prompts. Use with caution.
    
    .EXAMPLE
        Start-SpeculativeMitigation -Level 1
    
    .EXAMPLE
        Start-SpeculativeMitigation -CVE CVE-2017-5715,CVE-2018-3639 -WhatIf
    
    .EXAMPLE
        Start-SpeculativeMitigation -XmlReportPath "C:\Reports\speculation.xml" -InstallKBs
    #>

    [CmdletBinding(SupportsShouldProcess=$true, ConfirmImpact='High')]
    param(
        [ValidateSet(1, 2, 3)]
        [int]$Level = 1,
        [string[]]$CVE,
        [switch]$DisableHyperThreading,
        [switch]$Backup,
        [switch]$Restore,
        [switch]$InstallKBs,
        [switch]$ScanOnly,
        [string]$XmlReportPath,
        [switch]$Force
    )
    
    # Import the main script functions into the module scope
    $scriptPath = Join-Path $PSScriptRoot 'Fix-SpeculativeMitigations.ps1'
    if (-not (Test-Path $scriptPath)) {
        throw "Dependency script not found: $scriptPath"
    }
    
    # Dot-source the script to import all functions
    . $scriptPath
    
    # Call Main with the provided parameters
    $PSBoundParams['Level'] = $Level
    if ($CVE) { $PSBoundParams['CVE'] = $CVE }
    if ($DisableHyperThreading) { $PSBoundParams['DisableHyperThreading'] = $true }
    if ($Backup) { $PSBoundParams['Backup'] = $true }
    if ($Restore) { $PSBoundParams['Restore'] = $true }
    if ($InstallKBs) { $PSBoundParams['InstallKBs'] = $true }
    if ($ScanOnly) { $PSBoundParams['ScanOnly'] = $true }
    if ($XmlReportPath) { $PSBoundParams['XmlReportPath'] = $XmlReportPath }
    if ($Force) { $PSBoundParams['Force'] = $true }
    
    # Execute Main with bound parameters
    Main
}

function Get-SpeculativeMitigationStatus {
    <#
    .SYNOPSIS
        Runs Check-SpeculativeMitigations.ps1 and returns the report object.
    
    .DESCRIPTION
        Executes the check script and optionally exports XML output.
        Returns the mitigation status report for further processing.
    
    .PARAMETER OutputXml
        Export results to XML file.
    
    .PARAMETER OutputPath
        Path for XML output. Defaults to TEMP directory.
    
    .EXAMPLE
        $status = Get-SpeculativeMitigationStatus
    
    .EXAMPLE
        $status = Get-SpeculativeMitigationStatus -OutputXml -OutputPath "C:\Reports\speculation.xml"
    #>

    [CmdletBinding()]
    param(
        [switch]$OutputXml,
        [string]$OutputPath
    )
    
    $checkScript = Join-Path $PSScriptRoot 'Check-SpeculativeMitigations.ps1'
    if (-not (Test-Path $checkScript)) {
        throw "Check script not found: $checkScript"
    }
    
    $args = @()
    if ($OutputXml) { $args += '-OutputXml' }
    if ($OutputPath) { $args += "-OutputPath", $OutputPath }
    
    # Execute the check script
    & $checkScript @args
}

function Restore-SpeculativeMitigationRegistry {
    <#
    .SYNOPSIS
        Restores registry from the most recent backup.
    
    .DESCRIPTION
        Finds and restores the latest registry backup created by Start-SpeculativeMitigation -Backup.
    #>

    [CmdletBinding()]
    param()
    
    $scriptPath = Join-Path $PSScriptRoot 'Fix-SpeculativeMitigations.ps1'
    if (-not (Test-Path $scriptPath)) {
        throw "Dependency script not found: $scriptPath"
    }
    
    . $scriptPath
    
    $latestBackup = Get-ChildItem -Path $env:TEMP -Filter "SpeculativeMitigationBackup_*.json" -ErrorAction SilentlyContinue | 
                    Sort-Object LastWriteTime -Descending | Select-Object -First 1
    if ($latestBackup) {
        Restore-Registry -BackupFile $latestBackup.FullName
        Write-Host ''
        Write-Host 'A system reboot is required for changes to take effect.' -ForegroundColor Red
        Write-Host 'Reboot command: shutdown /r /t 0'
    } else {
        Write-Warning 'No backup found to restore.'
    }
}

function Test-SpeculativeMitigationAdministrator {
    <#
    .SYNOPSIS
        Tests if the current session has administrator privileges.
    
    .DESCRIPTION
        Returns True if running as Administrator, False otherwise.
        Required for all mitigation operations.
    #>

    [CmdletBinding()]
    param()
    
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]::new($identity)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Get-VulnerableCvesFromXml {
    <#
    .SYNOPSIS
        Parses XML report from Check-SpeculativeMitigations.ps1 and returns vulnerable CVEs.
    
    .DESCRIPTION
        Reads the XML output from the check script and extracts only CVEs
        that are vulnerable, disabled, unsupported, or require updates.
    
    .PARAMETER XmlPath
        Path to the XML report file.
    
    .EXAMPLE
        $vulns = Get-VulnerableCvesFromXml -XmlPath "C:\Reports\speculation.xml"
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)]
        [string]$XmlPath
    )
    
    $scriptPath = Join-Path $PSScriptRoot 'Fix-SpeculativeMitigations.ps1'
    if (-not (Test-Path $scriptPath)) {
        throw "Dependency script not found: $scriptPath"
    }
    
    . $scriptPath
    
    return Get-VulnerableCvesFromXml -XmlPath $XmlPath
}

#endregion

#region Module Initialization

Write-Verbose "PSModule_SpeculationControl module loaded. Run Get-Command -Module PSModule_SpeculationControl to see available functions."

#endregion