From 15fdf77ce49f50170a3499e89834213b0acd91b4 Mon Sep 17 00:00:00 2001
From: rbalsleyMSFT <53497092+rbalsleyMSFT@users.noreply.github.com>
Date: Mon, 15 Dec 2025 16:20:01 -0800
Subject: [PATCH] Refactors cleanup logic into shared module

Consolidates duplicated cleanup code by moving logic into a shared function, eliminating redundant implementations across multiple locations.

Removes standalone cleanup functions (Remove-FFU, Remove-Apps, Remove-Updates) and replaces scattered cleanup calls with a single invocation of Invoke-FFUPostBuildCleanup.

Enhances driver cleanup to preserve configuration files (Drivers.json and DriverMapping.json) while removing other contents, preventing loss of driver mapping data.

Improves maintainability by centralizing cleanup operations and reducing code duplication, making future updates easier to implement consistently.
---
 FFUDevelopment/BuildFFUVM.ps1                 | 150 +-----------------
 .../FFU.Common/FFU.Common.Cleanup.psm1        |  17 +-
 2 files changed, 18 insertions(+), 149 deletions(-)

diff --git a/FFUDevelopment/BuildFFUVM.ps1 b/FFUDevelopment/BuildFFUVM.ps1
index e0a253d..ef79cf7 100644
--- a/FFUDevelopment/BuildFFUVM.ps1
+++ b/FFUDevelopment/BuildFFUVM.ps1
@@ -3809,15 +3809,10 @@ function Get-FFUEnvironment {
         Remove-FFUUserShare
         WriteLog 'Removal complete'
     }
