Fix-SpeculativeMitigations.ps1

<#
.SYNOPSIS
    Comprehensive remediation for speculative execution vulnerabilities on Windows Server.

.DESCRIPTION
    Applies registry-based mitigations, installs required Windows KBs, and consumes
    scan results from Check-SpeculativeMitigations.ps1 XML output.

    Registry-fixable CVEs:
      CVE-2017-5715 (BTI), CVE-2017-5754 (RDCL / Meltdown),
      CVE-2018-3639 (SSBD), CVE-2018-3620/3646 (L1TF),
      CVE-2018-11091/12126/12127/12130 (MDS),
      CVE-2022-21123 (SBDR), CVE-2022-21125 (SBDS),
      CVE-2022-0001 (BHI), CVE-2022-23825 (Retbleed)

    OS-update-only CVEs (guidance only):
      CVE-2022-21127 (SRBDS), CVE-2022-21166 (DRPW),
      CVE-2023-20569 (SRSO), CVE-2023-20588 (DivByZero),
      CVE-2023-28746 (RFDS), CVE-2022-40982 (GDS)

.PARAMETER Level
    Mitigation level: 1, 2, or 3. Default is 1.
      1 = All registry-fixable CVEs, Hyper-Threading enabled (Recommended)
      2 = All registry-fixable CVEs, Hyper-Threading disabled
      3 = CVE-2022-0001 (BHI) only

.PARAMETER CVE
    Granular list of CVEs to fix. Overrides Level if specified.
    Example: -CVE CVE-2017-5715,CVE-2018-3639,CVE-2022-0001

.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 for OS-update-only CVEs.
    Requires user approval. May require system reboot.

.PARAMETER ScanOnly
    Scan for missing KBs and report status without installing.

.PARAMETER XmlReportPath
    Path to XML report from Check-SpeculativeMitigations.ps1.
    When provided, uses scan results to determine required actions.

.PARAMETER Force
    Skip confirmation prompts. Use with caution.

.PARAMETER WhatIf
    Shows what would be done without making changes.

.PARAMETER Confirm
    Prompts for confirmation before each change.

.EXAMPLE
    .\Fix-SpeculativeMitigations.ps1 -Level 1

.EXAMPLE
    .\Fix-SpeculativeMitigations.ps1 -CVE CVE-2017-5715,CVE-2018-3639 -WhatIf

.EXAMPLE
    .\Fix-SpeculativeMitigations.ps1 -Restore

.EXAMPLE
    .\Fix-SpeculativeMitigations.ps1 -Level 1 -InstallKBs

.EXAMPLE
    .\Fix-SpeculativeMitigations.ps1 -XmlReportPath "C:\Reports\speculation.xml" -InstallKBs

.NOTES
    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.
    Requires administrator privileges.
    A system reboot is required for registry changes to take effect.
    KB installation may also require system reboot.
    All changes require user approval unless -Force is specified.
#>


[CmdletBinding(SupportsShouldProcess=$true, ConfirmImpact='High')]
param(
    [Parameter(Mandatory=$false, ParameterSetName='Level')]
    [ValidateSet(1, 2, 3)]
    [int]$Level = 1,

    [Parameter(Mandatory=$true, ParameterSetName='CVE')]
    [string[]]$CVE,

    [Parameter(Mandatory=$false)]
    [switch]$DisableHyperThreading,

    [Parameter(Mandatory=$false)]
    [switch]$Backup,

    [Parameter(Mandatory=$false)]
    [switch]$Restore,

    [Parameter(Mandatory=$false)]
    [switch]$InstallKBs,

    [Parameter(Mandatory=$false)]
    [switch]$ScanOnly,

    [Parameter(Mandatory=$false)]
    [string]$XmlReportPath,

    [Parameter(Mandatory=$false)]
    [switch]$Force
)

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

$RegistryPath = 'HKLM:\SYSTEM\CurrentControlSet\Control\Session Manager\Memory Management'
$BackupPath = 'HKLM:\SYSTEM\SpeculativeMitigationBackup'

