Win11 Creator USB and Log Fixes (#4139)

* initial usb fixes

* fix full button width

* Cleanup and Verbose output for copy

* expand ui and fix clean and reset

* Add minimal driver injection

* initial driver support

* fix verbage

* fix syntax error

* create log file on iso generation

* inject to boot.wim for install

* fix single driver install issues

* fix first run probs

* cleanup injection

* improve clean up

* Fix OSCDIMG output and cleanup comments

* Fix Scrollviewer in Status Log

* large drive support and change to Exfat

* Fix BOOT for older UEFI and Add error checks for small usb drives

* fix single usb drive error

functions/private/Invoke-WinUtilISO.ps1

@@ -1,18 +1,12 @@
 function Write-Win11ISOLog {
-    <#
-    .SYNOPSIS
-        Appends a timestamped message to the Win11ISO status log TextBox.
-    .PARAMETER Message
-        The message to append.
-    #>
     param([string]$Message)
-    $timestamp = (Get-Date).ToString("HH:mm:ss")
+    $ts = (Get-Date).ToString("HH:mm:ss")
     $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
         $current = $sync["WPFWin11ISOStatusLog"].Text
         if ($current -eq "Ready. Please select a Windows 11 ISO to begin.") {
-            $sync["WPFWin11ISOStatusLog"].Text = "[$timestamp] $Message"
+            $sync["WPFWin11ISOStatusLog"].Text = "[$ts] $Message"
         } else {
-            $sync["WPFWin11ISOStatusLog"].Text += "`n[$timestamp] $Message"
+            $sync["WPFWin11ISOStatusLog"].Text += "`n[$ts] $Message"
         }
         $sync["WPFWin11ISOStatusLog"].CaretIndex = $sync["WPFWin11ISOStatusLog"].Text.Length
         $sync["WPFWin11ISOStatusLog"].ScrollToEnd()
@@ -20,11 +14,6 @@ function Write-Win11ISOLog {
 }
 
 function Invoke-WinUtilISOBrowse {
-    <#
-    .SYNOPSIS
-        Opens an OpenFileDialog so the user can choose a Windows 11 ISO file.
-        Populates WPFWin11ISOPath and reveals the Mount & Verify section (Step 2).
-    #>
     Add-Type -AssemblyName System.Windows.Forms
 
     $dlg = [System.Windows.Forms.OpenFileDialog]::new()
@@ -35,18 +24,12 @@ function Invoke-WinUtilISOBrowse {
     if ($dlg.ShowDialog() -ne [System.Windows.Forms.DialogResult]::OK) { return }
 
     $isoPath    = $dlg.FileName
-
-    # ── Basic size sanity-check (a Win11 ISO is typically > 4 GB) ──
     $fileSizeGB = [math]::Round((Get-Item $isoPath).Length / 1GB, 2)
 
     $sync["WPFWin11ISOPath"].Text           = $isoPath
     $sync["WPFWin11ISOFileInfo"].Text       = "File size: $fileSizeGB GB"
     $sync["WPFWin11ISOFileInfo"].Visibility = "Visible"
-
-    # Reveal Step 2
     $sync["WPFWin11ISOMountSection"].Visibility       = "Visible"
-
-    # Collapse all later steps whenever a new ISO is chosen
     $sync["WPFWin11ISOVerifyResultPanel"].Visibility  = "Collapsed"
     $sync["WPFWin11ISOModifySection"].Visibility      = "Collapsed"
     $sync["WPFWin11ISOOutputSection"].Visibility      = "Collapsed"
@@ -55,17 +38,10 @@ function Invoke-WinUtilISOBrowse {
 }
 
 function Invoke-WinUtilISOMountAndVerify {
-    <#
-    .SYNOPSIS
-        Mounts the selected ISO, verifies it is a valid Windows 11 image,
-        and populates the edition list.  Reveals Step 3 on success.
-    #>
     $isoPath = $sync["WPFWin11ISOPath"].Text
 
     if ([string]::IsNullOrWhiteSpace($isoPath) -or $isoPath -eq "No ISO selected...") {
-        [System.Windows.MessageBox]::Show(
-            "Please select an ISO file first.",
-            "No ISO Selected", "OK", "Warning")
+        [System.Windows.MessageBox]::Show("Please select an ISO file first.", "No ISO Selected", "OK", "Warning")
         return
     }
 
@@ -73,14 +49,12 @@ function Invoke-WinUtilISOMountAndVerify {
     Set-WinUtilProgressBar -Label "Mounting ISO..." -Percent 10
 
     try {
-        # Mount the ISO
         $diskImage   = Mount-DiskImage -ImagePath $isoPath -PassThru -ErrorAction Stop
         $driveLetter = ($diskImage | Get-Volume).DriveLetter + ":"
         Write-Win11ISOLog "Mounted at drive $driveLetter"
 
         Set-WinUtilProgressBar -Label "Verifying ISO contents..." -Percent 30
 
-        # ── Verify install.wim / install.esd presence ──
         $wimPath = Join-Path $driveLetter "sources\install.wim"
         $esdPath = Join-Path $driveLetter "sources\install.esd"
 
@@ -96,14 +70,10 @@ function Invoke-WinUtilISOMountAndVerify {
 
         $activeWim = if (Test-Path $wimPath) { $wimPath } else { $esdPath }
 
-        # ── Read edition / architecture info ──
         Set-WinUtilProgressBar -Label "Reading image metadata..." -Percent 55
-
         $imageInfo = Get-WindowsImage -ImagePath $activeWim | Select-Object ImageIndex, ImageName
 
-        # ── Verify at least one Win11 edition is present ──
-        $isWin11 = $imageInfo | Where-Object { $_.ImageName -match "Windows 11" }
-        if (-not $isWin11) {
+        if (-not ($imageInfo | Where-Object { $_.ImageName -match "Windows 11" })) {
             Dismount-DiskImage -ImagePath $isoPath | Out-Null
             Write-Win11ISOLog "ERROR: No 'Windows 11' edition found in the image."
             [System.Windows.MessageBox]::Show(
@@ -113,10 +83,8 @@ function Invoke-WinUtilISOMountAndVerify {
             return
         }
 
-        # Store edition info for later index lookup
         $sync["Win11ISOImageInfo"] = $imageInfo
 
-        # ── Populate UI ──
         $sync["WPFWin11ISOMountDriveLetter"].Text = "Mounted at: $driveLetter   |   Image file: $(Split-Path $activeWim -Leaf)"
         $sync["WPFWin11ISOEditionComboBox"].Dispatcher.Invoke([action]{
             $sync["WPFWin11ISOEditionComboBox"].Items.Clear()
@@ -124,12 +92,10 @@ function Invoke-WinUtilISOMountAndVerify {
                 [void]$sync["WPFWin11ISOEditionComboBox"].Items.Add("$($img.ImageIndex): $($img.ImageName)")
             }
             if ($sync["WPFWin11ISOEditionComboBox"].Items.Count -gt 0) {
-                # Default to Windows 11 Pro; fall back to first item if not found
                 $proIndex = -1
                 for ($i = 0; $i -lt $sync["WPFWin11ISOEditionComboBox"].Items.Count; $i++) {
                     if ($sync["WPFWin11ISOEditionComboBox"].Items[$i] -match "Windows 11 Pro(?![\w ])") {
-                        $proIndex = $i
-                        break
+                        $proIndex = $i; break
                     }
                 }
                 $sync["WPFWin11ISOEditionComboBox"].SelectedIndex = if ($proIndex -ge 0) { $proIndex } else { 0 }
@@ -137,40 +103,25 @@ function Invoke-WinUtilISOMountAndVerify {
         })
         $sync["WPFWin11ISOVerifyResultPanel"].Visibility = "Visible"
 
-        # Store for later steps
         $sync["Win11ISODriveLetter"] = $driveLetter
         $sync["Win11ISOWimPath"]     = $activeWim
         $sync["Win11ISOImagePath"]   = $isoPath
-
-        # Reveal Step 3
         $sync["WPFWin11ISOModifySection"].Visibility = "Visible"
 
-        Set-WinUtilProgressBar -Label "ISO verified ✔" -Percent 100
+        Set-WinUtilProgressBar -Label "ISO verified" -Percent 100
         Write-Win11ISOLog "ISO verified OK.  Editions found: $($imageInfo.Count)"
-    }
-    catch {
+    } catch {
         Write-Win11ISOLog "ERROR during mount/verify: $_"
         [System.Windows.MessageBox]::Show(
             "An error occurred while mounting or verifying the ISO:`n`n$_",
             "Error", "OK", "Error")
-    }
-    finally {
+    } finally {
         Start-Sleep -Milliseconds 800
         Set-WinUtilProgressBar -Label "" -Percent 0
     }
 }
 
 function Invoke-WinUtilISOModify {
-    <#
-    .SYNOPSIS
-        Extracts ISO contents to a temp working directory, modifies install.wim,
-        then repackages the image.  Reveals Step 4 (output options) on success.
-
-    .NOTES
-        This function runs inside a PowerShell runspace so the UI stays responsive.
-        Placeholder modification logic is provided; extend as needed.
-    #>
-
     $isoPath     = $sync["Win11ISOImagePath"]
     $driveLetter = $sync["Win11ISODriveLetter"]
     $wimPath     = $sync["Win11ISOWimPath"]
@@ -182,9 +133,8 @@ function Invoke-WinUtilISOModify {
         return
     }
 
-    # ── Resolve selected edition index from the ComboBox ──
     $selectedItem     = $sync["WPFWin11ISOEditionComboBox"].SelectedItem
-    $selectedWimIndex = 1  # default fallback
+    $selectedWimIndex = 1
     if ($selectedItem -and $selectedItem -match '^(\d+):') {
         $selectedWimIndex = [int]$Matches[1]
     } elseif ($sync["Win11ISOImageInfo"]) {
@@ -193,13 +143,10 @@ function Invoke-WinUtilISOModify {
     $selectedEditionName = if ($selectedItem) { ($selectedItem -replace '^\d+:\s*', '') } else { "Unknown" }
     Write-Win11ISOLog "Selected edition: $selectedEditionName (Index $selectedWimIndex)"
 
-    # Disable the modify button to prevent double-click
     $sync["WPFWin11ISOModifyButton"].IsEnabled = $false
 
     $existingWorkDir = Get-Item -Path (Join-Path $env:TEMP "WinUtil_Win11ISO*") -ErrorAction SilentlyContinue |
-        Where-Object { $_.PSIsContainer } |
-        Sort-Object LastWriteTime -Descending |
-        Select-Object -First 1
+        Where-Object { $_.PSIsContainer } | Sort-Object LastWriteTime -Descending | Select-Object -First 1
 
     $workDir = if ($existingWorkDir) {
         Write-Win11ISOLog "Reusing existing temp directory: $($existingWorkDir.FullName)"
@@ -208,9 +155,6 @@ function Invoke-WinUtilISOModify {
         Join-Path $env:TEMP "WinUtil_Win11ISO_$(Get-Date -Format 'yyyyMMdd_HHmmss')"
     }
 
-    # ── Resolve autounattend.xml content ──────────────────────────────────────
-    # Compiled winutil.ps1 sets $WinUtilAutounattendXml before main.ps1 runs.
-    # In dev/source mode fall back to reading tools\autounattend.xml directly.
     $autounattendContent = if ($WinUtilAutounattendXml) {
         $WinUtilAutounattendXml
     } else {
@@ -218,11 +162,12 @@ function Invoke-WinUtilISOModify {
         if (Test-Path $toolsXml) { Get-Content $toolsXml -Raw } else { "" }
     }
 
-    # ── Run modification in a background runspace ──
     $runspace = [Management.Automation.Runspaces.RunspaceFactory]::CreateRunspace()
     $runspace.ApartmentState = "STA"
     $runspace.ThreadOptions  = "ReuseThread"
     $runspace.Open()
+    $injectDrivers = $sync["WPFWin11ISOInjectDrivers"].IsChecked -eq $true
+
     $runspace.SessionStateProxy.SetVariable("sync",                $sync)
     $runspace.SessionStateProxy.SetVariable("isoPath",             $isoPath)
     $runspace.SessionStateProxy.SetVariable("driveLetter",         $driveLetter)
@@ -231,28 +176,18 @@ function Invoke-WinUtilISOModify {
     $runspace.SessionStateProxy.SetVariable("selectedWimIndex",    $selectedWimIndex)
     $runspace.SessionStateProxy.SetVariable("selectedEditionName", $selectedEditionName)
     $runspace.SessionStateProxy.SetVariable("autounattendContent", $autounattendContent)
+    $runspace.SessionStateProxy.SetVariable("injectDrivers",       $injectDrivers)
 
-    # Serialize functions so they are available inside the runspace
-    $isoScriptFuncDef = "function Invoke-WinUtilISOScript {`n" + `
-        ${function:Invoke-WinUtilISOScript}.ToString() + "`n}"
+    $isoScriptFuncDef   = "function Invoke-WinUtilISOScript {`n" + ${function:Invoke-WinUtilISOScript}.ToString() + "`n}"
+    $win11ISOLogFuncDef = "function Write-Win11ISOLog {`n"       + ${function:Write-Win11ISOLog}.ToString()       + "`n}"
     $runspace.SessionStateProxy.SetVariable("isoScriptFuncDef",   $isoScriptFuncDef)
-
-    $win11ISOLogFuncDef = "function Write-Win11ISOLog {`n" + `
-        ${function:Write-Win11ISOLog}.ToString() + "`n}"
     $runspace.SessionStateProxy.SetVariable("win11ISOLogFuncDef", $win11ISOLogFuncDef)
 
-    $refreshUSBFuncDef = "function Invoke-WinUtilISORefreshUSBDrives {`n" + `
-        ${function:Invoke-WinUtilISORefreshUSBDrives}.ToString() + "`n}"
-    $runspace.SessionStateProxy.SetVariable("refreshUSBFuncDef", $refreshUSBFuncDef)
-
     $script = [Management.Automation.PowerShell]::Create()
     $script.Runspace = $runspace
     $script.AddScript({
-
-        # Import helper functions into this runspace
         . ([scriptblock]::Create($isoScriptFuncDef))
         . ([scriptblock]::Create($win11ISOLogFuncDef))
-        . ([scriptblock]::Create($refreshUSBFuncDef))
 
         function Log($msg) {
             $ts = (Get-Date).ToString("HH:mm:ss")
@@ -261,6 +196,7 @@ function Invoke-WinUtilISOModify {
                 $sync["WPFWin11ISOStatusLog"].CaretIndex = $sync["WPFWin11ISOStatusLog"].Text.Length
                 $sync["WPFWin11ISOStatusLog"].ScrollToEnd()
             })
+            Add-Content -Path (Join-Path $workDir "WinUtil_Win11ISO.log") -Value "[$ts] $msg" -ErrorAction SilentlyContinue
         }
 
         function SetProgress($label, $pct) {
@@ -272,173 +208,111 @@ function Invoke-WinUtilISOModify {
         }
 
         try {
-            # ── Hide Steps 1-3 while modification is running; expand log to fill screen ──
             $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
                 $sync["WPFWin11ISOSelectSection"].Visibility = "Collapsed"
                 $sync["WPFWin11ISOMountSection"].Visibility  = "Collapsed"
                 $sync["WPFWin11ISOModifySection"].Visibility = "Collapsed"
-                $expandedHeight = [Math]::Max(400, $sync["Form"].ActualHeight - 100)
-                $sync["WPFWin11ISOStatusLog"].Height = $expandedHeight
-                $sync["Win11ISOLogExpanded"] = $true
-                # Register the resize handler once so the log tracks window resizes
-                if (-not $sync["Win11ISOResizeHandlerAdded"]) {
-                    $sync["Form"].add_SizeChanged({
-                        if ($sync["Win11ISOLogExpanded"]) {
-                            $sync["WPFWin11ISOStatusLog"].Height = [Math]::Max(400, $sync["Form"].ActualHeight - 100)
-                            $sync["WPFWin11ISOStatusLog"].CaretIndex = $sync["WPFWin11ISOStatusLog"].Text.Length
-                            $sync["WPFWin11ISOStatusLog"].ScrollToEnd()
-                        }
-                    })
-                    $sync["Win11ISOResizeHandlerAdded"] = $true
-                }
             })
 
-            # ── 1. Create working directory structure ──
             Log "Creating working directory: $workDir"
             $isoContents = Join-Path $workDir "iso_contents"
             $mountDir    = Join-Path $workDir "wim_mount"
             New-Item -ItemType Directory -Path $isoContents, $mountDir -Force | Out-Null
             SetProgress "Copying ISO contents..." 10
 
-            # ── 2. Copy all ISO contents to the working directory ──
             Log "Copying ISO contents from $driveLetter to $isoContents..."
-            $robocopyArgs = @($driveLetter, $isoContents, "/E", "/NFL", "/NDL", "/NJH", "/NJS")
-            & robocopy @robocopyArgs | Out-Null
+            & robocopy $driveLetter $isoContents /E /NFL /NDL /NJH /NJS | Out-Null
             Log "ISO contents copied."
             SetProgress "Mounting install.wim..." 25
 
-            # ── 3. Copy install.wim to working dir (it may be read-only on the DVD) ──
             $localWim = Join-Path $isoContents "sources\install.wim"
-            if (-not (Test-Path $localWim)) {
-                # ESD path
-                $localWim = Join-Path $isoContents "sources\install.esd"
-            }
-            # Ensure the file is writable
+            if (-not (Test-Path $localWim)) { $localWim = Join-Path $isoContents "sources\install.esd" }
             Set-ItemProperty -Path $localWim -Name IsReadOnly -Value $false
 
-            # ── 4. Mount the selected edition of install.wim ──
             Log "Mounting install.wim (Index ${selectedWimIndex}: $selectedEditionName) at $mountDir..."
             Mount-WindowsImage -ImagePath $localWim -Index $selectedWimIndex -Path $mountDir -ErrorAction Stop | Out-Null
             SetProgress "Modifying install.wim..." 45
 
-            # ── Apply all WinUtil modifications via Invoke-WinUtilISOScript ──
             Log "Applying WinUtil modifications to install.wim..."
-            Invoke-WinUtilISOScript -ScratchDir $mountDir -ISOContentsDir $isoContents -AutoUnattendXml $autounattendContent -Log { param($m) Log $m }
+            Invoke-WinUtilISOScript -ScratchDir $mountDir -ISOContentsDir $isoContents -AutoUnattendXml $autounattendContent -InjectCurrentSystemDrivers $injectDrivers -Log { param($m) Log $m }
 
-            # ── 4b. DISM component store cleanup ──
-            # /ResetBase removes all superseded component versions from WinSxS,
-            # which is the single largest space saving possible (typically 300–800 MB).
-            # This must be done while the image is still mounted.
             SetProgress "Cleaning up component store (WinSxS)..." 56
             Log "Running DISM component store cleanup (/ResetBase)..."
             & dism /English "/image:$mountDir" /Cleanup-Image /StartComponentCleanup /ResetBase | ForEach-Object { Log $_ }
             Log "Component store cleanup complete."
 
-            # ── 5. Save and dismount the WIM ──
             SetProgress "Saving modified install.wim..." 65
             Log "Dismounting and saving install.wim. This will take several minutes..."
             Dismount-WindowsImage -Path $mountDir -Save -ErrorAction Stop | Out-Null
             Log "install.wim saved."
 
-            # ── 5b. Strip unused editions — export only the selected index ──
-            # A standard multi-edition install.wim can be 4–5 GB; exporting a
-            # single index typically drops it to ~3 GB, saving 1–2 GB in the ISO.
             SetProgress "Removing unused editions from install.wim..." 70
             Log "Exporting edition '$selectedEditionName' (Index $selectedWimIndex) to a single-edition install.wim..."
             $exportWim = Join-Path $isoContents "sources\install_export.wim"
-            Export-WindowsImage `
-                -SourceImagePath $localWim `
-                -SourceIndex     $selectedWimIndex `
-                -DestinationImagePath $exportWim `
-                -ErrorAction Stop | Out-Null
+            Export-WindowsImage -SourceImagePath $localWim -SourceIndex $selectedWimIndex -DestinationImagePath $exportWim -ErrorAction Stop | Out-Null
             Remove-Item -Path $localWim -Force
             Rename-Item -Path $exportWim -NewName "install.wim" -Force
-            # Update local path so later steps (e.g. ISO build) reference the new file
             $localWim = Join-Path $isoContents "sources\install.wim"
             Log "Unused editions removed. install.wim now contains only '$selectedEditionName'."
 
             SetProgress "Dismounting source ISO..." 80
-
-            # ── 6. Dismount the original ISO ──
             Log "Dismounting original ISO..."
             Dismount-DiskImage -ImagePath $isoPath | Out-Null
 
-            # Store work directory for output steps
             $sync["Win11ISOWorkDir"]     = $workDir
             $sync["Win11ISOContentsDir"] = $isoContents
 
-            SetProgress "Modification complete ✔" 100
+            SetProgress "Modification complete" 100
             Log "install.wim modification complete. Choose an output option in Step 4."
 
-            # ── Reveal Step 4 on the UI thread ──
-            # Note: USB drive enumeration (Get-Disk) is intentionally deferred to
-            # when the user explicitly selects the USB option, to avoid blocking
-            # the UI thread here.
             $sync["WPFWin11ISOOutputSection"].Dispatcher.Invoke([action]{
                 $sync["WPFWin11ISOOutputSection"].Visibility = "Visible"
             })
-        }
-        catch {
+        } catch {
             Log "ERROR during modification: $_"
 
-            # ── Cleanup: dismount WIM if still mounted ──
             try {
                 if (Test-Path $mountDir) {
-                    $mountedImages = Get-WindowsImage -Mounted -ErrorAction SilentlyContinue |
-                                     Where-Object { $_.Path -eq $mountDir }
+                    $mountedImages = Get-WindowsImage -Mounted -ErrorAction SilentlyContinue | Where-Object { $_.Path -eq $mountDir }
                     if ($mountedImages) {
                         Log "Cleaning up: dismounting install.wim (discarding changes)..."
                         Dismount-WindowsImage -Path $mountDir -Discard -ErrorAction SilentlyContinue | Out-Null
                     }
                 }
-            } catch {
-                Log "Warning: could not dismount install.wim during cleanup: $_"
-            }
+            } catch { Log "Warning: could not dismount install.wim during cleanup: $_" }
 
-            # ── Cleanup: dismount the source ISO ──
             try {
                 $mountedISO = Get-DiskImage -ImagePath $isoPath -ErrorAction SilentlyContinue
                 if ($mountedISO -and $mountedISO.Attached) {
                     Log "Cleaning up: dismounting source ISO..."
                     Dismount-DiskImage -ImagePath $isoPath -ErrorAction SilentlyContinue | Out-Null
                 }
-            } catch {
-                Log "Warning: could not dismount ISO during cleanup: $_"
-            }
+            } catch { Log "Warning: could not dismount ISO during cleanup: $_" }
 
-            # ── Cleanup: remove temp working directory ──
             try {
                 if (Test-Path $workDir) {
                     Log "Cleaning up: removing temp directory $workDir..."
                     Remove-Item -Path $workDir -Recurse -Force -ErrorAction SilentlyContinue
                 }
-            } catch {
-                Log "Warning: could not remove temp directory during cleanup: $_"
-            }
+            } catch { Log "Warning: could not remove temp directory during cleanup: $_" }
 
             $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
                 [System.Windows.MessageBox]::Show(
                     "An error occurred during install.wim modification:`n`n$_",
                     "Modification Error", "OK", "Error")
             })
-        }
-        finally {
+        } finally {
             Start-Sleep -Milliseconds 800
             $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
                 $sync.progressBarTextBlock.Text    = ""
                 $sync.progressBarTextBlock.ToolTip = ""
                 $sync.ProgressBar.Value            = 0
                 $sync["WPFWin11ISOModifyButton"].IsEnabled = $true
-                # ── Only restore steps 1-3 if Step 4 was NOT successfully shown ──
-                # When modification succeeds, Step 4 is visible and steps 1-3 stay
-                # hidden until the user clicks Clean & Reset.
                 if ($sync["WPFWin11ISOOutputSection"].Visibility -ne "Visible") {
                     $sync["WPFWin11ISOSelectSection"].Visibility = "Visible"
                     $sync["WPFWin11ISOMountSection"].Visibility  = "Visible"
                     $sync["WPFWin11ISOModifySection"].Visibility = "Visible"
                 }
-                $sync["Win11ISOLogExpanded"] = $false
-                $sync["WPFWin11ISOStatusLog"].Height = 140
             })
         }
     }) | Out-Null
@@ -446,13 +320,36 @@ function Invoke-WinUtilISOModify {
     $script.BeginInvoke() | Out-Null
 }
 
-function Invoke-WinUtilISOCleanAndReset {
-    <#
-    .SYNOPSIS
-        Deletes the temporary working directory created during ISO modification
-        and resets the entire ISO UI back to its initial state (Step 1 only).
-    #>
+function Invoke-WinUtilISOCheckExistingWork {
+    if ($sync["Win11ISOContentsDir"] -and (Test-Path $sync["Win11ISOContentsDir"])) { return }
 
+    $existingWorkDir = Get-Item -Path (Join-Path $env:TEMP "WinUtil_Win11ISO*") -ErrorAction SilentlyContinue |
+        Where-Object { $_.PSIsContainer } | Sort-Object LastWriteTime -Descending | Select-Object -First 1
+
+    if (-not $existingWorkDir) { return }
+
+    $isoContents = Join-Path $existingWorkDir.FullName "iso_contents"
+    if (-not (Test-Path $isoContents)) { return }
+
+    $sync["Win11ISOWorkDir"]     = $existingWorkDir.FullName
+    $sync["Win11ISOContentsDir"] = $isoContents
+
+    $sync["WPFWin11ISOSelectSection"].Visibility = "Collapsed"
+    $sync["WPFWin11ISOMountSection"].Visibility  = "Collapsed"
+    $sync["WPFWin11ISOModifySection"].Visibility = "Collapsed"
+    $sync["WPFWin11ISOOutputSection"].Visibility = "Visible"
+
+    $modified = $existingWorkDir.LastWriteTime.ToString("yyyy-MM-dd HH:mm")
+    Write-Win11ISOLog "Existing working directory found: $($existingWorkDir.FullName)"
+    Write-Win11ISOLog "Last modified: $modified - Skipping Steps 1-3 and resuming at Step 4."
+    Write-Win11ISOLog "Click 'Clean & Reset' if you want to start over with a new ISO."
+
+    [System.Windows.MessageBox]::Show(
+        "A previous WinUtil ISO working directory was found:`n`n$($existingWorkDir.FullName)`n`n(Last modified: $modified)`n`nStep 4 (output options) has been restored so you can save the already-modified image.`n`nClick 'Clean & Reset' in Step 4 if you want to start over.",
+        "Existing Work Found", "OK", "Info")
+}
+
+function Invoke-WinUtilISOCleanAndReset {
     $workDir = $sync["Win11ISOWorkDir"]
 
     if ($workDir -and (Test-Path $workDir)) {
@@ -460,17 +357,104 @@ function Invoke-WinUtilISOCleanAndReset {
             "This will delete the temporary working directory:`n`n$workDir`n`nAnd reset the interface back to the start.`n`nContinue?",
             "Clean & Reset", "YesNo", "Warning")
         if ($confirm -ne "Yes") { return }
+    }
+
+    $sync["WPFWin11ISOCleanResetButton"].IsEnabled = $false
+
+    $runspace = [Management.Automation.Runspaces.RunspaceFactory]::CreateRunspace()
+    $runspace.ApartmentState = "STA"
+    $runspace.ThreadOptions  = "ReuseThread"
+    $runspace.Open()
+    $runspace.SessionStateProxy.SetVariable("sync",    $sync)
+    $runspace.SessionStateProxy.SetVariable("workDir", $workDir)
+
+    $script = [Management.Automation.PowerShell]::Create()
+    $script.Runspace = $runspace
+    $script.AddScript({
+
+        function Log($msg) {
+            $ts = (Get-Date).ToString("HH:mm:ss")
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                $sync["WPFWin11ISOStatusLog"].Text += "`n[$ts] $msg"
+                $sync["WPFWin11ISOStatusLog"].CaretIndex = $sync["WPFWin11ISOStatusLog"].Text.Length
+                $sync["WPFWin11ISOStatusLog"].ScrollToEnd()
+            })
+            Add-Content -Path (Join-Path $workDir "WinUtil_Win11ISO.log") -Value "[$ts] $msg" -ErrorAction SilentlyContinue
+        }
+
+        function SetProgress($label, $pct) {
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                $sync.progressBarTextBlock.Text    = $label
+                $sync.progressBarTextBlock.ToolTip = $label
+                $sync.ProgressBar.Value            = [Math]::Max($pct, 5)
+            })
+        }
 
         try {
-            Write-Win11ISOLog "Deleting temp directory: $workDir"
-            Remove-Item -Path $workDir -Recurse -Force -ErrorAction Stop
-            Write-Win11ISOLog "Temp directory deleted."
+            if ($workDir) {
+                $mountDir = Join-Path $workDir "wim_mount"
+                try {
+                    $mountedImages = Get-WindowsImage -Mounted -ErrorAction SilentlyContinue |
+                                     Where-Object { $_.Path -like "$workDir*" }
+                    if ($mountedImages) {
+                        foreach ($img in $mountedImages) {
+                            Log "Dismounting WIM at: $($img.Path) (discarding changes)..."
+                            SetProgress "Dismounting WIM image..." 3
+                            Dismount-WindowsImage -Path $img.Path -Discard -ErrorAction Stop | Out-Null
+                            Log "WIM dismounted successfully."
+                        }
+                    } elseif (Test-Path $mountDir) {
+                        Log "No mounted WIM reported by Get-WindowsImage, running DISM /Cleanup-Wim as a precaution..."
+                        SetProgress "Running DISM cleanup..." 3
+                        & dism /English /Cleanup-Wim 2>&1 | ForEach-Object { Log $_ }
+                    }
                 } catch {
-            Write-Win11ISOLog "WARNING: could not fully delete temp directory: $_"
+                    Log "Warning: could not dismount WIM cleanly, attempting DISM /Cleanup-Wim fallback: $_"
+                    try { & dism /English /Cleanup-Wim 2>&1 | ForEach-Object { Log $_ } }
+                    catch { Log "Warning: DISM /Cleanup-Wim also failed: $_" }
                 }
             }
 
-    # Clear all stored ISO state
+            if ($workDir -and (Test-Path $workDir)) {
+                Log "Scanning files to delete in: $workDir"
+                SetProgress "Scanning files..." 5
+
+                $allFiles = @(Get-ChildItem -Path $workDir -File -Recurse -Force -ErrorAction SilentlyContinue)
+                $allDirs  = @(Get-ChildItem -Path $workDir -Directory -Recurse -Force -ErrorAction SilentlyContinue |
+                    Sort-Object { $_.FullName.Length } -Descending)
+                $total   = $allFiles.Count
+                $deleted = 0
+
+                Log "Found $total files to delete."
+
+                foreach ($f in $allFiles) {
+                    try { Remove-Item -Path $f.FullName -Force -ErrorAction Stop } catch { Log "WARNING: could not delete $($f.FullName): $_" }
+                    $deleted++
+                    if ($deleted % 100 -eq 0 -or $deleted -eq $total) {
+                        $pct = [math]::Round(($deleted / [Math]::Max($total, 1)) * 85) + 5
+                        SetProgress "Deleting files in $($f.Directory.Name)... ($deleted / $total)" $pct
+                    }
+                }
+
+                foreach ($d in $allDirs) {
+                    try { Remove-Item -Path $d.FullName -Force -ErrorAction SilentlyContinue } catch {}
+                }
+
+                try { Remove-Item -Path $workDir -Recurse -Force -ErrorAction Stop } catch {}
+
+                if (Test-Path $workDir) {
+                    Log "WARNING: some items could not be deleted in $workDir"
+                } else {
+                    Log "Temp directory deleted successfully."
+                }
+            } else {
+                Log "No temp directory found — resetting UI."
+            }
+
+            SetProgress "Resetting UI..." 95
+            Log "Resetting interface..."
+
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
                 $sync["Win11ISOWorkDir"]     = $null
                 $sync["Win11ISOContentsDir"] = $null
                 $sync["Win11ISOImagePath"]   = $null
@@ -479,7 +463,6 @@ function Invoke-WinUtilISOCleanAndReset {
                 $sync["Win11ISOImageInfo"]   = $null
                 $sync["Win11ISOUSBDisks"]    = $null
 
-    # Reset the UI to the initial state
                 $sync["WPFWin11ISOPath"].Text                   = "No ISO selected..."
                 $sync["WPFWin11ISOFileInfo"].Visibility          = "Collapsed"
                 $sync["WPFWin11ISOVerifyResultPanel"].Visibility = "Collapsed"
@@ -488,23 +471,35 @@ function Invoke-WinUtilISOCleanAndReset {
                 $sync["WPFWin11ISOModifySection"].Visibility     = "Collapsed"
                 $sync["WPFWin11ISOMountSection"].Visibility      = "Collapsed"
                 $sync["WPFWin11ISOSelectSection"].Visibility     = "Visible"
-    $sync["WPFWin11ISOStatusLog"].Text                 = "Ready. Please select a Windows 11 ISO to begin."
-    $sync["WPFWin11ISOStatusLog"].Height               = 140
                 $sync["WPFWin11ISOModifyButton"].IsEnabled       = $true
+                $sync["WPFWin11ISOCleanResetButton"].IsEnabled   = $true
+
+                $sync.progressBarTextBlock.Text    = ""
+                $sync.progressBarTextBlock.ToolTip = ""
+                $sync.ProgressBar.Value            = 0
+
+                $sync["WPFWin11ISOStatusLog"].Text   = "Ready. Please select a Windows 11 ISO to begin."
+            })
+        } catch {
+            Log "ERROR during Clean & Reset: $_"
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                $sync.progressBarTextBlock.Text    = ""
+                $sync.progressBarTextBlock.ToolTip = ""
+                $sync.ProgressBar.Value            = 0
+                $sync["WPFWin11ISOCleanResetButton"].IsEnabled = $true
+            })
+        }
+    }) | Out-Null
+
+    $script.BeginInvoke() | Out-Null
 }
 
 function Invoke-WinUtilISOExport {
-    <#
-    .SYNOPSIS
-        Saves the modified ISO contents as a new bootable ISO file.
-        Uses oscdimg.exe (part of the Windows ADK) if present; falls back
-        to a reminder message if not installed.
-    #>
     $contentsDir = $sync["Win11ISOContentsDir"]
 
     if (-not $contentsDir -or -not (Test-Path $contentsDir)) {
         [System.Windows.MessageBox]::Show(
-            "No modified ISO content found.  Please complete Steps 1–3 first.",
+            "No modified ISO content found.  Please complete Steps 1-3 first.",
             "Not Ready", "OK", "Warning")
         return
     }
@@ -520,9 +515,6 @@ function Invoke-WinUtilISOExport {
     if ($dlg.ShowDialog() -ne [System.Windows.Forms.DialogResult]::OK) { return }
 
     $outputISO = $dlg.FileName
-    Write-Win11ISOLog "Exporting to ISO: $outputISO"
-
-    Set-WinUtilProgressBar -Label "Building ISO..." -Percent 10
 
     # Locate oscdimg.exe (Windows ADK or winget per-user install)
     $oscdimg = Get-ChildItem "C:\Program Files (x86)\Windows Kits" -Recurse -Filter "oscdimg.exe" -ErrorAction SilentlyContinue |
@@ -535,12 +527,10 @@ function Invoke-WinUtilISOExport {
 
     if (-not $oscdimg) {
         Write-Win11ISOLog "oscdimg.exe not found. Attempting to install via winget..."
-        Set-WinUtilProgressBar -Label "Installing oscdimg..." -Percent 5
         try {
             $winget = Get-Command winget -ErrorAction Stop
             $result = & $winget install -e --id Microsoft.OSCDIMG --accept-package-agreements --accept-source-agreements 2>&1
             Write-Win11ISOLog "winget output: $result"
-            # Re-scan after install
             $oscdimg = Get-ChildItem "$env:LOCALAPPDATA\Microsoft\WinGet\Packages" -Recurse -Filter "oscdimg.exe" -ErrorAction SilentlyContinue |
                        Where-Object { $_.FullName -match 'Microsoft\.OSCDIMG' } |
                        Select-Object -First 1 -ExpandProperty FullName
@@ -549,7 +539,6 @@ function Invoke-WinUtilISOExport {
         }
 
         if (-not $oscdimg) {
-            Set-WinUtilProgressBar -Label "" -Percent 0
             Write-Win11ISOLog "oscdimg.exe still not found after install attempt."
             [System.Windows.MessageBox]::Show(
                 "oscdimg.exe could not be found or installed automatically.`n`nPlease install it manually:`n  winget install -e --id Microsoft.OSCDIMG`n`nOr install the Windows ADK from:`nhttps://learn.microsoft.com/windows-hardware/get-started/adk-install",
@@ -559,21 +548,42 @@ function Invoke-WinUtilISOExport {
         Write-Win11ISOLog "oscdimg.exe installed successfully."
     }
 
-    # Build boot parameters (BIOS + UEFI dual-boot)
-    $bootData   = "2#p0,e,b`"$contentsDir\boot\etfsboot.com`"#pEF,e,b`"$contentsDir\efi\microsoft\boot\efisys.bin`""
-    $oscdimgArgs = @(
-        "-m",           # ignore source path max size
-        "-o",           # optimise storage
-        "-u2",          # UDF 2.01
-        "-udfver102",
-        "-bootdata:$bootData",
-        "-l`"CTOS_MODIFIED`"",
-        "`"$contentsDir`"",
-        "`"$outputISO`""
-    )
+    $sync["WPFWin11ISOChooseISOButton"].IsEnabled = $false
+
+    $runspace = [Management.Automation.Runspaces.RunspaceFactory]::CreateRunspace()
+    $runspace.ApartmentState = "STA"
+    $runspace.ThreadOptions  = "ReuseThread"
+    $runspace.Open()
+    $runspace.SessionStateProxy.SetVariable("sync",        $sync)
+    $runspace.SessionStateProxy.SetVariable("contentsDir", $contentsDir)
+    $runspace.SessionStateProxy.SetVariable("outputISO",   $outputISO)
+    $runspace.SessionStateProxy.SetVariable("oscdimg",     $oscdimg)
+
+    $win11ISOLogFuncDef = "function Write-Win11ISOLog {`n" + ${function:Write-Win11ISOLog}.ToString() + "`n}"
+    $runspace.SessionStateProxy.SetVariable("win11ISOLogFuncDef", $win11ISOLogFuncDef)
+
+    $script = [Management.Automation.PowerShell]::Create()
+    $script.Runspace = $runspace
+    $script.AddScript({
+        . ([scriptblock]::Create($win11ISOLogFuncDef))
+
+        function SetProgress($label, $pct) {
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                $sync.progressBarTextBlock.Text    = $label
+                $sync.progressBarTextBlock.ToolTip = $label
+                $sync.ProgressBar.Value            = [Math]::Max($pct, 5)
+            })
+        }
 
         try {
+            Write-Win11ISOLog "Exporting to ISO: $outputISO"
+            SetProgress "Building ISO..." 10
+
+            $bootData    = "2#p0,e,b`"$contentsDir\boot\etfsboot.com`"#pEF,e,b`"$contentsDir\efi\microsoft\boot\efisys.bin`""
+            $oscdimgArgs = @("-m", "-o", "-u2", "-udfver102", "-bootdata:$bootData", "-l`"CTOS_MODIFIED`"", "`"$contentsDir`"", "`"$outputISO`"")
+
             Write-Win11ISOLog "Running oscdimg..."
+
             $psi = [System.Diagnostics.ProcessStartInfo]::new()
             $psi.FileName               = $oscdimg
             $psi.Arguments              = $oscdimgArgs -join " "
@@ -586,230 +596,46 @@ function Invoke-WinUtilISOExport {
             $proc.StartInfo = $psi
             $proc.Start() | Out-Null
 
-        # Stream stdout and stderr line-by-line to the status log
-        $stdoutTask = $proc.StandardOutput.ReadToEndAsync()
-        $stderrTask = $proc.StandardError.ReadToEndAsync()
-        $proc.WaitForExit()
-        [System.Threading.Tasks.Task]::WaitAll($stdoutTask, $stderrTask)
-
-        foreach ($line in ($stdoutTask.Result -split "`r?`n")) {
+            # Stream stdout line-by-line as oscdimg runs
+            while (-not $proc.StandardOutput.EndOfStream) {
+                $line = $proc.StandardOutput.ReadLine()
                 if ($line.Trim()) { Write-Win11ISOLog $line }
             }
-        foreach ($line in ($stderrTask.Result -split "`r?`n")) {
+
+            $proc.WaitForExit()
+
+            # Flush any stderr after process exits
+            $stderr = $proc.StandardError.ReadToEnd()
+            foreach ($line in ($stderr -split "`r?`n")) {
                 if ($line.Trim()) { Write-Win11ISOLog "[stderr]$line" }
             }
 
             if ($proc.ExitCode -eq 0) {
-            Set-WinUtilProgressBar -Label "ISO exported ✔" -Percent 100
+                SetProgress "ISO exported" 100
                 Write-Win11ISOLog "ISO exported successfully: $outputISO"
-            [System.Windows.MessageBox]::Show(
-                "ISO exported successfully!`n`n$outputISO",
-                "Export Complete", "OK", "Info")
+                $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                    [System.Windows.MessageBox]::Show("ISO exported successfully!`n`n$outputISO", "Export Complete", "OK", "Info")
+                })
             } else {
                 Write-Win11ISOLog "oscdimg exited with code $($proc.ExitCode)."
+                $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
                     [System.Windows.MessageBox]::Show(
                         "oscdimg exited with code $($proc.ExitCode).`nCheck the status log for details.",
                         "Export Error", "OK", "Error")
+                })
             }
-    }
-    catch {
+        } catch {
             Write-Win11ISOLog "ERROR during ISO export: $_"
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
                 [System.Windows.MessageBox]::Show("ISO export failed:`n`n$_", "Error", "OK", "Error")
-    }
-    finally {
-        Start-Sleep -Milliseconds 800
-        Set-WinUtilProgressBar -Label "" -Percent 0
-    }
-}
-
-function Invoke-WinUtilISORefreshUSBDrives {
-    <#
-    .SYNOPSIS
-        Populates the USB drive ComboBox with all currently attached removable drives.
-    #>
-    $combo = $sync["WPFWin11ISOUSBDriveComboBox"]
-    $combo.Items.Clear()
-
-    $removable = Get-Disk | Where-Object { $_.BusType -eq "USB" } | Sort-Object Number
-
-    if ($removable.Count -eq 0) {
-        $combo.Items.Add("No USB drives detected")
-        $combo.SelectedIndex = 0
-        Write-Win11ISOLog "No USB drives detected."
-        return
-    }
-
-    foreach ($disk in $removable) {
-        $sizeGB    = [math]::Round($disk.Size / 1GB, 1)
-        $label     = "Disk $($disk.Number): $($disk.FriendlyName)  [$sizeGB GB]  — $($disk.PartitionStyle)"
-        $combo.Items.Add($label)
-    }
-    $combo.SelectedIndex = 0
-    Write-Win11ISOLog "Found $($removable.Count) USB drive(s)."
-
-    # Store disk objects for later use
-    $sync["Win11ISOUSBDisks"] = $removable
-}
-
-function Invoke-WinUtilISOWriteUSB {
-    <#
-    .SYNOPSIS
-        Erases the selected USB drive and writes the modified Windows 11 ISO
-        content as a bootable installation drive (using DISM / robocopy approach).
-    #>
-    $contentsDir = $sync["Win11ISOContentsDir"]
-    $usbDisks    = $sync["Win11ISOUSBDisks"]
-
-    if (-not $contentsDir -or -not (Test-Path $contentsDir)) {
-        [System.Windows.MessageBox]::Show(
-            "No modified ISO content found.  Please complete Steps 1–3 first.",
-            "Not Ready", "OK", "Warning")
-        return
-    }
-
-    $selectedIndex = $sync["WPFWin11ISOUSBDriveComboBox"].SelectedIndex
-    if ($selectedIndex -lt 0 -or -not $usbDisks -or $selectedIndex -ge $usbDisks.Count) {
-        [System.Windows.MessageBox]::Show(
-            "Please select a USB drive from the dropdown.",
-            "No Drive Selected", "OK", "Warning")
-        return
-    }
-
-    $targetDisk = $usbDisks[$selectedIndex]
-    $diskNum    = $targetDisk.Number
-    $sizeGB     = [math]::Round($targetDisk.Size / 1GB, 1)
-
-    $confirm = [System.Windows.MessageBox]::Show(
-        "ALL data on Disk $diskNum ($($targetDisk.FriendlyName), $sizeGB GB) will be PERMANENTLY ERASED.`n`nAre you sure you want to continue?",
-        "Confirm USB Erase", "YesNo", "Warning")
-
-    if ($confirm -ne "Yes") {
-        Write-Win11ISOLog "USB write cancelled by user."
-        return
-    }
-
-    $sync["WPFWin11ISOWriteUSBButton"].IsEnabled = $false
-    Write-Win11ISOLog "Starting USB write to Disk $diskNum..."
-
-    $runspace = [Management.Automation.Runspaces.RunspaceFactory]::CreateRunspace()
-    $runspace.ApartmentState = "STA"
-    $runspace.ThreadOptions  = "ReuseThread"
-    $runspace.Open()
-    $runspace.SessionStateProxy.SetVariable("sync",         $sync)
-    $runspace.SessionStateProxy.SetVariable("diskNum",      $diskNum)
-    $runspace.SessionStateProxy.SetVariable("contentsDir",  $contentsDir)
-
-    $script = [Management.Automation.PowerShell]::Create()
-    $script.Runspace = $runspace
-    $script.AddScript({
-
-        function Log($msg) {
-            $ts = (Get-Date).ToString("HH:mm:ss")
-            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
-                $sync["WPFWin11ISOStatusLog"].Text += "`n[$ts] $msg"
-                $sync["WPFWin11ISOStatusLog"].CaretIndex = $sync["WPFWin11ISOStatusLog"].Text.Length
-                $sync["WPFWin11ISOStatusLog"].ScrollToEnd()
             })
-        }
-        function SetProgress($label, $pct) {
-            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
-                $sync.progressBarTextBlock.Text    = $label
-                $sync.progressBarTextBlock.ToolTip = $label
-                $sync.ProgressBar.Value            = [Math]::Max($pct, 5)
-            })
-        }
-
-        try {
-            SetProgress "Formatting USB drive..." 10
-
-            # ── Diskpart script: clean, GPT, create ESP + data partitions ──
-            $dpScript = @"
-select disk $diskNum
-clean
-convert gpt
-create partition efi size=512
-format quick fs=fat32 label="SYSTEM"
-assign
-create partition primary
-format quick fs=fat32 label="WINPE"
-assign
-exit
-"@
-            $dpFile = Join-Path $env:TEMP "winutil_diskpart_$(Get-Random).txt"
-            $dpScript | Set-Content -Path $dpFile -Encoding ASCII
-            Log "Running diskpart on Disk $diskNum..."
-            diskpart /s $dpFile | Out-Null
-            Remove-Item $dpFile -Force
-
-            SetProgress "Identifying USB partitions..." 30
-            Start-Sleep -Seconds 3   # let Windows assign drive letters
-
-            # Find newly assigned drive letter for the data partition
-            $usbVol = Get-Partition -DiskNumber $diskNum |
-                      Where-Object { $_.Type -eq "Basic" } |
-                      Get-Volume |
-                      Where-Object { $_.FileSystemLabel -eq "WINPE" } |
-                      Select-Object -First 1
-
-            if (-not $usbVol) {
-                throw "Could not locate the formatted USB data partition.  Drive letter may not have been assigned automatically."
-            }
-
-            $usbDrive = "$($usbVol.DriveLetter):"
-            Log "USB data partition: $usbDrive"
-            SetProgress "Copying Windows 11 files to USB..." 45
-
-            # ── Copy files (split large install.wim if > 4 GB for FAT32) ──
-            $installWim = Join-Path $contentsDir "sources\install.wim"
-            if (Test-Path $installWim) {
-                $wimSizeMB = [math]::Round((Get-Item $installWim).Length / 1MB)
-                if ($wimSizeMB -gt 3800) {
-                    # FAT32 limit – split with DISM
-                    Log "install.wim is $wimSizeMB MB – splitting for FAT32 compatibility..."
-                    $splitDest = Join-Path $usbDrive "sources\install.swm"
-                    New-Item -ItemType Directory -Path (Split-Path $splitDest) -Force | Out-Null
-                    Split-WindowsImage -ImagePath $installWim `
-                                       -SplitImagePath $splitDest `
-                                       -FileSize 3800 -CheckIntegrity | Out-Null
-                    Log "install.wim split complete."
-
-                    # Copy everything else (exclude install.wim)
-                    $robocopyArgs = @($contentsDir, $usbDrive, "/E", "/XF", "install.wim", "/NFL", "/NDL", "/NJH", "/NJS")
-                    & robocopy @robocopyArgs | Out-Null
-                } else {
-                    & robocopy $contentsDir $usbDrive /E /NFL /NDL /NJH /NJS | Out-Null
-                }
-            } else {
-                & robocopy $contentsDir $usbDrive /E /NFL /NDL /NJH /NJS | Out-Null
-            }
-
-            SetProgress "Finalising USB drive..." 90
-            Log "Files copied to USB."
-
-            SetProgress "USB write complete ✔" 100
-            Log "USB drive is ready for use."
-
-            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
-                [System.Windows.MessageBox]::Show(
-                    "USB drive created successfully!`n`nYou can now boot from this drive to install Windows 11.",
-                    "USB Ready", "OK", "Info")
-            })
-        }
-        catch {
-            Log "ERROR during USB write: $_"
-            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
-                [System.Windows.MessageBox]::Show(
-                    "USB write failed:`n`n$_",
-                    "USB Write Error", "OK", "Error")
-            })
-        }
-        finally {
+        } finally {
             Start-Sleep -Milliseconds 800
             $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
                 $sync.progressBarTextBlock.Text    = ""
                 $sync.progressBarTextBlock.ToolTip = ""
                 $sync.ProgressBar.Value            = 0
-                $sync["WPFWin11ISOWriteUSBButton"].IsEnabled = $true
+                $sync["WPFWin11ISOChooseISOButton"].IsEnabled = $true
             })
         }
     }) | Out-Null