-    if ($RemoveApps) {
-        WriteLog "Removing Apps in $AppsPath"
-        Remove-Apps
-    }
-    #Remove updates
-    if ($RemoveUpdates) {
-        WriteLog "Removing updates"
-        Remove-Updates
-    }
+
+    #Run shared cleanup to avoid duplicated logic
+    Invoke-FFUPostBuildCleanup -RootPath $FFUDevelopmentPath -AppsPath $AppsPath -DriversPath $DriversFolder -FFUCapturePath $FFUCaptureLocation -CaptureISOPath $CaptureISO -DeployISOPath $DeployISO -AppsISOPath $AppsISO -RemoveCaptureISO:$CleanupCaptureISO -RemoveDeployISO:$CleanupDeployISO -RemoveAppsISO:$CleanupAppsISO -RemoveDrivers:$CleanupDrivers -RemoveFFU:$RemoveFFU -RemoveApps:$RemoveApps -RemoveUpdates:$RemoveUpdates
+
     #Clean up $KBPath
     If (Test-Path -Path $KBPath) {
         WriteLog "Removing $KBPath"
@@ -3841,12 +3836,6 @@ function Get-FFUEnvironment {
     Remove-Item -Path "$FFUDevelopmentPath\dirty.txt" -Force
     WriteLog "Cleanup complete"
 }
-function Remove-FFU {
-    #Remove all FFU files in the FFUCaptureLocation
-    WriteLog "Removing all FFU files in $FFUCaptureLocation"
-    Remove-Item -Path $FFUCaptureLocation\*.ffu -Force
-    WriteLog "Removal complete"
-}
 Function Remove-DisabledArtifacts {
     # Remove Office artifacts if Install Office is disabled
     if (-not $InstallOffice) {
@@ -3932,101 +3921,6 @@ Function Remove-DisabledArtifacts {
     }
 }
 
-Function Remove-Updates {
-    if ($UpdateLatestDefender) {
-        #Clean up $installDefenderPath
-        WriteLog "Removing $installDefenderPath"
-        If (Test-Path -Path $installDefenderPath) {
-            Remove-Item -Path $installDefenderPath -Force -ErrorAction SilentlyContinue
-            WriteLog 'Removal complete'
-        }
-        #Clean up $DefenderPath
-        If (Test-Path -Path $DefenderPath) {
-            WriteLog "Removing $DefenderPath"
-            Remove-Item -Path $DefenderPath -Recurse -Force -ErrorAction SilentlyContinue
-            WriteLog 'Removal complete'
-        }
-    }
-    if ($UpdateLatestMSRT) {
-        # Clean up Update-MSRT.ps1
-        WriteLog "Removing $installMSRTPath"
-        If (Test-Path -Path $installMSRTPath) {
-            Remove-Item -Path $installMSRTPath -Force -ErrorAction SilentlyContinue
-            WriteLog 'Removal complete'
-        }
-        #Clean up $MSRTPath
-        If (Test-Path -Path $MSRTPath) {
-            WriteLog "Removing $MSRTPath"
-            Remove-Item -Path $MSRTPath -Recurse -Force -ErrorAction SilentlyContinue
-            WriteLog 'Removal complete'
-        }
-    }
-    if ($UpdateOneDrive) {
-        # Clean up Update-OneDrive.ps1
-        WriteLog "Removing $installODPath"
-        If (Test-Path -Path $installODPath) {
-            Remove-Item -Path $installODPath -Force -ErrorAction SilentlyContinue
-            WriteLog 'Removal complete'
-        }
-        #Clean up $OneDrivePath
-        If (Test-Path -Path $OneDrivePath) {
-            WriteLog "Removing $OneDrivePath"
-            Remove-Item -Path $OneDrivePath -Recurse -Force -ErrorAction SilentlyContinue
-            WriteLog 'Removal complete'
-        }
-    }
-    if ($UpdateEdge) {
-        # Clean up Update-Edge.ps1
-        WriteLog "Removing $installEdgePath"
-        If (Test-Path -Path $installEdgePath) {
-            Remove-Item -Path $installEdgePath -Force -ErrorAction SilentlyContinue
-            WriteLog 'Removal complete'
-        }
-        #Clean up $EdgePath
-        If (Test-Path -Path $EdgePath) {
-            WriteLog "Removing $EdgePath"
-            Remove-Item -Path $EdgePath -Recurse -Force -ErrorAction SilentlyContinue
-            WriteLog 'Removal complete'
-        }
-    }
-
-}
-function Remove-Apps {
-   
-    # Check if the file exists before attempting to clear it
-    if (Test-Path -Path $wingetWin32jsonFile) {
-        WriteLog "Removing $wingetWin32jsonFile"
-        Remove-Item -Path $wingetWin32jsonFile -Force -ErrorAction SilentlyContinue
-        WriteLog 'Removal complete'
-    }
-    # Clean up Win32 and MSStore folders
-    if (Test-Path -Path "$AppsPath\Win32" -PathType Container) {
-        WriteLog "Cleaning up Win32 folder"
-        Remove-Item -Path "$AppsPath\Win32" -Recurse -Force
-    }
-    if (Test-Path -Path "$AppsPath\MSStore" -PathType Container) {
-        WriteLog "Cleaning up MSStore folder"
-        Remove-Item -Path "$AppsPath\MSStore" -Recurse -Force
-    }
-
-    #Remove the Office Download and ODT
-    if ($InstallOffice) {
-        $ODTPath = "$AppsPath\Office"
-        $OfficeDownloadPath = "$ODTPath\Office"
-        WriteLog 'Removing Office and ODT download'
-        Remove-Item -Path $OfficeDownloadPath -Recurse -Force
-        Remove-Item -Path "$ODTPath\setup.exe"
-        Remove-Item -Path "$orchestrationPath\Install-Office.ps1"
-        WriteLog 'Removal complete'
-    }
-
-    #Remove AppsISO
-    if ($CleanupAppsISO) {
-        WriteLog "Removing $AppsISO"
-        Remove-Item -Path $AppsISO -Force -ErrorAction SilentlyContinue
-        WriteLog 'Removal complete'
-    }
-}
 function Export-ConfigFile {
     [CmdletBinding()]
     param (
@@ -6220,30 +6114,6 @@ If ($InstallApps) {
         Remove-FFUVM -VMName $VMName
         throw $_
     }
-    #Clean up Apps
-    if ($RemoveApps) {
-        try {
-            WriteLog "Cleaning up $AppsPath"
-            Remove-Apps
-        }
-        catch {
-            Write-Host 'Cleaning up Apps failed'
-            Writelog "Cleaning up Apps failed with error $_"
-            throw $_
-        }
-    }
-    #Clean up Updates
-    if ($RemoveUpdates) {
-        try {
-            WriteLog "Cleaning up downloaded update files"
-            Remove-Updates
-        }
-        catch {
-            Write-Host 'Cleaning up downloaded update files failed'
-            Writelog "Cleaning up downloaded update files failed with error $_"
-            throw $_
-        }
-    }
 }
 #Clean up VM or VHDX
 try {
@@ -6315,17 +6185,7 @@ If ($BuildUSBDrive) {
         throw $_
     }
 }
-If ($RemoveFFU) {
-    try {
-        Remove-FFU
-    }
-    catch {
-        Write-Host 'Removing FFU files failed'
-        Writelog "Removing FFU files failed with error $_"
-        throw $_
-    }
-       
-}
+
 Set-Progress -Percentage 99 -Message "Finalizing and cleaning up..."
 # Delegated post-build cleanup to common module
 Invoke-FFUPostBuildCleanup -RootPath $FFUDevelopmentPath -AppsPath $AppsPath -DriversPath $Driversfolder -FFUCapturePath $FFUCaptureLocation -CaptureISOPath $CaptureISO -DeployISOPath $DeployISO -AppsISOPath $AppsISO -RemoveCaptureISO:$CleanupCaptureISO -RemoveDeployISO:$CleanupDeployISO -RemoveAppsISO:$CleanupAppsISO -RemoveDrivers:$CleanupDrivers -RemoveFFU:$RemoveFFU -RemoveApps:$RemoveApps -RemoveUpdates:$RemoveUpdates
diff --git a/FFUDevelopment/FFU.Common/FFU.Common.Cleanup.psm1 b/FFUDevelopment/FFU.Common/FFU.Common.Cleanup.psm1
index bb8334f..bb4fe35 100644
--- a/FFUDevelopment/FFU.Common/FFU.Common.Cleanup.psm1
+++ b/FFUDevelopment/FFU.Common/FFU.Common.Cleanup.psm1
@@ -49,8 +49,15 @@ function Invoke-FFUPostBuildCleanup {
         }
 
         if ($RemoveDrivers -and -not [string]::IsNullOrWhiteSpace($DriversPath) -and (Test-Path -LiteralPath $DriversPath -PathType Container)) {
-            WriteLog "CommonCleanup: Removing contents of $DriversPath"
-            try { Get-ChildItem -LiteralPath $DriversPath -Force -ErrorAction SilentlyContinue | Remove-Item -Force -Recurse -ErrorAction SilentlyContinue } catch { WriteLog "CommonCleanup: Driver content cleanup issue: $($_.Exception.Message)" }
+            WriteLog "CommonCleanup: Removing contents of $DriversPath (preserving Drivers.json and DriverMapping.json)"
+            try {
+                # Preserve drivers json files
+                $driverItems = Get-ChildItem -LiteralPath $DriversPath -Force -ErrorAction SilentlyContinue | Where-Object { @('Drivers.json', 'DriverMapping.json') -notcontains $_.Name }
+                if ($driverItems) {
+                    $driverItems | Remove-Item -Force -Recurse -ErrorAction SilentlyContinue
+                }
+            }
+            catch { WriteLog "CommonCleanup: Driver content cleanup issue: $($_.Exception.Message)" }
         }
 
         if ($RemoveFFU -and -not [string]::IsNullOrWhiteSpace($FFUCapturePath) -and (Test-Path -LiteralPath $FFUCapturePath -PathType Container)) {
@@ -72,14 +79,16 @@ function Invoke-FFUPostBuildCleanup {
                 try { Remove-Item -LiteralPath $store -Recurse -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $store : $($_.Exception.Message)" }
             }
             $office = Join-Path $AppsPath 'Office'
-            if (Test-Path -LiteralPath $office) {
-                WriteLog "CommonCleanup: Cleaning Office artifacts"
+            if ((Test-Path -LiteralPath $office) -and $InstallOffice) {
+                WriteLog "CommonCleanup: Checking for Office artifacts in $office"
                 $officeSub = Join-Path $office 'Office'
                 if (Test-Path -LiteralPath $officeSub) {
+                    WriteLog "CommonCleanup: Removing $officeSub"
                     try { Remove-Item -LiteralPath $officeSub -Recurse -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $officeSub : $($_.Exception.Message)" }
                 }
                 $setupExe = Join-Path $office 'setup.exe'
                 if (Test-Path -LiteralPath $setupExe) {
+                    WriteLog "CommonCleanup: Removing $setupExe"
                     try { Remove-Item -LiteralPath $setupExe -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $setupExe : $($_.Exception.Message)" }
                 }
             }