# Microsoft-recommended registry presets for Windows Server
$Presets = [ordered]@{
    1 = [ordered]@{
        OverrideMask = 3
        Override     = 0x00800048
        Description  = 'Level 1: All speculative execution mitigations with Hyper-Threading enabled (Recommended for most servers)'
    }
    2 = [ordered]@{
        OverrideMask = 3
        Override     = 0x00802048
        Description  = 'Level 2: All speculative execution mitigations with Hyper-Threading disabled'
    }
    3 = [ordered]@{
        OverrideMask = 3
        Override     = 0x00800000
        Description  = 'Level 3: CVE-2022-0001 (Intel BHI) only. Use only if other mitigations are managed separately.'
    }
}

# CVE to FeatureSettingsOverride bit mapping
$CveToBit = [ordered]@{
    'CVE-2017-5715'  = 0x00800000  # BTI - part of BHI bundle
    'CVE-2017-5754'  = 0x00000001  # RDCL - KVA shadow
    'CVE-2018-3639'  = 0x00000002  # SSBD
    'CVE-2018-3620'  = 0x00000004  # L1TF
    'CVE-2018-11091' = 0x00000008  # MDS
    'CVE-2018-12126' = 0x00000008  # MDS
    'CVE-2018-12127' = 0x00000008  # MDS
    'CVE-2018-12130' = 0x00000008  # MDS
    'CVE-2022-21123' = 0x00000010  # SBDR
    'CVE-2022-21125' = 0x00000010  # SBDS
    'CVE-2022-0001'  = 0x00800000  # BHI
    'CVE-2022-23825' = 0x00000020  # Retbleed
}

# KB requirements per mitigation level
$KbRequirements = @{
    1 = @{
        KbNumbers = @('KB5006670', 'KB5006669', 'KB5006672')
        Description = 'Mitigations for Spectre, Meltdown, and L1TF'
    }
    2 = @{
        KbNumbers = @('KB5006670', 'KB5006669')
        Description = 'Mitigations for Spectre and Meltdown'
    }
    3 = @{
        KbNumbers = @()
        Description = 'No additional KB requirements for BHI-only mitigation'
    }
}

# OS-specific KB requirements for SRBDS/DRPW
$OsKbRequirements = @{
    'Windows Server 2022' = 'KB5015018'
    'Windows Server 2019' = 'KB5015017'
    'Windows Server 2016' = 'KB5015019'
    'Windows 10 21H2'     = 'KB5015020'
    'Windows 11 21H2'     = 'KB5015020'
}

function Test-Administrator {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]::new($identity)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Get-RegistryValue {
    param([string]$Name)
    try {
        $val = Get-ItemProperty -Path $RegistryPath -Name $Name -ErrorAction SilentlyContinue
        if ($null -ne $val -and $null -ne $val.$Name) {
            return [uint32]$val.$Name
        }
    }
    catch { }
    return $null
}

function Set-RegistryDword {
    param(
        [string]$Name,
        [uint32]$Value
    )
    try {
        if (-not (Test-Path $RegistryPath)) {
            New-Item -Path $RegistryPath -Force | Out-Null
        }
        Set-ItemProperty -Path $RegistryPath -Name $Name -Value $Value -Type DWord -Force
        return $true
    }
    catch {
        Write-Error "Failed to set $Name`: $_"
        return $false
    }
}

function Backup-Registry {
    Write-Host 'Creating registry backup...' -ForegroundColor Yellow
    try {
        if (-not (Test-Path $BackupPath)) {
            New-Item -Path $BackupPath -Force | Out-Null
        }
        
        $backupData = [ordered]@{
            Timestamp = Get-Date
            ComputerName = $env:COMPUTERNAME
            FeatureSettingsOverride = Get-RegistryValue -Name 'FeatureSettingsOverride'
            FeatureSettingsOverrideMask = Get-RegistryValue -Name 'FeatureSettingsOverrideMask'
            SpeculationControlFlags = Get-RegistryValue -Name 'SpeculationControlFlags'
        }
        
        $backupJson = $backupData | ConvertTo-Json
        $backupFile = Join-Path $env:TEMP "SpeculativeMitigationBackup_$(Get-Date -Format 'yyyyMMdd_HHmmss').json"
        $backupJson | Out-File -FilePath $backupFile -Encoding UTF8
        
        Write-Host "Backup saved to: $backupFile" -ForegroundColor Green
        return $backupFile
    }
    catch {
        Write-Warning "Backup failed: $_"
        return $null
    }
}

