From 49a9fd49c157e549a5dbf7a87584f339799c71f4 Mon Sep 17 00:00:00 2001
From: Zehadi Alam <63765084+zehadialam@users.noreply.github.com>
Date: Mon, 20 May 2024 20:09:26 -0400
Subject: [PATCH] Add Optimize-FFUCaptureDrive function and disk cleanup to
 InstallAppsandSysprep.cmd file

---
 FFUDevelopment/Apps/InstallAppsandSysprep.cmd | 17 ++++++++++++
 FFUDevelopment/BuildFFUVM.ps1                 | 26 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/FFUDevelopment/Apps/InstallAppsandSysprep.cmd b/FFUDevelopment/Apps/InstallAppsandSysprep.cmd
index 1d2e97f..31cf62d 100644
--- a/FFUDevelopment/Apps/InstallAppsandSysprep.cmd
+++ b/FFUDevelopment/Apps/InstallAppsandSysprep.cmd
@@ -18,5 +18,22 @@ timeout /t 10
 REM Run Component Cleanup since dism /online /cleanup-image /analyzecomponentcleanup recommends it
 REM If adding latest CU, definitely need to do this to keep FFU size smaller
 dism /online /cleanup-image /startcomponentcleanup /resetbase
+REM Run disk cleanup (cleanmgr.exe) with all options enabled: https://learn.microsoft.com/en-us/troubleshoot/windows-server/backup-and-storage/automating-disk-cleanup-tool
+set rootkey=HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\VolumeCaches
+REM Per above doc, the Offline Pages Files subkey does not have stateflags value
+for /f "tokens=*" %%K in ('reg query "%rootkey%"') do (
+    echo %%K | findstr /i /c:"Offline Pages Files"
+    if errorlevel 1 (
+        reg add "%%K" /v StateFlags0000 /t REG_DWORD /d 2 /f
+    )
+)
+cleanmgr.exe /sagerun:0
+REM Remove the StateFlags0000 registry value
+for /f "tokens=*" %%K in ('reg query "%rootkey%"') do (
+    echo %%K | findstr /i /c:"Offline Pages Files"
+    if errorlevel 1 (
+        reg delete "%%K" /v StateFlags0000 /f
+    )
+)
 REM Sysprep/Generalize
 c:\windows\system32\sysprep\sysprep.exe /quiet /generalize /oobe
diff --git a/FFUDevelopment/BuildFFUVM.ps1 b/FFUDevelopment/BuildFFUVM.ps1
index 7a24450..aceaa03 100644
--- a/FFUDevelopment/BuildFFUVM.ps1
+++ b/FFUDevelopment/BuildFFUVM.ps1
@@ -1438,6 +1438,31 @@ function New-PEMedia {
     Remove-Item -Path "$WinPEFFUPath" -Recurse -Force
     WriteLog 'Cleanup complete'
 }
+
+function Optimize-FFUCaptureDrive {
+    param (
+        [string]$VhdxPath
+    )
+    try {
+        WriteLog 'Mounting VHDX for volume optimization'
+        Mount-VHD -Path $VhdxPath
+        WriteLog 'Defragmenting Windows partition...'
+        Optimize-Volume -DriveLetter W -Defrag -NormalPriority -Verbose
+        WriteLog 'Performing slab consolidation on Windows partition...'
+        Optimize-Volume -DriveLetter W -SlabConsolidate -NormalPriority -Verbose
+        WriteLog 'Dismounting VHDX'
+        Dismount-ScratchVhdx -VhdxPath $VhdxPath
+        WriteLog 'Mounting VHDX as read-only for optimization'
+        Mount-VHD -Path $VhdxPath -NoDriveLetter -ReadOnly
+        WriteLog 'Optimizing VHDX in full mode...'
+        Optimize-VHD -Path $VhdxPath -Mode Full
+        WriteLog 'Dismounting VHDX'
+        Dismount-ScratchVhdx -VhdxPath $VhdxPath
+    } catch {
+        throw $_
+    }
+}
+
 function New-FFU {
     param (
         [Parameter(Mandatory = $false)]
@@ -2364,6 +2389,7 @@ try {
             WriteLog 'Waiting for VM to shutdown'
         } while ($FFUVM.State -ne 'Off')
         WriteLog 'VM Shutdown'
+        Optimize-FFUCaptureDrive -VhdxPath $VHDXPath
         #Capture FFU file
         New-FFU $FFUVM.Name
     }