functions/private/Invoke-WinUtilISOScript.ps1

@@ -4,50 +4,38 @@ function Invoke-WinUtilISOScript {
         Applies WinUtil modifications to a mounted Windows 11 install.wim image.
 
     .DESCRIPTION
-        Performs the following operations against an already-mounted WIM image:
+        Removes AppX bloatware and OneDrive, optionally injects all drivers exported from
+        the running system into install.wim and boot.wim (controlled by the
+        -InjectCurrentSystemDrivers switch), applies offline registry tweaks (hardware
+        bypass, privacy, OOBE, telemetry, update suppression), deletes CEIP/WU
+        scheduled-task definition files, and optionally writes autounattend.xml to the ISO
+        root and removes the support\ folder from the ISO contents directory.
 
-        1. Removes provisioned AppX bloatware packages via DISM.
-        2. Removes OneDriveSetup.exe from the system image.
-        3. Loads offline registry hives (COMPONENTS, DEFAULT, NTUSER, SOFTWARE, SYSTEM)
-           and applies the following tweaks:
-             - Bypasses hardware requirement checks (CPU, RAM, SecureBoot, Storage, TPM).
-             - Disables sponsored-app delivery and ContentDeliveryManager features.
-             - Enables local-account OOBE path (BypassNRO).
-             - Writes autounattend.xml to the Sysprep directory inside the WIM and,
-               optionally, to the ISO/USB root so Windows Setup picks it up at boot.
-             - Disables reserved storage.
-             - Disables BitLocker device encryption.
-             - Hides the Chat (Teams) taskbar icon.
-             - Disables OneDrive folder backup (KFM).
-             - Disables telemetry, advertising ID, and input personalization.
-             - Blocks post-install delivery of DevHome, Outlook, and Teams.
-             - Disables Windows Copilot.
-             - Disables Windows Update during OOBE.
-        4. Deletes unwanted scheduled-task XML definition files (CEIP, Appraiser, etc.).
-        5. Removes the support\ folder from the ISO contents directory (if supplied).
+        All setup scripts embedded in the autounattend.xml <Extensions><File> nodes are
+        written directly into the WIM at their target paths under C:\Windows\Setup\Scripts\
+        to ensure they survive Windows Setup stripping unrecognised-namespace XML elements
+        from the Panther copy of the answer file.
 
-        Mounting and dismounting the WIM is the responsibility of the caller
-        (e.g. Invoke-WinUtilISO).
+        Mounting/dismounting the WIM is the caller's responsibility (e.g. Invoke-WinUtilISO).
 
     .PARAMETER ScratchDir
         Mandatory. Full path to the directory where the Windows image is currently mounted.
-        Example: C:\Users\USERNAME\AppData\Local\Temp\WinUtil_Win11ISO_20260222\wim_mount
 
     .PARAMETER ISOContentsDir
-        Optional. Root directory of the extracted ISO contents.
-        When supplied, autounattend.xml is also written here so Windows Setup picks it
-        up automatically at boot, and the support\ folder is deleted from that location.
+        Optional. Root directory of the extracted ISO contents. When supplied,
+        autounattend.xml is written here and the support\ folder is removed.
 
     .PARAMETER AutoUnattendXml
-        Optional. Full XML content for autounattend.xml.
-        In compiled winutil.ps1 this is the embedded $WinUtilAutounattendXml here-string;
-        in dev mode it is read from tools\autounattend.xml.
-        If empty, the OOBE bypass file is skipped and a warning is logged.
+        Optional. Full XML content for autounattend.xml. If empty, the OOBE bypass
+        file is skipped and a warning is logged.
+
+    .PARAMETER InjectCurrentSystemDrivers
+        Optional. When $true, exports all drivers from the running system and injects
+        them into install.wim and boot.wim index 2 (Windows Setup PE).
+        Defaults to $false.
 
     .PARAMETER Log
-        Optional ScriptBlock used for progress/status logging.
-        Receives a single [string] message argument.
-        Defaults to { param($m) Write-Output $m } when not supplied.
+        Optional ScriptBlock for progress/status logging. Receives a single [string] argument.
 
     .EXAMPLE
         Invoke-WinUtilISOScript -ScratchDir "C:\Temp\wim_mount"
@@ -62,24 +50,19 @@ function Invoke-WinUtilISOScript {
     .NOTES
         Author  : Chris Titus @christitustech
         GitHub  : https://github.com/ChrisTitusTech
-        Version : 26.02.22
+        Version : 26.03.02
     #>
     param (
         [Parameter(Mandatory)][string]$ScratchDir,
-        # Root directory of the extracted ISO contents. When supplied, autounattend.xml
-        # is written here so Windows Setup picks it up automatically at boot.
         [string]$ISOContentsDir = "",
-        # Autounattend XML content. In compiled winutil.ps1 this comes from the embedded
-        # $WinUtilAutounattendXml here-string; in dev mode it is read from tools\autounattend.xml.
         [string]$AutoUnattendXml = "",
+        [bool]$InjectCurrentSystemDrivers = $false,
         [scriptblock]$Log = { param($m) Write-Output $m }
     )
 
-    # ── Resolve admin group name (for takeown / icacls) ──────────────────────
     $adminSID   = New-Object System.Security.Principal.SecurityIdentifier('S-1-5-32-544')
     $adminGroup = $adminSID.Translate([System.Security.Principal.NTAccount])
 
-    # ── Local helpers ─────────────────────────────────────────────────────────
     function Set-ISOScriptReg {
         param ([string]$path, [string]$name, [string]$type, [string]$value)
         try {
@@ -100,15 +83,37 @@ function Invoke-WinUtilISOScript {
         }
     }
 
-    # ═════════════════════════════════════════════════════════════════════════
-    #  1. Remove provisioned AppX packages
-    # ═════════════════════════════════════════════════════════════════════════
+    function Add-DriversToImage {
+        param ([string]$MountPath, [string]$DriverDir, [string]$Label = "image", [scriptblock]$Logger)
+        & dism /English "/image:$MountPath" /Add-Driver "/Driver:$DriverDir" /Recurse 2>&1 |
+            ForEach-Object { & $Logger "  dism[$Label]: $_" }
+    }
+
+    function Invoke-BootWimInject {
+        param ([string]$BootWimPath, [string]$DriverDir, [scriptblock]$Logger)
+        Set-ItemProperty -Path $BootWimPath -Name IsReadOnly -Value $false -ErrorAction SilentlyContinue
+        $mountDir = Join-Path $env:TEMP "WinUtil_BootMount_$(Get-Random)"
+        New-Item -Path $mountDir -ItemType Directory -Force | Out-Null
+        try {
+            & $Logger "Mounting boot.wim (index 2) for driver injection..."
+            Mount-WindowsImage -ImagePath $BootWimPath -Index 2 -Path $mountDir -ErrorAction Stop | Out-Null
+            Add-DriversToImage -MountPath $mountDir -DriverDir $DriverDir -Label "boot" -Logger $Logger
+            & $Logger "Saving boot.wim..."
+            Dismount-WindowsImage -Path $mountDir -Save -ErrorAction Stop | Out-Null
+            & $Logger "boot.wim driver injection complete."
+        } catch {
+            & $Logger "Warning: boot.wim driver injection failed: $_"
+            try { Dismount-WindowsImage -Path $mountDir -Discard -ErrorAction SilentlyContinue | Out-Null } catch {}
+        } finally {
+            Remove-Item -Path $mountDir -Recurse -Force -ErrorAction SilentlyContinue
+        }
+    }
+
+    # ── 1. Remove provisioned AppX packages ──────────────────────────────────
     & $Log "Removing provisioned AppX packages..."
 
     $packages = & dism /English "/image:$ScratchDir" /Get-ProvisionedAppxPackages |
-        ForEach-Object {
-            if ($_ -match 'PackageName : (.*)') { $matches[1] }
-        }
+        ForEach-Object { if ($_ -match 'PackageName : (.*)') { $matches[1] } }
 
     $packagePrefixes = @(
         'AppUp.IntelManagementandSecurityStatus',
@@ -155,25 +160,46 @@ function Invoke-WinUtilISOScript {
         'MicrosoftTeams'
     )
 
-    $packagesToRemove = $packages | Where-Object {
-        $pkg = $_
-        $packagePrefixes | Where-Object { $pkg -like "*$_*" }
+    $packages | Where-Object { $pkg = $_; $packagePrefixes | Where-Object { $pkg -like "*$_*" } } |
+        ForEach-Object { & dism /English "/image:$ScratchDir" /Remove-ProvisionedAppxPackage "/PackageName:$_" }
+
+    # ── 2. Inject current system drivers (optional) ───────────────────────────
+    if ($InjectCurrentSystemDrivers) {
+        & $Log "Exporting all drivers from running system..."
+        $driverExportRoot = Join-Path $env:TEMP "WinUtil_DriverExport_$(Get-Random)"
+        New-Item -Path $driverExportRoot -ItemType Directory -Force | Out-Null
+        try {
+            Export-WindowsDriver -Online -Destination $driverExportRoot | Out-Null
+
+            & $Log "Injecting current system drivers into install.wim..."
+            Add-DriversToImage -MountPath $ScratchDir -DriverDir $driverExportRoot -Label "install" -Logger $Log
+            & $Log "install.wim driver injection complete."
+
+            if ($ISOContentsDir -and (Test-Path $ISOContentsDir)) {
+                $bootWim = Join-Path $ISOContentsDir "sources\boot.wim"
+                if (Test-Path $bootWim) {
+                    & $Log "Injecting current system drivers into boot.wim..."
+                    Invoke-BootWimInject -BootWimPath $bootWim -DriverDir $driverExportRoot -Logger $Log
+                } else {
+                    & $Log "Warning: boot.wim not found — skipping boot.wim driver injection."
                 }
-    foreach ($package in $packagesToRemove) {
-        & dism /English "/image:$ScratchDir" /Remove-ProvisionedAppxPackage "/PackageName:$package"
+            }
+        } catch {
+            & $Log "Error during driver export/injection: $_"
+        } finally {
+            Remove-Item -Path $driverExportRoot -Recurse -Force -ErrorAction SilentlyContinue
+        }
+    } else {
+        & $Log "Driver injection skipped."
     }
 
-    # ═════════════════════════════════════════════════════════════════════════
-    #  2. Remove OneDrive
-    # ═════════════════════════════════════════════════════════════════════════
+    # ── 3. Remove OneDrive ────────────────────────────────────────────────────
     & $Log "Removing OneDrive..."
     & takeown /f "$ScratchDir\Windows\System32\OneDriveSetup.exe" | Out-Null
     & icacls    "$ScratchDir\Windows\System32\OneDriveSetup.exe" /grant "$($adminGroup.Value):(F)" /T /C | Out-Null
     Remove-Item -Path "$ScratchDir\Windows\System32\OneDriveSetup.exe" -Force -ErrorAction SilentlyContinue
 
-    # ═════════════════════════════════════════════════════════════════════════
-    #  3. Registry tweaks
-    # ═════════════════════════════════════════════════════════════════════════
+    # ── 4. Registry tweaks ────────────────────────────────────────────────────
     & $Log "Loading offline registry hives..."
     reg load HKLM\zCOMPONENTS "$ScratchDir\Windows\System32\config\COMPONENTS"
     reg load HKLM\zDEFAULT    "$ScratchDir\Windows\System32\config\default"
@@ -222,14 +248,37 @@ function Invoke-WinUtilISOScript {
     Set-ISOScriptReg 'HKLM\zSOFTWARE\Microsoft\Windows\CurrentVersion\OOBE' 'BypassNRO' 'REG_DWORD' '1'
 
     if ($AutoUnattendXml) {
-        # ── Place autounattend.xml inside the WIM (Sysprep) ──────────────────
-        $sysprepDest = "$ScratchDir\Windows\System32\Sysprep\autounattend.xml"
-        Set-Content -Path $sysprepDest -Value $AutoUnattendXml -Encoding UTF8 -Force
-        & $Log "Written autounattend.xml to Sysprep directory."
+        try {
+            $xmlDoc = [xml]::new()
+            $xmlDoc.LoadXml($AutoUnattendXml)
+
+            $nsMgr = New-Object System.Xml.XmlNamespaceManager($xmlDoc.NameTable)
+            $nsMgr.AddNamespace("sg", "https://schneegans.de/windows/unattend-generator/")
+
+            $fileNodes = $xmlDoc.SelectNodes("//sg:File", $nsMgr)
+            if ($fileNodes -and $fileNodes.Count -gt 0) {
+                foreach ($fileNode in $fileNodes) {
+                    $absPath  = $fileNode.GetAttribute("path")
+                    $relPath  = $absPath -replace '^[A-Za-z]:[/\\]', ''
+                    $destPath = Join-Path $ScratchDir $relPath
+                    New-Item -Path (Split-Path $destPath -Parent) -ItemType Directory -Force -ErrorAction SilentlyContinue | Out-Null
+
+                    $ext = [IO.Path]::GetExtension($destPath).ToLower()
+                    $encoding = switch ($ext) {
+                        { $_ -in '.ps1', '.xml' }        { [System.Text.Encoding]::UTF8 }
+                        { $_ -in '.reg', '.vbs', '.js' } { [System.Text.UnicodeEncoding]::new($false, $true) }
+                        default                          { [System.Text.Encoding]::Default }
+                    }
+                    [System.IO.File]::WriteAllBytes($destPath, ($encoding.GetPreamble() + $encoding.GetBytes($fileNode.InnerText.Trim())))
+                    & $Log "Pre-staged setup script: $relPath"
+                }
+            } else {
+                & $Log "Warning: no <Extensions><File> nodes found in autounattend.xml — setup scripts not pre-staged."
+            }
+        } catch {
+            & $Log "Warning: could not pre-stage setup scripts from autounattend.xml: $_"
+        }
 
-        # ── Place autounattend.xml at the ISO / USB root ──────────────────────
-        # Windows Setup reads this file first (before booting into the OS),
-        # which is what drives the local-account / OOBE bypass at install time.
         if ($ISOContentsDir -and (Test-Path $ISOContentsDir)) {
             $isoDest = Join-Path $ISOContentsDir "autounattend.xml"
             Set-Content -Path $isoDest -Value $AutoUnattendXml -Encoding UTF8 -Force
@@ -304,12 +353,9 @@ function Invoke-WinUtilISOScript {
     reg unload HKLM\zSOFTWARE
     reg unload HKLM\zSYSTEM
 
-    # ═════════════════════════════════════════════════════════════════════════
-    #  4. Delete scheduled task definition files
-    # ═════════════════════════════════════════════════════════════════════════
+    # ── 5. Delete scheduled task definition files ─────────────────────────────
     & $Log "Deleting scheduled task definition files..."
     $tasksPath = "$ScratchDir\Windows\System32\Tasks"
-
     Remove-Item "$tasksPath\Microsoft\Windows\Application Experience\Microsoft Compatibility Appraiser" -Force -ErrorAction SilentlyContinue
     Remove-Item "$tasksPath\Microsoft\Windows\Customer Experience Improvement Program"                  -Recurse -Force -ErrorAction SilentlyContinue
     Remove-Item "$tasksPath\Microsoft\Windows\Application Experience\ProgramDataUpdater"               -Force -ErrorAction SilentlyContinue
@@ -321,12 +367,9 @@ function Invoke-WinUtilISOScript {
     Remove-Item "$tasksPath\Microsoft\Windows\WaaSMedic"                                               -Recurse -Force -ErrorAction SilentlyContinue
     Remove-Item "$tasksPath\Microsoft\Windows\WindowsUpdate"                                           -Recurse -Force -ErrorAction SilentlyContinue
     Remove-Item "$tasksPath\Microsoft\WindowsUpdate"                                                   -Recurse -Force -ErrorAction SilentlyContinue
-
     & $Log "Scheduled task files deleted."
 
-    # ═════════════════════════════════════════════════════════════════════════
-    #  5. Remove ISO support folder (fresh-install only; not needed)
-    # ═════════════════════════════════════════════════════════════════════════
+    # ── 6. Remove ISO support folder ─────────────────────────────────────────
     if ($ISOContentsDir -and (Test-Path $ISOContentsDir)) {
         & $Log "Removing ISO support\ folder..."
         Remove-Item -Path (Join-Path $ISOContentsDir "support") -Recurse -Force -ErrorAction SilentlyContinue

functions/private/Invoke-WinUtilISOUSB.ps1

@@ -0,0 +1,240 @@
+function Invoke-WinUtilISORefreshUSBDrives {
+    $combo    = $sync["WPFWin11ISOUSBDriveComboBox"]
+    $removable = @(Get-Disk | Where-Object { $_.BusType -eq "USB" } | Sort-Object Number)
+
+    $combo.Items.Clear()
+
+    if ($removable.Count -eq 0) {
+        $combo.Items.Add("No USB drives detected")
+        $combo.SelectedIndex = 0
+        $sync["Win11ISOUSBDisks"] = @()
+        Write-Win11ISOLog "No USB drives detected."
+        return
+    }
+
+    foreach ($disk in $removable) {
+        $sizeGB = [math]::Round($disk.Size / 1GB, 1)
+        $combo.Items.Add("Disk $($disk.Number): $($disk.FriendlyName)  [$sizeGB GB] - $($disk.PartitionStyle)")
+    }
+    $combo.SelectedIndex = 0
+    Write-Win11ISOLog "Found $($removable.Count) USB drive(s)."
+    $sync["Win11ISOUSBDisks"] = $removable
+}
+
+function Invoke-WinUtilISOWriteUSB {
+    $contentsDir = $sync["Win11ISOContentsDir"]
+    $usbDisks    = $sync["Win11ISOUSBDisks"]
+
+    if (-not $contentsDir -or -not (Test-Path $contentsDir)) {
+        [System.Windows.MessageBox]::Show("No modified ISO content found.  Please complete Steps 1-3 first.", "Not Ready", "OK", "Warning")
+        return
+    }
+
+    $combo = $sync["WPFWin11ISOUSBDriveComboBox"]
+    $selectedIndex = $combo.SelectedIndex
+    $selectedItemText = [string]$combo.SelectedItem
+    $usbDisks = @($usbDisks)
+
+    $targetDisk = $null
+    if ($selectedIndex -ge 0 -and $selectedIndex -lt $usbDisks.Count) {
+        $targetDisk = $usbDisks[$selectedIndex]
+    } elseif ($selectedItemText -match 'Disk\s+(\d+):') {
+        $selectedDiskNum = [int]$matches[1]
+        $targetDisk = $usbDisks | Where-Object { $_.Number -eq $selectedDiskNum } | Select-Object -First 1
+    }
+
+    if (-not $targetDisk) {
+        [System.Windows.MessageBox]::Show("Please select a USB drive from the dropdown.", "No Drive Selected", "OK", "Warning")
+        return
+    }
+
+    $diskNum    = $targetDisk.Number
+    $sizeGB     = [math]::Round($targetDisk.Size / 1GB, 1)
+
+    $confirm = [System.Windows.MessageBox]::Show(
+        "ALL data on Disk $diskNum ($($targetDisk.FriendlyName), $sizeGB GB) will be PERMANENTLY ERASED.`n`nAre you sure you want to continue?",
+        "Confirm USB Erase", "YesNo", "Warning")
+
+    if ($confirm -ne "Yes") {
+        Write-Win11ISOLog "USB write cancelled by user."
+        return
+    }
+
+    $sync["WPFWin11ISOWriteUSBButton"].IsEnabled = $false
+    Write-Win11ISOLog "Starting USB write to Disk $diskNum..."
+
+    $runspace = [Management.Automation.Runspaces.RunspaceFactory]::CreateRunspace()
+    $runspace.ApartmentState = "STA"
+    $runspace.ThreadOptions  = "ReuseThread"
+    $runspace.Open()
+    $runspace.SessionStateProxy.SetVariable("sync",        $sync)
+    $runspace.SessionStateProxy.SetVariable("diskNum",     $diskNum)
+    $runspace.SessionStateProxy.SetVariable("contentsDir", $contentsDir)
+
+    $script = [Management.Automation.PowerShell]::Create()
+    $script.Runspace = $runspace
+    $script.AddScript({
+
+        function Log($msg) {
+            $ts = (Get-Date).ToString("HH:mm:ss")
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                $sync["WPFWin11ISOStatusLog"].Text += "`n[$ts] $msg"
+                $sync["WPFWin11ISOStatusLog"].CaretIndex = $sync["WPFWin11ISOStatusLog"].Text.Length
+                $sync["WPFWin11ISOStatusLog"].ScrollToEnd()
+            })
+        }
+
+        function SetProgress($label, $pct) {
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                $sync.progressBarTextBlock.Text    = $label
+                $sync.progressBarTextBlock.ToolTip = $label
+                $sync.ProgressBar.Value            = [Math]::Max($pct, 5)
+            })
+        }
+
+        function Get-FreeDriveLetter {
+            $used = (Get-PSDrive -PSProvider FileSystem -ErrorAction SilentlyContinue).Name
+            foreach ($c in [char[]](68..90)) {
+                if ($used -notcontains [string]$c) { return $c }
+            }
+            return $null
+        }
+
+        try {
+            SetProgress "Formatting USB drive..." 10
+
+            # Phase 1: Clean disk via diskpart
+            $dpFile1 = Join-Path $env:TEMP "winutil_diskpart_$(Get-Random).txt"
+            "select disk $diskNum`nclean`nexit" | Set-Content -Path $dpFile1 -Encoding ASCII
+            Log "Running diskpart clean on Disk $diskNum..."
+            diskpart /s $dpFile1 2>&1 | Where-Object { $_ -match '\S' } | ForEach-Object { Log "  diskpart: $_" }
+            Remove-Item $dpFile1 -Force -ErrorAction SilentlyContinue
+
+            # Phase 2: Initialize as GPT
+            Start-Sleep -Seconds 2
+            Update-Disk -Number $diskNum -ErrorAction SilentlyContinue
+            $diskObj = Get-Disk -Number $diskNum -ErrorAction Stop
+            if ($diskObj.PartitionStyle -eq 'RAW') {
+                Initialize-Disk -Number $diskNum -PartitionStyle GPT -ErrorAction Stop
+                Log "Disk $diskNum initialized as GPT."
+            } else {
+                Set-Disk -Number $diskNum -PartitionStyle GPT -ErrorAction Stop
+                Log "Disk $diskNum converted to GPT (was $($diskObj.PartitionStyle))."
+            }
+
+            # Phase 3: Create FAT32 partition via diskpart
+            $volLabel = "W11-" + (Get-Date).ToString('yyMMdd')
+            $dpFile2  = Join-Path $env:TEMP "winutil_diskpart2_$(Get-Random).txt"
+            $maxFat32PartitionMB = 32768
+            $diskSizeMB = [int][Math]::Floor((Get-Disk -Number $diskNum -ErrorAction Stop).Size / 1MB)
+            $createPartitionCommand = "create partition primary"
+            if ($diskSizeMB -gt $maxFat32PartitionMB) {
+                $createPartitionCommand = "create partition primary size=$maxFat32PartitionMB"
+                Log "Disk $diskNum is $diskSizeMB MB; creating FAT32 partition capped at $maxFat32PartitionMB MB (32 GB)."
+            }
+
+            @(
+                "select disk $diskNum"
+                $createPartitionCommand
+                "format quick fs=fat32 label=`"$volLabel`""
+                "exit"
+            ) | Set-Content -Path $dpFile2 -Encoding ASCII
+            Log "Creating partitions on Disk $diskNum..."
+            diskpart /s $dpFile2 2>&1 | Where-Object { $_ -match '\S' } | ForEach-Object { Log "  diskpart: $_" }
+            Remove-Item $dpFile2 -Force -ErrorAction SilentlyContinue
+
+            SetProgress "Assigning drive letters..." 30
+            Start-Sleep -Seconds 3
+            Update-Disk -Number $diskNum -ErrorAction SilentlyContinue
+
+            $partitions = Get-Partition -DiskNumber $diskNum -ErrorAction Stop
+            Log "Partitions on Disk $diskNum after format: $($partitions.Count)"
+            foreach ($p in $partitions) {
+                Log "  Partition $($p.PartitionNumber)  Type=$($p.Type)  Letter=$($p.DriveLetter)  Size=$([math]::Round($p.Size/1MB))MB"
+            }
+
+            $winpePart = $partitions | Where-Object { $_.Type -eq "Basic" } | Select-Object -Last 1
+            if (-not $winpePart) {
+                throw "Could not find the WINPE (Basic) partition on Disk $diskNum after format."
+            }
+
+            try { Remove-PartitionAccessPath -DiskNumber $diskNum -PartitionNumber $winpePart.PartitionNumber -AccessPath "$($winpePart.DriveLetter):" -ErrorAction SilentlyContinue } catch {}
+            $usbLetter = Get-FreeDriveLetter
+            if (-not $usbLetter) { throw "No free drive letters (D-Z) available to assign to the USB data partition." }
+            Set-Partition -DiskNumber $diskNum -PartitionNumber $winpePart.PartitionNumber -NewDriveLetter $usbLetter
+            Log "Assigned drive letter $usbLetter to WINPE partition (Partition $($winpePart.PartitionNumber))."
+            Start-Sleep -Seconds 2
+
+            $usbDrive = "${usbLetter}:"
+            if (-not (Test-Path $usbDrive)) { throw "Drive $usbDrive is not accessible after letter assignment." }
+            Log "USB data partition: $usbDrive"
+
+            $contentSizeBytes = (Get-ChildItem -LiteralPath $contentsDir -File -Recurse -Force -ErrorAction Stop | Measure-Object -Property Length -Sum).Sum
+            if (-not $contentSizeBytes) { $contentSizeBytes = 0 }
+            $usbVolume = Get-Volume -DriveLetter $usbLetter -ErrorAction Stop
+            $partitionCapacityBytes = [int64]$usbVolume.Size
+            $partitionFreeBytes = [int64]$usbVolume.SizeRemaining
+
+            $contentSizeGB = [math]::Round($contentSizeBytes / 1GB, 2)
+            $partitionCapacityGB = [math]::Round($partitionCapacityBytes / 1GB, 2)
+            $partitionFreeGB = [math]::Round($partitionFreeBytes / 1GB, 2)
+
+            Log "Source content size: $contentSizeGB GB. USB partition capacity: $partitionCapacityGB GB, free: $partitionFreeGB GB."
+
+            if ($contentSizeBytes -gt $partitionCapacityBytes) {
+                throw "ISO content ($contentSizeGB GB) is larger than the USB partition capacity ($partitionCapacityGB GB). Use a larger USB drive or reduce image size."
+            }
+
+            if ($contentSizeBytes -gt $partitionFreeBytes) {
+                throw "Insufficient free space on USB partition. Required: $contentSizeGB GB, available: $partitionFreeGB GB."
+            }
+
+            SetProgress "Copying Windows 11 files to USB..." 45
+
+            # Copy files; split install.wim if > 4 GB (FAT32 limit)
+            $installWim = Join-Path $contentsDir "sources\install.wim"
+            if (Test-Path $installWim) {
+                $wimSizeMB = [math]::Round((Get-Item $installWim).Length / 1MB)
+                if ($wimSizeMB -gt 3800) {
+                    Log "install.wim is $wimSizeMB MB - splitting for FAT32 compatibility... This will take several minutes."
+                    $splitDest = Join-Path $usbDrive "sources\install.swm"
+                    New-Item -ItemType Directory -Path (Split-Path $splitDest) -Force | Out-Null
+                    Split-WindowsImage -ImagePath $installWim -SplitImagePath $splitDest -FileSize 3800 -CheckIntegrity
+                    Log "install.wim split complete."
+                    Log "Copying remaining files to USB..."
+                    & robocopy $contentsDir $usbDrive /E /XF install.wim /NFL /NDL /NJH /NJS
+                } else {
+                    & robocopy $contentsDir $usbDrive /E /NFL /NDL /NJH /NJS
+                }
+            } else {
+                & robocopy $contentsDir $usbDrive /E /NFL /NDL /NJH /NJS
+            }
+
+            SetProgress "Finalising USB drive..." 90
+            Log "Files copied to USB."
+            SetProgress "USB write complete" 100
+            Log "USB drive is ready for use."
+
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                [System.Windows.MessageBox]::Show(
+                    "USB drive created successfully!`n`nYou can now boot from this drive to install Windows 11.",
+                    "USB Ready", "OK", "Info")
+            })
+        } catch {
+            Log "ERROR during USB write: $_"
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                [System.Windows.MessageBox]::Show("USB write failed:`n`n$_", "USB Write Error", "OK", "Error")
+            })
+        } finally {
+            Start-Sleep -Milliseconds 800
+            $sync["WPFWin11ISOStatusLog"].Dispatcher.Invoke([action]{
+                $sync.progressBarTextBlock.Text    = ""
+                $sync.progressBarTextBlock.ToolTip = ""
+                $sync.ProgressBar.Value            = 0
+                $sync["WPFWin11ISOWriteUSBButton"].IsEnabled = $true
+            })
+        }
+    }) | Out-Null
+
+    $script.BeginInvoke() | Out-Null
+}

scripts/main.ps1

@@ -540,6 +540,10 @@ $sync["FontScalingApplyButton"].Add_Click({
 
 # ── Win11ISO Tab button handlers ──────────────────────────────────────────────
 
+$sync["WPFTab5BT"].Add_Click({
+    $sync["Form"].Dispatcher.BeginInvoke([System.Windows.Threading.DispatcherPriority]::Background, [action]{ Invoke-WinUtilISOCheckExistingWork }) | Out-Null
+})
+
 $sync["WPFWin11ISOBrowseButton"].Add_Click({
     Write-Debug "WPFWin11ISOBrowseButton clicked"
     Invoke-WinUtilISOBrowse

xaml/inputXML.xaml

@@ -273,22 +273,43 @@
         <Style TargetType="ComboBox">
             <Setter Property="Foreground" Value="{DynamicResource ComboBoxForegroundColor}" />
             <Setter Property="Background" Value="{DynamicResource ComboBoxBackgroundColor}" />
+            <Setter Property="MinWidth"   Value="{DynamicResource ButtonWidth}" />
             <Setter Property="Template">
                 <Setter.Value>
                     <ControlTemplate TargetType="ComboBox">
                         <Grid>
+                            <Border x:Name="OuterBorder"
+                                    BorderBrush="{DynamicResource BorderColor}"
+                                    BorderThickness="1"
+                                    CornerRadius="{DynamicResource ButtonCornerRadius}"
+                                    Background="{TemplateBinding Background}">
                                 <ToggleButton x:Name="ToggleButton"
-                                          Background="{TemplateBinding Background}"
-                                          BorderBrush="{TemplateBinding Background}"
+                                              Background="Transparent"
                                               BorderThickness="0"
                                               IsChecked="{Binding IsDropDownOpen, Mode=TwoWay, RelativeSource={RelativeSource TemplatedParent}}"
                                               ClickMode="Press">
-                                <TextBlock Text="{TemplateBinding SelectionBoxItem}"
+                                    <Grid>
+                                        <Grid.ColumnDefinitions>
+                                            <ColumnDefinition Width="*"/>
+                                            <ColumnDefinition Width="Auto"/>
+                                        </Grid.ColumnDefinitions>
+                                        <TextBlock Grid.Column="0"
+                                                   Text="{TemplateBinding SelectionBoxItem}"
                                                    Foreground="{TemplateBinding Foreground}"
                                                    Background="Transparent"
-                                           HorizontalAlignment="Center" VerticalAlignment="Center" Margin="2"
-                                           />
+                                                   HorizontalAlignment="Left" VerticalAlignment="Center"
+                                                   Margin="6,3,2,3"/>
+                                        <Path Grid.Column="1"
+                                              Data="M 0,0 L 8,0 L 4,5 Z"
+                                              Fill="{TemplateBinding Foreground}"
+                                              Width="8" Height="5"
+                                              VerticalAlignment="Center"
+                                              HorizontalAlignment="Center"
+                                              Stretch="Uniform"
+                                              Margin="4,0,6,0"/>
+                                    </Grid>
                                 </ToggleButton>
+                            </Border>
                             <Popup x:Name="Popup"
                                    IsOpen="{TemplateBinding IsDropDownOpen}"
                                    Placement="Bottom"
@@ -297,11 +318,11 @@
                                    PopupAnimation="Slide">
                                 <Border x:Name="DropDownBorder"
                                         Background="{TemplateBinding Background}"
-                                        BorderBrush="{TemplateBinding Foreground}"
+                                        BorderBrush="{DynamicResource BorderColor}"
                                         BorderThickness="1"
                                         CornerRadius="4">
                                     <ScrollViewer>
-                                        <ItemsPresenter HorizontalAlignment="Center" VerticalAlignment="Center" Margin="2"/>
+                                        <ItemsPresenter HorizontalAlignment="Left" VerticalAlignment="Center" Margin="4,2"/>
                                     </ScrollViewer>
                                 </Border>
                             </Popup>
@@ -1340,21 +1361,17 @@
                 </ScrollViewer>
             </TabItem>
             <TabItem Header="Win11ISO" Visibility="Collapsed" Name="WPFTab5">
-                <ScrollViewer VerticalScrollBarVisibility="Auto" HorizontalScrollBarVisibility="Disabled" Margin="{DynamicResource TabContentMargin}">
-                    <Grid Background="Transparent" Name="Win11ISOPanel">
+                <Grid Name="Win11ISOPanel" Margin="{DynamicResource TabContentMargin}" Background="Transparent">
                     <Grid.RowDefinitions>
-                            <RowDefinition Height="Auto"/>  <!-- Step 1: Select ISO -->
-                            <RowDefinition Height="Auto"/>  <!-- Step 2: Mount & Verify -->
-                            <RowDefinition Height="Auto"/>  <!-- Step 3: Modify install.wim -->
-                            <RowDefinition Height="Auto"/>  <!-- Step 4: Output Options -->
-                            <RowDefinition Height="Auto"/>  <!-- Log / Status -->
+                        <RowDefinition Height="Auto"/>  <!-- Steps 1-4 -->
+                        <RowDefinition Height="*"/>     <!-- Log / Status -->
                     </Grid.RowDefinitions>
 
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <!-- STEP 1 : Select Windows 11 ISO                             -->
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <Border Grid.Row="0" Name="WPFWin11ISOSelectSection" Style="{StaticResource BorderStyle}">
-                            <Grid Margin="5">
+                    <!-- Steps 1-4 -->
+                    <StackPanel Grid.Row="0">
+
+                            <!-- ─── STEP 1 : Select Windows 11 ISO ─────────────── -->
+                            <Grid Name="WPFWin11ISOSelectSection" Margin="5" HorizontalAlignment="Left" MinWidth="{DynamicResource ButtonWidth}">
                                 <Grid.ColumnDefinitions>
                                     <ColumnDefinition Width="*"/>
                                     <ColumnDefinition Width="*"/>
@@ -1441,16 +1458,12 @@
                                     </StackPanel>
                                 </Border>
                             </Grid>
-                        </Border>
 
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <!-- STEP 2 : Mount & Verify ISO                                -->
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <Border Grid.Row="1"
-                                Name="WPFWin11ISOMountSection"
-                                Style="{StaticResource BorderStyle}"
-                                Visibility="Collapsed">
-                            <Grid Margin="5">
+                            <!-- ─── STEP 2 : Mount & Verify ISO ──────────────────── -->
+                            <Grid Name="WPFWin11ISOMountSection"
+                                  Margin="5"
+                                  Visibility="Collapsed"
+                                  HorizontalAlignment="Left" MinWidth="{DynamicResource ButtonWidth}">
                                 <Grid.ColumnDefinitions>
                                     <ColumnDefinition Width="Auto"/>
                                     <ColumnDefinition Width="*"/>
@@ -1472,6 +1485,13 @@
                                             HorizontalAlignment="Left"
                                             Width="Auto" Padding="12,0"
                                             Height="{DynamicResource ButtonHeight}"/>
+                                    <CheckBox Name="WPFWin11ISOInjectDrivers"
+                                              Content="Inject current system drivers"
+                                              FontSize="{DynamicResource FontSize}"
+                                              Foreground="{DynamicResource MainForegroundColor}"
+                                              IsChecked="False"
+                                              Margin="0,8,0,0"
+                                              ToolTip="Exports all drivers from this machine and injects them into install.wim and boot.wim. Recommended for systems with unsupported NVMe or network controllers."/>
                                 </StackPanel>
 
                                 <!-- Verification results panel -->
@@ -1500,22 +1520,17 @@
                                                   FontSize="{DynamicResource FontSize}"
                                                   Foreground="{DynamicResource MainForegroundColor}"
                                                   Background="{DynamicResource MainBackgroundColor}"
-                                                  BorderBrush="{DynamicResource BorderColor}"
-                                                  HorizontalAlignment="Stretch"
+                                                  HorizontalAlignment="Left"
                                                   Margin="0,0,0,0"/>
                                     </StackPanel>
                                 </Border>
                             </Grid>
-                        </Border>
 
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <!-- STEP 3 : Modify install.wim                                -->
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <Border Grid.Row="2"
-                                Name="WPFWin11ISOModifySection"
-                                Style="{StaticResource BorderStyle}"
-                                Visibility="Collapsed">
-                            <StackPanel Margin="5">
+                            <!-- ─── STEP 3 : Modify install.wim ───────────────────── -->
+                            <StackPanel Name="WPFWin11ISOModifySection"
+                                        Margin="5"
+                                        Visibility="Collapsed"
+                                        HorizontalAlignment="Left" MinWidth="{DynamicResource ButtonWidth}">
                                 <TextBlock FontSize="{DynamicResource FontSize}" FontWeight="Bold"
                                            Foreground="{DynamicResource MainForegroundColor}" Margin="0,0,0,8">
                                     Step 3 - Modify install.wim
@@ -1534,15 +1549,12 @@
                                         Width="Auto" Padding="12,0"
                                         Height="{DynamicResource ButtonHeight}"/>
                             </StackPanel>
-                        </Border>
 
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <!-- STEP 4 : Output Options                                    -->
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <Border Grid.Row="3"
-                                Name="WPFWin11ISOOutputSection"
-                                Style="{StaticResource BorderStyle}">
-                            <StackPanel Margin="5">
+                            <!-- ─── STEP 4 : Output Options ───────────────────────── -->
+                            <StackPanel Name="WPFWin11ISOOutputSection"
+                                        Margin="5"
+                                        Visibility="Collapsed"
+                                        HorizontalAlignment="Left" MinWidth="{DynamicResource ButtonWidth}">
                                 <!-- Header row: title + Clean & Reset button -->
                                 <Grid Margin="0,0,0,12">
                                     <Grid.ColumnDefinitions>
@@ -1579,7 +1591,7 @@
                                             Height="{DynamicResource ButtonHeight}"/>
                                     <Button Grid.Column="2"
                                             Name="WPFWin11ISOChooseUSBButton"
-                                            Content="Write Directly to a USB Drive  (erases drive)"
+                                            Content="Write Directly to a USB Drive (ERASES DRIVE)"
                                             Foreground="OrangeRed"
                                             HorizontalAlignment="Stretch"
                                             Width="Auto" Padding="12,0"
@@ -1613,7 +1625,7 @@
                                                       Margin="0,0,6,0"/>
                                             <Button Grid.Column="1"
                                                     Name="WPFWin11ISORefreshUSBButton"
-                                                    Content="↻ Refresh"
+                                                    Content="Refresh"
                                                     Width="Auto" Padding="8,0"
                                                     Height="{DynamicResource ButtonHeight}"/>
                                         </Grid>
@@ -1626,34 +1638,37 @@
                                                 Margin="0,0,0,10"/>
                                     </StackPanel>
                                 </Border>
+                            </StackPanel>
 
                     </StackPanel>
-                        </Border>
 
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <!-- Status / Log Output                                        -->
-                        <!-- ═══════════════════════════════════════════════════════════ -->
-                        <Border Grid.Row="4" Style="{StaticResource BorderStyle}">
-                            <StackPanel>
-                                <TextBlock FontSize="{DynamicResource FontSize}" FontWeight="Bold"
-                                           Foreground="{DynamicResource MainForegroundColor}" Margin="0,0,0,6">
+                    <!-- Status Log (fills remaining height) -->
+                    <Grid Grid.Row="1" Margin="5">
+                        <Grid.RowDefinitions>
+                            <RowDefinition Height="Auto"/>
+                            <RowDefinition Height="*"/>
+                        </Grid.RowDefinitions>
+                        <TextBlock Grid.Row="0"
+                                   FontSize="{DynamicResource FontSize}" FontWeight="Bold"
+                                   Foreground="{DynamicResource MainForegroundColor}"
+                                   Margin="0,0,0,4">
                             Status Log
                         </TextBlock>
-                                <TextBox Name="WPFWin11ISOStatusLog"
+                        <TextBox Grid.Row="1"
+                                 Name="WPFWin11ISOStatusLog"
                                  IsReadOnly="True"
                                  TextWrapping="Wrap"
-                                         VerticalScrollBarVisibility="Auto"
-                                         Height="140" Padding="6"
+                                 VerticalScrollBarVisibility="Visible"
+                                 VerticalAlignment="Stretch"
+                                 Padding="6"
                                  Background="{DynamicResource MainBackgroundColor}"
                                  Foreground="{DynamicResource MainForegroundColor}"
                                  BorderBrush="{DynamicResource BorderColor}"
                                  BorderThickness="1"
                                  Text="Ready. Please select a Windows 11 ISO to begin."/>
-                            </StackPanel>
-                        </Border>
+                    </Grid>
 
                 </Grid>
-                </ScrollViewer>
             </TabItem>
         </TabControl>
     </Grid>