function Restore-Registry {
    param([string]$BackupFile)
    
    if (-not (Test-Path $BackupFile)) {
        Write-Error "Backup file not found: $BackupFile"
        return $false
    }
    
    Write-Host "Restoring from backup: $BackupFile" -ForegroundColor Yellow
    try {
        $backupData = Get-Content -Path $BackupFile -Encoding UTF8 | ConvertFrom-Json
        
        if ($backupData.FeatureSettingsOverride -ne $null) {
            Set-RegistryDword -Name 'FeatureSettingsOverride' -Value $backupData.FeatureSettingsOverride
        }
        if ($backupData.FeatureSettingsOverrideMask -ne $null) {
            Set-RegistryDword -Name 'FeatureSettingsOverrideMask' -Value $backupData.FeatureSettingsOverrideMask
        }
        if ($backupData.SpeculationControlFlags -ne $null) {
            Set-RegistryDword -Name 'SpeculationControlFlags' -Value $backupData.SpeculationControlFlags
        }
        
        Write-Host 'Registry restored successfully.' -ForegroundColor Green
        return $true
    }
    catch {
        Write-Error "Restore failed: $_"
        return $false
    }
}

function Test-HyperThreadingEnabled {
    $processors = Get-CimInstance Win32_Processor -ErrorAction SilentlyContinue
    if ($null -eq $processors) {
        return $false
    }
    $totalCores = 0
    $totalLogical = 0
    foreach ($proc in $processors) {
        $totalCores += [uint32]$proc.NumberOfCores
        $totalLogical += [uint32]$proc.NumberOfLogicalProcessors
    }
    return ($totalLogical -gt $totalCores)
}

function Get-CveRecommendation {
    param([string]$CveId)
    
    $recommendations = @{
        'CVE-2017-5715'  = 'Registry: Set FeatureSettingsOverride/Mask (Level 1, 2, or 3)'
        'CVE-2017-5753'  = 'Application: Recompile with /Qspectre or /guard:cf'
        'CVE-2017-5754'  = 'Registry: Enable KVA shadow via FeatureSettingsOverride'
        'CVE-2018-3639'  = 'Registry: Enable SSBD system-wide'
        'CVE-2018-3615'  = 'Registry: Enable L1TF mitigation'
        'CVE-2018-3620'  = 'Registry: Enable L1TF mitigation + disable Hyper-Threading'
        'CVE-2018-3646'  = 'Registry: Enable L1TF mitigation'
        'CVE-2018-11091' = 'Registry: Enable MDS mitigation (MB clear)'
        'CVE-2018-12126' = 'Registry: Enable MDS mitigation (MB clear)'
        'CVE-2018-12127' = 'Registry: Enable MDS mitigation (MB clear)'
        'CVE-2018-12130' = 'Registry: Enable MDS mitigation (MB clear)'
        'CVE-2022-21123' = 'Registry: Enable FB clear for SBDR'
        'CVE-2022-21125' = 'Registry: Enable FB clear for SBDS'
        'CVE-2022-21127' = 'Windows Update: Install KB5015018 (Server 2022), KB5015017 (Server 2019), or KB5015019 (Server 2016)'
        'CVE-2022-21166' = 'Windows Update: Install KB5015018 (Server 2022), KB5015017 (Server 2019), or KB5015019 (Server 2016)'
        'CVE-2022-0001'  = 'Registry: Enable BHI via FeatureSettingsOverride (Level 1 or 3)'
        'CVE-2022-23825' = 'Registry: Enable BTI + kernel retpoline'
        'CVE-2023-20569' = 'OS Update: Install latest cumulative update + microcode'
        'CVE-2023-20588' = 'OS Update: Install latest cumulative update'
        'CVE-2023-28746' = 'OS Update: Install latest cumulative update + microcode'
        'CVE-2022-40982' = 'BIOS/Microcode: Update CPU microcode if hardware supports it'
    }
    
    if ($recommendations.ContainsKey($CveId)) {
        return $recommendations[$CveId]
    }
    return 'Unknown CVE - consult Microsoft documentation'
}

