KB5034957: Updating the WinRE partition on deployed devices to address security vulnerabilities in CVE-2024-20666

#################################################################################

#

# Copyright (c) Microsoft Corporation.

# Licensed under the MIT License.

#

# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

# SOFTWARE.

#

#################################################################################

Param (
    [Parameter(Mandatory=$true, HelpMessage="Path to target package")][string]$PackagePath,
    [Parameter(Mandatory=$false, HelpMessage="Directory to dism get-packages logs")][string]$LogDir
)

Set-StrictMode -Version Latest;

# Function to get WinRE path
function GetWinREPath {
    $WinRELocation = (reagentc /info | Select-String "Windows RE location")
    if ($WinRELocation) {
        return $WinRELocation.ToString().Split(':')[-1].Trim()
    } else {
        Write-Host "Failed to find WinRE path" -ForegroundColor Red
        exit 1
    }
}

# Function to get WinRE version
function GetWinREVersion {
    $filePath = "C:\mnt\Windows\System32\winpeshl.exe"
    $WinREVersion = (Get-Item $filePath).VersionInfo.FileVersionRaw.Revision
    return [int]$WinREVersion
}

# Function to get package version
function GetPackageVersion {
    $PackageInfo = dism /Online /Get-PackageInfo /PackagePath:"$PackagePath" | Select-String "Version :"
    if ($PackageInfo) {
        $VersionString = ($PackageInfo -split ':')[-1].Trim()
        if ($VersionString -match "\d+\.\d+\.\d+\.(\d+)") {
            return [int]$matches[1]  # Extract the last part (build number)
        } else {
            Write-Host "Failed to parse package version" -ForegroundColor Red
            exit 1
        }
    } else {
        Write-Host "Failed to retrieve package version" -ForegroundColor Red
        exit 1
    }
}

# Function to ensure log directory access
function EnsureLogDirAccess {
    param([string]$LogDir)
    if (Test-Path $LogDir) {
        try {
            $TestFile = "$LogDir\test_write_access.txt"
            Set-Content -Path $TestFile -Value "Test" -ErrorAction Stop
            Remove-Item -Path $TestFile -Force -ErrorAction Stop
            Write-Host "Log directory access verified." -ForegroundColor Green
        } catch {
            Write-Host "Insufficient permissions for log directory: $LogDir. Attempting to gain access..." -ForegroundColor Yellow
            try {
                $acl = Get-Acl $LogDir
                $rule = New-Object System.Security.AccessControl.FileSystemAccessRule("Everyone", "FullControl", "ContainerInherit,ObjectInherit", "None", "Allow")
                $acl.SetAccessRule($rule)
                Set-Acl -Path $LogDir -AclObject $acl
                Write-Host "Successfully gained access to log directory: $LogDir" -ForegroundColor Green
            } catch {
                Write-Host "Failed to gain access to log directory: $LogDir" -ForegroundColor Red
                exit 1
            }
        }
    } else {
        Write-Host "Log directory does not exist: $LogDir. Creating directory..." -ForegroundColor Yellow
        try {
            New-Item -ItemType Directory -Path $LogDir -Force | Out-Null
            Write-Host "Log directory created: $LogDir" -ForegroundColor Green
        } catch {
            Write-Host "Failed to create log directory: $LogDir" -ForegroundColor Red
            exit 1
        }
    }
}

# If file is other than .cab then exit
if ($PackagePath -notmatch "\.cab$") {
    Write-Host "Invalid package format. Only .cab files are supported." -ForegroundColor Red
    exit 1
}

# Main Execution
$WinREPath = GetWinREPath
$PackageVersion = GetPackageVersion

Write-Host "WinRE Path: $WinREPath" -ForegroundColor Cyan
Write-Host "Package Version: $PackageVersion" -ForegroundColor Cyan

if ($LogDir) {
    EnsureLogDirAccess -LogDir $LogDir
}

# Make dir C:\mnt if not exists
$TempDir = "C:\mnt"

# Get the read write permission for this directory
if (-not (Test-Path $TempDir)) {
    New-Item -ItemType Directory -Path $TempDir -Force | Out-Null
}

# Mount WinRE image
dism /Mount-Image /ImageFile:"$WinREPath\winre.wim" /Index:1 /MountDir:"$TempDir"

$WinREVersion = GetWinREVersion
Write-Host "WinRE Version: $WinREVersion" -ForegroundColor Cyan

if ($PackageVersion -gt $WinREVersion)
{
    Write-Host "Applying patch..." -ForegroundColor Yellow
    if ($LogDir) {
        dism /Image:"$TempDir" /Add-Package /PackagePath:"$PackagePath" /LogPath:"$LogDir\PatchWinRE.log"
    } else {
        dism /Image:"$TempDir" /Add-Package /PackagePath:"$PackagePath"
    }
    Write-Host "Committing changes..." -ForegroundColor Yellow
    dism /Unmount-Image /MountDir:"$TempDir" /Commit
    Write-Host "Patch applied and committed successfully!" -ForegroundColor Green
}
else{
    Write-Host "Already have a greater or equal version, no update needed." -ForegroundColor Green
    dism /Unmount-Image /MountDir:"$TempDir" /Discard
}

# Disable WinRE and re-enable it to let new WinRE be trusted by BitLocker
Write-Host "Disable WinRE"
reagentc /disable
Write-Host "Re-enable WinRE"
reagentc /enable
reagentc /info

# Cleanup
Remove-Item -Path $TempDir -Force -Recurse

## Usage
# .\Patch_WinRE_Generic.ps1 -PackagePath <path_to_cab_file>
# .\Patch_WinRE_Generic.ps1 -PackagePath <path_to_cab_file> -LogDir <path_to_custom_log_folder>