feat: Defer cleanup of compressed driver source folders

Implements a deferred cleanup mechanism for driver source folders when they are compressed to a WIM and also used for WinPE. When drivers are compressed, the original source folders are now preserved if they are also needed for WinPE driver injection. A marker file is created in these preserved folders. A new cleanup step is added after the WinPE media creation to remove these preserved folders, ensuring they are available when needed but not left behind permanently.
2026-06-14 02:09:35 -06:00 · 2025-09-12 15:08:48 -07:00
parent 50713188bf
commit c30ed923b6
8 changed files with 115 additions and 43 deletions
@@ -2728,7 +2728,8 @@ function Copy-Drivers {
                foreach ($sourceDiskFile in $sourceDiskFiles.Keys) {
                    if (!$sourceDiskFiles[$sourceDiskFile].Contains(",")) {
                        Copy-Item -Path "$infPath\$sourceDiskFile" -Destination $targetPath -Force
-                    } else {
+                    }
+                    else {
                        $subdir = ($sourceDiskFiles[$sourceDiskFile] -split ",")[1]
                        [void](New-Item -Path "$targetPath\$subdir" -ItemType Directory -Force)
                        Copy-Item -Path "$infPath\$subdir\$sourceDiskFile" -Destination "$targetPath\$subdir" -Force
@@ -2740,7 +2741,8 @@ function Copy-Drivers {
                foreach ($sourceDiskFile in $sourceDiskFiles.Keys) {
                    if (!$sourceDiskFiles[$sourceDiskFile].Contains(",")) {
                        Copy-Item -Path "$infPath\$sourceDiskFile" -Destination $targetPath -Force
-                    } else {
+                    }
+                    else {
                        $subdir = ($sourceDiskFiles[$sourceDiskFile] -split ",")[1]
                        [void](New-Item -Path "$targetPath\$subdir" -ItemType Directory -Force)
                        Copy-Item -Path "$infPath\$subdir\$sourceDiskFile" -Destination "$targetPath\$subdir" -Force
@@ -2902,6 +2904,41 @@ function New-PEMedia {
    WriteLog "Cleaning up $WinPEFFUPath"
    Remove-Item -Path "$WinPEFFUPath" -Recurse -Force
    WriteLog 'Cleanup complete'
+    # Deferred cleanup of preserved driver model folders (only after WinPE Deploy media is created)
+    if ($UseDriversAsPEDrivers -and $CompressDownloadedDriversToWim -and $Deploy -and $CopyPEDrivers) {
+        WriteLog "Beginning deferred cleanup of preserved driver model folders (UseDriversAsPEDrivers + compression scenario)."
+        $removedCount = 0
+        $skippedCount = 0
+        if (Test-Path -Path $DriversFolder) {
+            Get-ChildItem -Path $DriversFolder -Directory -ErrorAction SilentlyContinue | ForEach-Object {
+                $makeDir = $_.FullName
+                Get-ChildItem -Path $makeDir -Directory -ErrorAction SilentlyContinue | ForEach-Object {
+                    $modelDir = $_.FullName
+                    $markerFile = Join-Path -Path $modelDir -ChildPath '__PreservedForPEDrivers.txt'
+                    $leaf = Split-Path -Path $modelDir -Leaf
+                    $wimPath = Join-Path -Path $makeDir -ChildPath ($leaf + '.wim')
+                    if ((Test-Path -Path $markerFile -PathType Leaf) -and (Test-Path -Path $wimPath -PathType Leaf)) {
+                        try {
+                            WriteLog "Removing preserved driver folder: $modelDir (WIM located at $wimPath)"
+                            Remove-Item -Path $modelDir -Recurse -Force -ErrorAction Stop
+                            $removedCount++
+                        }
+                        catch {
+                            WriteLog "Warning: Failed to remove preserved folder $modelDir : $($_.Exception.Message)"
+                            $skippedCount++
+                        }
+                    }
+                    else {
+                        $skippedCount++
+                    }
+                }
+            }
+            WriteLog "Deferred driver cleanup complete. Removed: $removedCount; Skipped: $skippedCount"
+        }
+        else {
+            WriteLog "Drivers folder $DriversFolder not found during deferred cleanup."
+        }
+    }
 }

 function Optimize-FFUCaptureDrive {
@@ -4812,6 +4849,7 @@ if ($driversJsonPath -and (Test-Path $driversJsonPath) -and ($InstallDrivers -or
    else {
        WriteLog "Found $($driversToProcess.Count) driver entries to process from $driversJsonPath."

+        $preserveSourceOnCompress = ($UseDriversAsPEDrivers -and $CompressDownloadedDriversToWim)
        $taskArguments = @{
            DriversFolder            = $DriversFolder
            WindowsRelease           = $WindowsRelease
@@ -4820,6 +4858,7 @@ if ($driversJsonPath -and (Test-Path $driversJsonPath) -and ($InstallDrivers -or
            Headers                  = $Headers
            UserAgent                = $UserAgent
            CompressToWim            = $CompressDownloadedDriversToWim
+            PreserveSourceOnCompress = $preserveSourceOnCompress
        }
        
        WriteLog "Starting parallel driver processing using Invoke-ParallelProcessing..."
@@ -22,7 +22,10 @@ function Compress-DriverFolderToWim {
        [string]$WimName, # Optional, defaults to folder name

        [Parameter()]
-        [string]$WimDescription # Optional, defaults to folder name
+        [string]$WimDescription, # Optional, defaults to folder name
+
+        [Parameter()]
+        [bool]$PreserveSource = $false # When $true, do not delete source folder; create marker for deferred cleanup
    )

    WriteLog "Starting compression of folder '$SourceFolderPath' to '$DestinationWimPath'."
@@ -66,6 +69,20 @@ function Compress-DriverFolderToWim {
            WriteLog "Successfully compressed '$SourceFolderPath' to '$DestinationWimPath' using dism.exe."
            
            # Remove the source folder after successful compression
+            if ($PreserveSource) {
+                WriteLog "Preserving source driver folder for deferred WinPE driver harvesting: $SourceFolderPath"
+                try {
+                    $markerFile = Join-Path -Path $SourceFolderPath -ChildPath '__PreservedForPEDrivers.txt'
+                    if (-not (Test-Path -Path $markerFile -PathType Leaf)) {
+                        New-Item -Path $markerFile -ItemType File -Force | Out-Null
+                        WriteLog "Created preservation marker file: $markerFile"
+                    }
+                }
+                catch {
+                    WriteLog "Warning: Failed to create preservation marker in $SourceFolderPath. Error: $($_.Exception.Message)"
+                }
+            }
+            else {
                WriteLog "Removing source driver folder: $SourceFolderPath"
                try {
                    Remove-Item -Path $SourceFolderPath -Recurse -Force -ErrorAction Stop
@@ -75,6 +92,7 @@ function Compress-DriverFolderToWim {
                    WriteLog "Warning: Failed to remove source folder '$SourceFolderPath'. Error: $($_.Exception.Message)"
                    # Do not fail the whole operation, just log a warning.
                }
+            }

            return $true # Indicate success
        }
@@ -209,7 +209,8 @@ function Invoke-ParallelProcessing {
                                    -Headers $localJobArgs['Headers'] `
                                    -UserAgent $localJobArgs['UserAgent'] `
                                    -ProgressQueue $localProgressQueue `
-                                    -CompressToWim $localJobArgs['CompressToWim']
+                                    -CompressToWim $localJobArgs['CompressToWim'] `
+                                    -PreserveSourceOnCompress $localJobArgs['PreserveSourceOnCompress']
                            }
                            'Dell' {
                                $taskResult = Save-DellDriversTask -DriverItemData $currentItem `
@@ -217,7 +218,8 @@ function Invoke-ParallelProcessing {
                                    -WindowsArch $localJobArgs['WindowsArch'] `
                                    -WindowsRelease $localJobArgs['WindowsRelease'] `
                                    -ProgressQueue $localProgressQueue `
-                                    -CompressToWim $localJobArgs['CompressToWim']
+                                    -CompressToWim $localJobArgs['CompressToWim'] `
+                                    -PreserveSourceOnCompress $localJobArgs['PreserveSourceOnCompress']
                            }
                            'HP' {
                                $taskResult = Save-HPDriversTask -DriverItemData $currentItem `
@@ -226,7 +228,8 @@ function Invoke-ParallelProcessing {
                                    -WindowsRelease $localJobArgs['WindowsRelease'] `
                                    -WindowsVersion $localJobArgs['WindowsVersion'] `
                                    -ProgressQueue $localProgressQueue `
-                                    -CompressToWim $localJobArgs['CompressToWim']
+                                    -CompressToWim $localJobArgs['CompressToWim'] `
+                                    -PreserveSourceOnCompress $localJobArgs['PreserveSourceOnCompress']
                            }
                            'Lenovo' {
                                $taskResult = Save-LenovoDriversTask -DriverItemData $currentItem `
@@ -235,7 +238,8 @@ function Invoke-ParallelProcessing {
                                    -Headers $localJobArgs['Headers'] `
                                    -UserAgent $localJobArgs['UserAgent'] `
                                    -ProgressQueue $localProgressQueue `
-                                    -CompressToWim $localJobArgs['CompressToWim']
+                                    -CompressToWim $localJobArgs['CompressToWim'] `
+                                    -PreserveSourceOnCompress $localJobArgs['PreserveSourceOnCompress']
                            }
                            default {
                                $unsupportedMakeMessage = "Error: Unsupported Make '$make' for driver download."
@@ -167,10 +167,12 @@ function Save-DellDriversTask {
        [string]$WindowsArch,
        [Parameter(Mandatory = $true)]
        [int]$WindowsRelease,
-        [Parameter()] # Made optional
+        [Parameter()]
        [System.Collections.Concurrent.ConcurrentQueue[hashtable]]$ProgressQueue = $null, # Default to null
        [Parameter()]
-        [bool]$CompressToWim = $false # New parameter for compression
+        [bool]$CompressToWim = $false, # New parameter for compression
+        [Parameter()]
+        [bool]$PreserveSourceOnCompress = $false
    )
        
    $modelName = $DriverItemData.Model
@@ -201,7 +203,7 @@ function Save-DellDriversTask {
                WriteLog "Attempting compression of existing folder '$sourceFolderPath' to '$wimFilePath'."
                if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status "Compressing existing..." }
                try {
-                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $modelName -WimDescription "Drivers for $modelName" -ErrorAction Stop
+                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $modelName -WimDescription "Drivers for $modelName" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                    $existingDriver.Status = "Already downloaded & Compressed"
                    $existingDriver.DriverPath = Join-Path -Path $make -ChildPath "$($modelName).wim"
                    $existingDriver.Success = $true
@@ -664,7 +666,7 @@ function Save-DellDriversTask {
            $driverRelativePath = Join-Path -Path $make -ChildPath $wimFileName # Update relative path to the WIM file
            WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
            try {
-                $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -ErrorAction Stop
+                $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                if ($compressResult) {
                    WriteLog "Compression successful for '$modelName'."
                    $status = "Completed & Compressed"
@@ -118,7 +118,9 @@ function Save-HPDriversTask {
        [Parameter()] # Made optional
        [System.Collections.Concurrent.ConcurrentQueue[hashtable]]$ProgressQueue = $null, # Default to null
        [Parameter()]
-        [bool]$CompressToWim = $false # New parameter for compression
+        [bool]$CompressToWim = $false, # New parameter for compression
+        [Parameter()]
+        [bool]$PreserveSourceOnCompress = $false
    )
            
    $modelName = $DriverItemData.Model
@@ -150,7 +152,7 @@ function Save-HPDriversTask {
                WriteLog "Attempting compression of existing folder '$sourceFolderPath' to '$wimFilePath'."
                if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status "Compressing existing..." }
                try {
-                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -ErrorAction Stop
+                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                    $existingDriver.Status = "Already downloaded & Compressed"
                    $existingDriver.DriverPath = Join-Path -Path $make -ChildPath "$($sanitizedModelName).wim"
                    $existingDriver.Success = $true
@@ -362,7 +364,7 @@ function Save-HPDriversTask {
            $wimFilePath = Join-Path -Path $hpDriversBaseFolder -ChildPath "$($identifier).wim"
            WriteLog "Compressing '$modelSpecificFolder' to '$wimFilePath'..."
            try {
-                Compress-DriverFolderToWim -SourceFolderPath $modelSpecificFolder -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -ErrorAction Stop
+                Compress-DriverFolderToWim -SourceFolderPath $modelSpecificFolder -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                WriteLog "Compression successful for '$identifier'."
                $finalStatus = "Completed & Compressed"
                $driverRelativePath = Join-Path -Path $make -ChildPath "$($identifier).wim" # Update relative path to the WIM
@@ -94,10 +94,12 @@ function Save-LenovoDriversTask {
        [hashtable]$Headers,
        [Parameter(Mandatory = $true)]
        [string]$UserAgent,
-        [Parameter()] # Made optional
+        [Parameter()]
        [System.Collections.Concurrent.ConcurrentQueue[hashtable]]$ProgressQueue = $null,
        [Parameter()]
-        [bool]$CompressToWim = $false
+        [bool]$CompressToWim = $false,
+        [Parameter()]
+        [bool]$PreserveSourceOnCompress = $false
    )
            
    # The Model property from the UI already contains the combined "ProductName (MachineType)" string
@@ -133,7 +135,7 @@ function Save-LenovoDriversTask {
                WriteLog "Attempting compression of existing folder '$sourceFolderPath' to '$wimFilePath'."
                if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status "Compressing existing..." }
                try {
-                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -ErrorAction Stop
+                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                    $existingDriver.Status = "Already downloaded & Compressed"
                    $existingDriver.DriverPath = Join-Path -Path $make -ChildPath "$($sanitizedIdentifier).wim"
                    $existingDriver.Success = $true
@@ -424,7 +426,7 @@ function Save-LenovoDriversTask {
            $driverRelativePath = Join-Path -Path $make -ChildPath $wimFileName # Update relative path to the WIM file
            WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
            try {
-                $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $identifier -WimDescription $identifier -ErrorAction Stop
+                $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $identifier -WimDescription $identifier -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                if ($compressResult) {
                    WriteLog "Compression successful for '$identifier'."
                    $status = "Completed & Compressed"
@@ -94,7 +94,9 @@ function Save-MicrosoftDriversTask {
        [Parameter()] # Made optional
        [System.Collections.Concurrent.ConcurrentQueue[hashtable]]$ProgressQueue = $null, # Default to null
        [Parameter()]
-        [bool]$CompressToWim = $false # New parameter for compression
+        [bool]$CompressToWim = $false, # New parameter for compression
+        [Parameter()]
+        [bool]$PreserveSourceOnCompress = $false
        # REMOVED: UI-related parameters
    )
        
@@ -125,7 +127,7 @@ function Save-MicrosoftDriversTask {
                WriteLog "Attempting compression of existing folder '$sourceFolderPath' to '$wimFilePath'."
                if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status "Compressing existing..." }
                try {
-                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $modelName -WimDescription "Drivers for $modelName" -ErrorAction Stop
+                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $modelName -WimDescription "Drivers for $modelName" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                    $existingDriver.Status = "Already downloaded & Compressed"
                    $existingDriver.DriverPath = Join-Path -Path $make -ChildPath "$($modelName).wim"
                    $existingDriver.Success = $true
@@ -377,7 +379,7 @@ function Save-MicrosoftDriversTask {
                    WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
                    try {
                        # Use the function from the imported common module
-                        $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -ErrorAction Stop
+                        $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                        if ($compressResult) {
                            WriteLog "Compression successful for '$modelName'."
                            $status = "Completed & Compressed"
@@ -571,6 +571,8 @@ function Invoke-DownloadSelectedDrivers {
    $localHeaders = $coreStaticVars.Headers
    $localUserAgent = $coreStaticVars.UserAgent
    $compressDrivers = $State.Controls.chkCompressDriversToWIM.IsChecked
+    # Determine if we must preserve source folders (used later for PE driver harvesting)
+    $preserveSource = ($State.Controls.chkUseDriversAsPEDrivers.IsChecked -and $State.Controls.chkCompressDriversToWIM.IsChecked)

    $State.Controls.txtStatus.Text = "Processing all selected drivers..."
    WriteLog "Processing all selected drivers: $($selectedDrivers.Model -join ', ')"
@@ -620,7 +622,7 @@ function Invoke-DownloadSelectedDrivers {
            return
        }
    }
-
+    $preserveSource = ($State.Controls.chkUseDriversAsPEDrivers.IsChecked -and $State.Controls.chkCompressDriversToWIM.IsChecked)
    $taskArguments = @{
        DriversFolder            = $localDriversFolder
        WindowsRelease           = $localWindowsRelease
@@ -629,6 +631,7 @@ function Invoke-DownloadSelectedDrivers {
        Headers                  = $localHeaders
        UserAgent                = $localUserAgent
        CompressToWim            = $compressDrivers
+        PreserveSourceOnCompress = $preserveSource
    }

    $parallelResults = Invoke-ParallelProcessing -ItemsToProcess $selectedDrivers `