function Get-InstalledKbNumbers {
    try {
        $kbNumbers = @()
        $updates = Get-CimInstance Win32_QuickFixEngineering -ErrorAction SilentlyContinue
        if ($updates) {
            foreach ($update in $updates) {
                if ($update.HotFixID) {
                    $kbNumbers += $update.HotFixID
                }
            }
        }
        return $kbNumbers | Sort-Object
    }
    catch {
        return @()
    }
}

function Test-KbInstalled {
    param([string]$kbNumber)
    $installedKbs = Get-InstalledKbNumbers
    return $installedKbs -contains $kbNumber
}

function Test-PendingReboot {
    try {
        $rebootRequired = $false
        
        $lockReg = Get-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Control\Session Manager' -Name 'PendingFileRenameOperations' -ErrorAction SilentlyContinue
        if ($lockReg -and $lockReg.'PendingFileRenameOperations') {
            $rebootRequired = $true
        }
        
        return $rebootRequired
    }
    catch {
        return $false
    }
}

function Test-DiskSpace {
    param([int]$MbRequired)
    try {
        $cDrive = Get-PSDrive -Name C -ErrorAction SilentlyContinue
        if ($cDrive) {
            $freeSpaceMb = [math]::Floor($cDrive.Free / 1MB)
            return $freeSpaceMb -ge $MbRequired
        }
        return $false
    }
    catch {
        return $false
    }
}

function Install-KbUpdate {
    param([string]$kbNumber)

    Write-Host "Checking prerequisites for KB$kbNumber installation..." -ForegroundColor Cyan

    $pendingReboot = Test-PendingReboot
    if ($pendingReboot) {
        Write-Host 'WARNING: Pending system reboot detected.' -ForegroundColor Yellow
        Write-Host 'It is recommended to reboot the system before installing KB$kbNumber.' -ForegroundColor Yellow
        if (-not $Force) {
            $rebootChoice = Read-Host 'Do you want to reboot now? (Y/N)'
            if ($rebootChoice -match '^[Yy]$') {
                Write-Host 'Rebooting system...' -ForegroundColor Cyan
                shutdown /r /t 0
            }
            else {
                Write-Host 'Continuing without reboot. Installation may fail.' -ForegroundColor Yellow
            }
        }
        else {
            Write-Host 'Force mode: Rebooting without prompt...' -ForegroundColor Yellow
            shutdown /r /t 0
        }
    }

    if (-not (Test-DiskSpace 200)) {
        Write-Host 'WARNING: Insufficient disk space on C: drive for KB installation.' -ForegroundColor Red
        Write-Host 'At least 200 MB free space required.' -ForegroundColor Red
        return $false
    }

    Write-Host "Attempting to install KB$kbNumber..." -ForegroundColor Cyan

    try {
        $updateSession = New-Object -ComObject "Microsoft.Update.Session"
        $updateSearcher = $updateSession.CreateUpdateSearcher()
        $searchResult = $updateSearcher.Search("IsInstalled=0 AND Title:*$kbNumber*")

        if ($searchResult.Updates.Count -gt 0) {
            $update = $searchResult.Updates[0]
            $update.AcceptEndUserLicenseAgreement = $true

            $downloader = $updateSession.CreateUpdateDownloader()
            $downloader.Updates.Add($update)
            $downloadResult = $downloader.Download()

            Write-Host "Downloaded KB$kbNumber updates." -ForegroundColor Green

            $installer = $updateSession.CreateUpdateInstaller()
            $installResult = $installer.Install($update)

            Write-Host "KB$kbNumber installation result: $($installResult.ResultCode)" -ForegroundColor Green
            return $true
        }
        else {
            Write-Host "KB$kbNumber not found via Windows Update search." -ForegroundColor Yellow
            return $false
        }
    }
    catch {
        Write-Error "KB$kbNumber installation failed: $_"
        return $false
    }
}

