feat: Add restore defaults and centralize cleanup logic

Introduces a "Restore Defaults" feature in the UI to reset the environment. This action removes generated configuration files, ISOs, downloaded apps, updates, drivers, and FFUs.

The post-build cleanup logic is refactored from the main build script into a new common function. This new function is used by both the standard build process and the new restore defaults feature, promoting code reuse and simplifying maintenance.

FFUDevelopment/BuildFFUVM.ps1

@@ -5893,59 +5893,10 @@ If ($RemoveFFU) {
        
 }
 Set-Progress -Percentage 99 -Message "Finalizing and cleaning up..."
-If ($CleanupCaptureISO) {
+# Delegated post-build cleanup to common module
-    try {
+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
-        If (Test-Path -Path $CaptureISO) {
+
-            WriteLog "Removing $CaptureISO"
+# Remove KBPath for cached vhdx files
-            Remove-Item -Path $CaptureISO -Force
-            WriteLog "Removal complete"
-        }     
-    }
-    catch {
-        Writelog "Removing $CaptureISO failed with error $_"
-        throw $_
-    }
-}
-If ($CleanupDeployISO) {
-    try {
-        If (Test-Path -Path $DeployISO) {
-            WriteLog "Removing $DeployISO"
-            Remove-Item -Path $DeployISO -Force
-            WriteLog "Removal complete"
-        }     
-    }
-    catch {
-        Writelog "Removing $DeployISO failed with error $_"
-        throw $_
-    }
-}
-If ($CleanupAppsISO) {
-    try {
-        If (Test-Path -Path $AppsISO) {
-            WriteLog "Removing $AppsISO"
-            Remove-Item -Path $AppsISO -Force
-            WriteLog "Removal complete"
-        }     
-    }
-    catch {
-        Writelog "Removing $AppsISO failed with error $_"
-        throw $_
-    }
-}
-If ($CleanupDrivers) {
-    try {
-        #Remove files in $Driversfolder, but keep $DriversFolder
-        If (Test-Path -Path $Driversfolder) {
-            WriteLog "Removing files in $Driversfolder"
-            Remove-Item -Path $Driversfolder\* -Force -Recurse
-            WriteLog "Removal complete"
-        }  
-    }
-    catch {
-        Writelog "Removing $Driversfolder\* failed with error $_"
-        throw $_
-    }
-}
 if ($AllowVHDXCaching) {
     try {
         If (Test-Path -Path $KBPath) {

FFUDevelopment/BuildFFUVM_UI.xaml

@@ -816,6 +816,7 @@
         <TextBlock x:Name="txtStatus" Grid.Row="2" Margin="0,5,0,0"/>
         <!-- Buttons (Build Config File / Load Config File / Build FFU) -->
         <StackPanel Grid.Row="3" Orientation="Horizontal" HorizontalAlignment="Right" Margin="0,0,20,20">
+            <Button x:Name="btnRestoreDefaults" Content="Restore Defaults" Width="150" Margin="0,0,10,0" FontSize="14" Padding="10,5"/>
             <Button x:Name="btnBuildConfig" Content="Save Config File" Width="150" Margin="0,0,10,0" FontSize="14" Padding="10,5"/>
             <Button x:Name="btnLoadConfig" Content="Load Config File" Width="150" Margin="0,0,10,0" FontSize="14" Padding="10,5"/>
             <Button x:Name="btnRun" Content="Build FFU" Width="120" FontSize="14" Padding="10,5"/>

FFUDevelopment/FFU.Common/FFU.Common.Cleanup.psm1

@@ -0,0 +1,109 @@
+# Provides shared cleanup functionality for both UI and build script.
+
+function Invoke-FFUPostBuildCleanup {
+    param(
+        [string]$RootPath,
+        [string]$AppsPath,
+        [string]$DriversPath,
+        [string]$FFUCapturePath,
+        [string]$CaptureISOPath,
+        [string]$DeployISOPath,
+        [string]$AppsISOPath,
+        [bool]$RemoveCaptureISO = $false,
+        [bool]$RemoveDeployISO = $false,
+        [bool]$RemoveAppsISO = $false,
+        [bool]$RemoveDrivers = $false,
+        [bool]$RemoveFFU = $false,
+        [bool]$RemoveApps = $false,
+        [bool]$RemoveUpdates = $false
+    )
+    $originalProgressPreference = $ProgressPreference
+    $ProgressPreference = 'SilentlyContinue'
+    try {
+        WriteLog "CommonCleanup: Starting cleanup (CaptureISO=$RemoveCaptureISO DeployISO=$RemoveDeployISO AppsISO=$RemoveAppsISO Drivers=$RemoveDrivers FFU=$RemoveFFU Apps=$RemoveApps Updates=$RemoveUpdates)."
+
+        # Primary ISO paths (new naming/location)
+        if ($RemoveCaptureISO -and -not [string]::IsNullOrWhiteSpace($CaptureISOPath) -and (Test-Path -LiteralPath $CaptureISOPath)) {
+            WriteLog "CommonCleanup: Removing $CaptureISOPath"
+            try { Remove-Item -LiteralPath $CaptureISOPath -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $CaptureISOPath : $($_.Exception.Message)" }
+        }
+        if ($RemoveDeployISO -and -not [string]::IsNullOrWhiteSpace($DeployISOPath) -and (Test-Path -LiteralPath $DeployISOPath)) {
+            WriteLog "CommonCleanup: Removing $DeployISOPath"
+            try { Remove-Item -LiteralPath $DeployISOPath -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $DeployISOPath : $($_.Exception.Message)" }
+        }
+        if ($RemoveAppsISO -and -not [string]::IsNullOrWhiteSpace($AppsISOPath) -and (Test-Path -LiteralPath $AppsISOPath)) {
+            WriteLog "CommonCleanup: Removing $AppsISOPath"
+            try { Remove-Item -LiteralPath $AppsISOPath -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $AppsISOPath : $($_.Exception.Message)" }
+        }
+
+        # Legacy / root-level WinPE ISOs (pattern-based)
+        if ($RemoveCaptureISO) {
+            Get-ChildItem -LiteralPath $RootPath -Filter 'WinPE_FFU_Capture*.iso' -ErrorAction SilentlyContinue | ForEach-Object {
+                try { WriteLog "CommonCleanup: Removing legacy capture ISO $($_.FullName)"; Remove-Item -LiteralPath $_.FullName -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing legacy capture ISO $($_.FullName) : $($_.Exception.Message)" }
+            }
+        }
+        if ($RemoveDeployISO) {
+            Get-ChildItem -LiteralPath $RootPath -Filter 'WinPE_FFU_Deploy*.iso' -ErrorAction SilentlyContinue | ForEach-Object {
+                try { WriteLog "CommonCleanup: Removing legacy deploy ISO $($_.FullName)"; Remove-Item -LiteralPath $_.FullName -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing legacy deploy ISO $($_.FullName) : $($_.Exception.Message)" }
+            }
+        }
+
+        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)" }
+        }
+
+        if ($RemoveFFU -and -not [string]::IsNullOrWhiteSpace($FFUCapturePath) -and (Test-Path -LiteralPath $FFUCapturePath -PathType Container)) {
+            WriteLog "CommonCleanup: Removing FFU files in $FFUCapturePath"
+            Get-ChildItem -LiteralPath $FFUCapturePath -Filter *.ffu -ErrorAction SilentlyContinue | ForEach-Object {
+                try { Remove-Item -LiteralPath $_.FullName -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing FFU $($_.FullName) : $($_.Exception.Message)" }
+            }
+        }
+
+        if ($RemoveApps -and -not [string]::IsNullOrWhiteSpace($AppsPath) -and (Test-Path -LiteralPath $AppsPath -PathType Container)) {
+            $win32 = Join-Path $AppsPath 'Win32'
+            $store = Join-Path $AppsPath 'MSStore'
+            if (Test-Path -LiteralPath $win32) {
+                WriteLog "CommonCleanup: Removing $win32"
+                try { Remove-Item -LiteralPath $win32 -Recurse -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $win32 : $($_.Exception.Message)" }
+            }
+            if (Test-Path -LiteralPath $store) {
+                WriteLog "CommonCleanup: Removing $store"
+                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"
+                $officeSub = Join-Path $office 'Office'
+                if (Test-Path -LiteralPath $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) {
+                    try { Remove-Item -LiteralPath $setupExe -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $setupExe : $($_.Exception.Message)" }
+                }
+            }
+        }
+
+        if ($RemoveUpdates -and -not [string]::IsNullOrWhiteSpace($AppsPath) -and (Test-Path -LiteralPath $AppsPath)) {
+            $updateDirs = @('Defender', 'Edge', 'MSRT', 'OneDrive', '.NET', 'CU', 'Microcode')
+            foreach ($d in $updateDirs) {
+                $target = Join-Path $AppsPath $d
+                if (Test-Path -LiteralPath $target) {
+                    WriteLog "CommonCleanup: Removing update folder $target"
+                    try { Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction Stop } catch { WriteLog "CommonCleanup: Failed removing $target : $($_.Exception.Message)" }
+                }
+            }
+        }
+
+        WriteLog "CommonCleanup: Completed."
+    }
+    catch {
+        WriteLog "CommonCleanup: Fatal cleanup error $($_.Exception.Message)"
+    }
+    finally {
+        $ProgressPreference = $originalProgressPreference
+    }
+}
+
+Export-ModuleMember -Function Invoke-FFUPostBuildCleanup

FFUDevelopment/FFU.Common/FFU.Common.psd1

@@ -68,7 +68,8 @@ Description = 'Common functions shared between FFU Builder UI and the BuildFFUVM
 # Modules to import as nested modules of the module specified in RootModule/ModuleToProcess
 NestedModules = @('FFU.Common.Drivers.psm1', 
                'FFU.Common.Winget.psm1', 
-               'FFU.Common.Parallel.psm1')
+               'FFU.Common.Parallel.psm1',
+               'FFU.Common.Cleanup.psm1')
 
 # Functions to export from this module, for best performance, do not use wildcards and do not delete the entry, use an empty array if there are no functions to export.
 FunctionsToExport = '*'

FFUDevelopment/FFUUI.Core/FFUUI.Core.Config.psm1

@@ -664,6 +664,126 @@ function Invoke-SaveConfiguration {
     }
 }
 
+function Invoke-RestoreDefaults {
+    param(
+        [Parameter(Mandatory = $true)]
+        [psobject]$State
+    )
+    try {
+        $rootPath = $State.FFUDevelopmentPath
+
+        # Normalize potential array values to single strings
+        function Normalize-PathScalar {
+            param([object]$value)
+            if ($null -eq $value) { return $null }
+            if ($value -is [System.Array]) {
+                foreach ($v in $value) {
+                    if (-not [string]::IsNullOrWhiteSpace([string]$v)) {
+                        return [string]$v
+                    }
+                }
+                return $null
+            }
+            return [string]$value
+        }
+
+        $appsPath = Join-Path $rootPath 'Apps'
+        $driversRaw = Normalize-PathScalar -value $State.Controls.txtDriversFolder.Text
+        if ([string]::IsNullOrWhiteSpace($driversRaw)) {
+            $driversPath = Join-Path $rootPath 'Drivers'
+        }
+        else {
+            $driversPath = $driversRaw
+        }
+        $ffuCaptureRaw = Normalize-PathScalar -value $State.Controls.txtFFUCaptureLocation.Text
+        $ffuCapturePath = if ([string]::IsNullOrWhiteSpace($ffuCaptureRaw)) { Join-Path $rootPath 'FFU' } else { $ffuCaptureRaw }
+
+        $captureISOPath = Join-Path $rootPath 'WinPECaptureFFUFiles\WinPE-Capture.iso'
+        $deployISOPath = Join-Path $rootPath 'WinPEDeployFFUFiles\WinPE-Deploy.iso'
+        $appsISOPath = Join-Path $rootPath 'Apps.iso'
+        
+        $msg = "Restore Defaults will:`n`n- Delete generated config and app/driver list JSON files`n- Remove ISO files (Capture, Deploy, Apps) if present`n- Remove Apps/Update/downloaded artifacts`n- Remove driver folder contents (not the folder)`n- Remove FFU files in the capture folder`n`nSample/template files and VM/VHDX cache are NOT removed.`n`nProceed?"
+        $result = [System.Windows.MessageBox]::Show($msg, "Confirm Restore Defaults", "YesNo", "Warning")
+        if ($result -ne [System.Windows.MessageBoxResult]::Yes) {
+            WriteLog "RestoreDefaults: User cancelled."
+            return
+        }
+
+        WriteLog "RestoreDefaults: Starting environment reset."
+        WriteLog "RestoreDefaults: Paths -> Apps=$appsPath Drivers=$driversPath FFUCapture=$ffuCapturePath"
+
+        # Remove JSON artifact files if present
+        $artifactFiles = @(
+            (Join-Path $rootPath 'config\FFUConfig.json'),
+            (Join-Path $appsPath 'AppList.json'),
+            (Join-Path $driversPath 'Drivers.json'),
+            (Join-Path $appsPath 'UserAppList.json')
+        ) | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }
+
+        foreach ($file in $artifactFiles) {
+            if ((-not [string]::IsNullOrWhiteSpace($file)) -and (Test-Path -LiteralPath $file)) {
+                try {
+                    WriteLog "RestoreDefaults: Removing $file"
+                    Remove-Item -LiteralPath $file -Force -ErrorAction Stop
+                }
+                catch {
+                    WriteLog "RestoreDefaults: Failed removing $file : $($_.Exception.Message)"
+                }
+            }
+        }
+
+        # Force all cleanup flags true
+        Invoke-FFUPostBuildCleanup `
+            -RootPath $rootPath `
+            -AppsPath $appsPath `
+            -DriversPath $driversPath `
+            -FFUCapturePath $ffuCapturePath `
+            -CaptureISOPath $captureISOPath `
+            -DeployISOPath $deployISOPath `
+            -AppsISOPath $appsISOPath `
+            -RemoveCaptureISO:$true `
+            -RemoveDeployISO:$true `
+            -RemoveAppsISO:$true `
+            -RemoveDrivers:$true `
+            -RemoveFFU:$true `
+            -RemoveApps:$true `
+            -RemoveUpdates:$true
+
+        # Clear UI lists / state
+        if ($null -ne $State.Data.allDriverModels) { $State.Data.allDriverModels.Clear() }
+        if ($null -ne $State.Controls.lstDriverModels) { $State.Controls.lstDriverModels.Items.Refresh() }
+        if ($null -ne $State.Controls.lstApplications) {
+            try {
+                if ($State.Controls.lstApplications.ItemsSource) { $State.Controls.lstApplications.ItemsSource = $null }
+                $State.Controls.lstApplications.Items.Clear()
+            } catch {}
+        }
+        if ($null -ne $State.Controls.lstWingetResults) {
+            try { 
+                if ($State.Controls.lstWingetResults.ItemsSource) { $State.Controls.lstWingetResults.ItemsSource = $null }
+                $State.Controls.lstWingetResults.Items.Clear() 
+            } catch {}
+        }
+        if ($null -ne $State.Controls.lstAppsScriptVariables) {
+            try {
+                if ($State.Controls.lstAppsScriptVariables.ItemsSource) { $State.Controls.lstAppsScriptVariables.ItemsSource = $null }
+                $State.Controls.lstAppsScriptVariables.Items.Clear()
+            } catch {}
+        }
+
+        $State.Data.lastConfigFilePath = $null
+
+        Initialize-UIDefaults -State $State
+
+        WriteLog "RestoreDefaults: Completed."
+        [System.Windows.MessageBox]::Show("Environment restored to defaults.", "Restore Defaults", "OK", "Information")
+    }
+    catch {
+        WriteLog "RestoreDefaults: Failed with $($_.Exception.Message)"
+        [System.Windows.MessageBox]::Show("Restore Defaults failed:`n$($_.Exception.Message)", "Error", "OK", "Error")
+    }
+}
+
 function Invoke-AutoLoadPreviousEnvironment {
     [CmdletBinding()]
     param(

FFUDevelopment/FFUUI.Core/FFUUI.Core.Handlers.psm1

@@ -216,11 +216,11 @@ function Register-EventHandlers {
             }
             else {
                 $localState.Controls.txtCustomVMSwitchName.Visibility = 'Collapsed'
-                if ($localState.Data.vmSwitchMap.ContainsKey($selectedItem)) {
+                if ($null -ne $selectedItem -and $localState.Data.vmSwitchMap.ContainsKey($selectedItem)) {
                     $localState.Controls.txtVMHostIPAddress.Text = $localState.Data.vmSwitchMap[$selectedItem]
                 }
                 else {
-                    $localState.Controls.txtVMHostIPAddress.Text = '' # Clear IP if not found in map
+                    $localState.Controls.txtVMHostIPAddress.Text = '' # Clear IP if not found or key null
                 }
             }
         })
@@ -950,6 +950,12 @@ function Register-EventHandlers {
             $localState = $window.Tag
             Invoke-LoadConfiguration -State $localState
         })
+    $State.Controls.btnRestoreDefaults.Add_Click({
+            param($eventSource, $routedEventArgs)
+            $window = [System.Windows.Window]::GetWindow($eventSource)
+            $localState = $window.Tag
+            Invoke-RestoreDefaults -State $localState
+        })
     $State.Controls.btnBuildConfig.Add_Click({
             param($eventSource, $routedEventArgs)
             $window = [System.Windows.Window]::GetWindow($eventSource)

FFUDevelopment/FFUUI.Core/FFUUI.Core.Initialize.psm1

@@ -170,6 +170,7 @@ function Initialize-UIControls {
     $State.Controls.btnBrowseDriversJsonPath = $window.FindName('btnBrowseDriversJsonPath')
     $State.Controls.chkUpdateADK = $window.FindName('chkUpdateADK')
     $State.Controls.btnLoadConfig = $window.FindName('btnLoadConfig')
+    $State.Controls.btnRestoreDefaults = $window.FindName('btnRestoreDefaults')
     $State.Controls.btnBuildConfig = $window.FindName('btnBuildConfig')
 
     # Monitor Tab
@@ -198,11 +199,11 @@ function Initialize-VMSwitchData {
     if ($State.Controls.cmbVMSwitchName.Items.Count -gt 1) {
         $State.Controls.cmbVMSwitchName.SelectedIndex = 0
         $firstSwitch = $State.Controls.cmbVMSwitchName.SelectedItem
-        if ($State.Data.vmSwitchMap.ContainsKey($firstSwitch)) {
+        if ($null -ne $firstSwitch -and $State.Data.vmSwitchMap.ContainsKey($firstSwitch)) {
             $State.Controls.txtVMHostIPAddress.Text = $State.Data.vmSwitchMap[$firstSwitch]
         }
         else {
-            $State.Controls.txtVMHostIPAddress.Text = $State.Defaults.generalDefaults.VMHostIPAddress # Use default if IP not found
+            $State.Controls.txtVMHostIPAddress.Text = $State.Defaults.generalDefaults.VMHostIPAddress # Use default if IP not found or key null
         }
         $State.Controls.txtCustomVMSwitchName.Visibility = 'Collapsed'
     }