function Get-VulnerableCvesFromXml {
    param([string]$XmlPath)
    
    $vulnerableCves = @()
    
    if (-not (Test-Path $XmlPath)) {
        Write-Warning "XML report not found at: $XmlPath"
        return ,$vulnerableCves
    }
    
    try {
        [xml]$xmlDoc = Get-Content $XmlPath -Raw
        
        foreach ($cveNode in $xmlDoc.SpeculativeMitigationReport.Mitigations.CVE) {
            $status = $cveNode.Status
            if ($status -notmatch '^(Mitigated|Hardware Immune|Not Vulnerable)$') {
                $cveName = $cveNode.Name
                $requiredKb = $cveNode.RequiredKB
                $minBuild = $cveNode.MinimumBuild
                $curBuild = $cveNode.CurrentBuild
                $vulnerableCves += [PSCustomObject]@{
                    CveName = $cveName
                    Status = $status
                    RequiredKB = $requiredKb
                    MinimumBuild = $minBuild
                    CurrentBuild = $curBuild
                }
            }
        }
    }
    catch {
        Write-Warning "Failed to parse XML report: $_"
    }
    
    return ,$vulnerableCves
}

function Main {
    Write-Host '=== Windows Speculative Execution Mitigation Fixer ===' -ForegroundColor Cyan
    Write-Host ''
    
    if (-not (Test-Administrator)) {
        Write-Error 'This script requires administrator privileges. Please run as Administrator.'
        exit 1
    }
    
    if ($Restore) {
        $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.'
        }
        return
    }
    
    if ($Backup) {
        $backupFile = Backup-Registry
        if ($backupFile) {
            Write-Host "Backup created: $backupFile" -ForegroundColor Green
        }
        if (-not $Force -and -not $PSCmdlet.ShouldContinue('Continue with applying fixes?', 'Apply Mitigations')) {
            return
        }
    }
    
    $htEnabled = Test-HyperThreadingEnabled
    Write-Host "Hyper-Threading Detected: $(if ($htEnabled) { 'ENABLED' } else { 'DISABLED' })" -ForegroundColor $(if ($htEnabled) { 'Green' } else { 'Yellow' })
    Write-Host ''
    
    $currentOverride = Get-RegistryValue -Name 'FeatureSettingsOverride'
    $currentMask = Get-RegistryValue -Name 'FeatureSettingsOverrideMask'
    
    Write-Host 'Current Registry Settings' -ForegroundColor Yellow
    Write-Host " FeatureSettingsOverrideMask : 0x$($currentMask.ToString('X8'))"
    Write-Host " FeatureSettingsOverride : 0x$($currentOverride.ToString('X8'))"
    Write-Host ''
    
    # Handle XML report consumption
    $xmlVulnerableCves = @()
    if ($XmlReportPath) {
        Write-Host "Loading scan results from: $XmlReportPath" -ForegroundColor Cyan
        $xmlVulnerableCves = Get-VulnerableCvesFromXml -XmlPath $XmlReportPath
        if ($xmlVulnerableCves.Count -gt 0) {
            Write-Host "Found $($xmlVulnerableCves.Count) vulnerable CVEs in scan results:" -ForegroundColor Yellow
            foreach ($cve in $xmlVulnerableCves) {
                Write-Host " - $($cve.CveName): $($cve.Status)" -ForegroundColor Red
            }
            Write-Host ''
        }
        else {
            Write-Host 'No vulnerable CVEs found in scan results. System appears compliant.' -ForegroundColor Green
            if (-not $InstallKBs) {
                return
            }
        }
    }
    
    $targetOverride = $null
    $targetMask = $null
    $appliedPreset = $false
    
    if ($PSCmdlet.ParameterSetName -eq 'Level') {
        $preset = $Presets[$Level]
        $targetOverride = $preset.Override
        $targetMask = $preset.OverrideMask
        
        Write-Host "Applying preset: $($preset.Description)" -ForegroundColor Cyan
        $appliedPreset = $true
    }
    elseif ($PSCmdlet.ParameterSetName -eq 'CVE') {
        Write-Host "Applying granular fixes for CVEs: $($CVE -join ', ')" -ForegroundColor Cyan
        
        $calculatedOverride = 0
        foreach ($cve in $CVE) {
            if ($CveToBit.ContainsKey($cve)) {
                $calculatedOverride = $calculatedOverride -bor $CveToBit[$cve]
                Write-Host " $cve : 0x$($CveToBit[$cve].ToString('X8'))" -ForegroundColor Green
            } else {
                Write-Warning " $cve : No direct registry mapping. Requires OS update or application fix."
                Write-Host " Recommendation: $(Get-CveRecommendation -CveId $cve)" -ForegroundColor Yellow
            }
        }
        
        if ($calculatedOverride -eq 0) {
            Write-Warning 'No registry-fixable CVEs selected. Nothing to apply.'
            return
        }
        
        $targetOverride = $calculatedOverride
        $targetMask = 0xFFFFFFFF
        
        if ($DisableHyperThreading) {
            $targetOverride = $targetOverride -bor 0x00802000
            Write-Host ' Hyper-Threading disable flag added.' -ForegroundColor Yellow
        }
    }
    
    if ($null -eq $targetOverride) {
        Write-Warning 'No changes specified. Use -Level or -CVE parameter.'
        return
    }
    
    # Show planned changes
    Write-Host ''
    Write-Host 'Planned Changes' -ForegroundColor Yellow
    Write-Host " FeatureSettingsOverrideMask : 0x$($targetMask.ToString('X8')) (current: 0x$($currentMask.ToString('X8')))"
    Write-Host " FeatureSettingsOverride : 0x$($targetOverride.ToString('X8')) (current: 0x$($currentOverride.ToString('X8')))"
    
    if ($currentOverride -eq $targetOverride -and $currentMask -eq $targetMask) {
        Write-Host ''
        Write-Host 'System is already configured with the requested settings. No changes needed.' -ForegroundColor Green
    }
    else {
        # Confirm
        if (-not $Force -and -not $PSCmdlet.ShouldContinue('Apply these registry changes? A reboot is required.', 'Confirm Mitigation Changes')) {
            return
        }
        
        # Apply changes
        Write-Host ''
        Write-Host 'Applying changes...' -ForegroundColor Yellow
        
        $maskResult = $true
        $overrideResult = $true
        
        if ($currentMask -ne $targetMask) {
            if ($PSCmdlet.ShouldProcess("FeatureSettingsOverrideMask", "Set to 0x$($targetMask.ToString('X8'))")) {
                $maskResult = Set-RegistryDword -Name 'FeatureSettingsOverrideMask' -Value $targetMask
                if ($maskResult) {
                    Write-Host " FeatureSettingsOverrideMask set to 0x$($targetMask.ToString('X8'))" -ForegroundColor Green
                }
            }
        } else {
            Write-Host " FeatureSettingsOverrideMask already set to 0x$($targetMask.ToString('X8')). Skipping." -ForegroundColor Gray
        }
        
        if ($currentOverride -ne $targetOverride) {
            if ($PSCmdlet.ShouldProcess("FeatureSettingsOverride", "Set to 0x$($targetOverride.ToString('X8'))")) {
                $overrideResult = Set-RegistryDword -Name 'FeatureSettingsOverride' -Value $targetOverride
                if ($overrideResult) {
                    Write-Host " FeatureSettingsOverride set to 0x$($targetOverride.ToString('X8'))" -ForegroundColor Green
                }
            }
        } else {
            Write-Host " FeatureSettingsOverride already set to 0x$($targetOverride.ToString('X8')). Skipping." -ForegroundColor Gray
        }
        
        Write-Host ''
        
        if ($maskResult -and $overrideResult) {
            Write-Host 'All settings applied successfully.' -ForegroundColor Green
            Write-Host ''
            Write-Host 'IMPORTANT: A system reboot is required for these changes to take effect.' -ForegroundColor Red
            Write-Host 'Reboot command: shutdown /r /t 0'
            Write-Host ''
            
            if ($appliedPreset) {
                Write-Host 'Verification' -ForegroundColor Yellow
                Write-Host ' After reboot, run Check-SpeculativeMitigations.ps1 to verify mitigations are active.' -ForegroundColor Gray
            }
        }
        else {
            Write-Error 'One or more settings failed to apply. Please review the errors above.'
            exit 1
        }
    }
    
    # Handle KB installation if requested
    if ($InstallKBs -or $ScanOnly) {
        $levelKbs = $KbRequirements[$Level].KbNumbers
        
        if ($levelKbs.Count -gt 0) {
            Write-Host ''
            Write-Host "=== KB Requirements for Level ${Level}: $($KbRequirements[$Level].Description) ===" -ForegroundColor Cyan
            Write-Host ''
            
            $kbNeedsInstall = $false
            foreach ($kb in $levelKbs) {
                if (-not (Test-KbInstalled -kbNumber $kb)) {
                    Write-Host "KB$kb : NOT installed" -ForegroundColor Red
                    $kbNeedsInstall = $true
                }
                else {
                    Write-Host "KB$kb : Already installed" -ForegroundColor Green
                }
            }
            Write-Host ''
            
            if ($kbNeedsInstall -and -not $ScanOnly) {
                Write-Host 'Missing KBs detected. Initiating installation...' -ForegroundColor Yellow
                
                if (-not $Force) {
                    $approve = Read-Host 'Do you want to proceed with KB installation? (Y/N)'
                    if ($approve -ne 'Y' -and $approve -ne 'y') {
                        Write-Host 'KB installation skipped.' -ForegroundColor Yellow
                        return
                    }
                }
                
                foreach ($kb in $levelKbs) {
                    if (-not (Test-KbInstalled -kbNumber $kb)) {
                        Write-Host "Installing KB$kb..." -ForegroundColor Cyan
                        $pendingReboot = Test-PendingReboot
                        if ($pendingReboot) {
                            Write-Host 'Pending system reboot detected. Rebooting before KB installation...' -ForegroundColor Yellow
                            shutdown /r /t 0
                        }
                        $installed = Install-KbUpdate -kbNumber $kb
                        if (-not $installed) {
                            Write-Error "Failed to install KB$kb." -ForegroundColor Red
                        }
                    }
                }
            }
            elseif ($kbNeedsInstall -and $ScanOnly) {
                Write-Host 'Scan mode: KBs would be installed. Use -InstallKBs to proceed.' -ForegroundColor Yellow
            }
        }
        else {
            Write-Host 'No additional KB requirements for this mitigation level.' -ForegroundColor Green
        }
        
        # Check OS-specific KBs (SRBDS/DRPW)
        $os = Get-CimInstance Win32_OperatingSystem -ErrorAction SilentlyContinue
        if ($os) {
            $osCaption = $os.Caption
            foreach ($osKey in $OsKbRequirements.Keys) {
                if ($osCaption -like "*$osKey*") {
                    $requiredKb = $OsKbRequirements[$osKey]
                    Write-Host ''
                    Write-Host "=== OS-Specific KB Check ===" -ForegroundColor Cyan
                    if (-not (Test-KbInstalled -kbNumber $requiredKb)) {
                        Write-Host "KB$requiredKb (for $osCaption) : NOT installed" -ForegroundColor Red
                        if (-not $ScanOnly) {
                            if (-not $Force) {
                                $approve = Read-Host "Install KB$requiredKb? (Y/N)"
                                if ($approve -match '^[Yy]$') {
                                    Install-KbUpdate -kbNumber $requiredKb
                                }
                            }
                            else {
                                Install-KbUpdate -kbNumber $requiredKb
                            }
                        }
                    }
                    else {
                        Write-Host "KB$requiredKb (for $osCaption) : Already installed" -ForegroundColor Green
                    }
                    break
                }
            }
        }
    }
    
    # OS-only guidance
    Write-Host ''
    Write-Host 'OS-Only CVEs (Require Windows Update or Microcode)' -ForegroundColor Yellow
    $osOnlyCves = @(
        'CVE-2022-21127',
        'CVE-2022-21166',
        'CVE-2023-20569',
        'CVE-2023-20588',
        'CVE-2023-28746',
        'CVE-2022-40982'
    )
    
    foreach ($cve in $osOnlyCves) {
        Write-Host " $cve : $(Get-CveRecommendation -CveId $cve)" -ForegroundColor Gray
    }
}

if ($MyInvocation.InvocationName -eq $MyInvocation.MyCommand.Name) {
    Main
}