Refactor: Unify driver WIM and folder detection

Streamlines the driver discovery process by searching for both WIM files and driver folders at the same time. This replaces the previous sequential logic that checked for WIMs first, then folders.

If multiple driver sources are found, they are now presented in a single, unified list for user selection, regardless of whether they are a WIM or a folder. This simplifies the script's control flow and improves the user experience.

FFUDevelopment/WinPEDeployFFUFiles/ApplyFFU.ps1

@@ -410,78 +410,72 @@ else {
 
 #Find Drivers
 $DriversPath = $USBDrive + "Drivers"
-$WimToInstall = $null
+$DriverSourcePath = $null
-$DriverFolderToInstall = $null
+$DriverSourceType = $null # Will be 'WIM' or 'Folder'
+
 If (Test-Path -Path $DriversPath)
 {
-    WriteLog "Searching for driver WIMs in $DriversPath"
+    WriteLog "Searching for driver WIMs and folders in $DriversPath"
-    $WimFiles = Get-ChildItem -Path $DriversPath -Filter *.wim -Recurse
-    $WimFilesCount = $WimFiles.Count
     
-    if ($WimFilesCount -gt 0) {
+    # Get all WIM files
-        WriteLog "Found $WimFilesCount driver WIM file(s)."
+    $WimFiles = Get-ChildItem -Path $DriversPath -Filter *.wim -Recurse
-        if ($WimFilesCount -eq 1) {
+    
-            $WimToInstall = $WimFiles[0].FullName
+    # Get all top-level driver folders
-            WriteLog "Single driver WIM found, will install: $WimToInstall"
+    $DriverFolders = Get-ChildItem -Path $DriversPath -Directory
-        } else {
+
-            # Multiple WIMs found, prompt user
+    # Create a combined list
-            WriteLog "Multiple driver WIMs found. Prompting for selection."
+    $DriverSources = @()
-            $array = @()
+    $WimFiles | ForEach-Object {
-            for($i=0; $i -le $WimFilesCount -1; $i++){
+        $DriverSources += [PSCustomObject]@{
-                $Properties = [ordered]@{Number = $i + 1; WimFile = $WimFiles[$i].FullName}
+            Type = 'WIM'
-                $array += New-Object PSObject -Property $Properties
+            Path = $_.FullName
         }
-            $array | Format-Table -AutoSize -Property Number, WimFile
+    }
+    $DriverFolders | ForEach-Object {
+        $DriverSources += [PSCustomObject]@{
+            Type = 'Folder'
+            Path = $_.FullName
+        }
+    }
+
+    $DriverSourcesCount = $DriverSources.Count
+
+    if ($DriverSourcesCount -gt 0) {
+        WriteLog "Found $DriverSourcesCount total driver sources (WIMs and folders)."
+        if ($DriverSourcesCount -eq 1) {
+            $DriverSourcePath = $DriverSources[0].Path
+            $DriverSourceType = $DriverSources[0].Type
+            WriteLog "Single driver source found. Type: $DriverSourceType, Path: $DriverSourcePath"
+        } else {
+            # Multiple sources found, prompt user
+            WriteLog "Multiple driver sources found. Prompting for selection."
+            $displayArray = @()
+            for($i=0; $i -lt $DriverSourcesCount; $i++){
+                $displayArray += [PSCustomObject]@{
+                    Number = $i + 1
+                    Type   = $DriverSources[$i].Type
+                    Path   = $DriverSources[$i].Path
+                }
+            }
+            $displayArray | Format-Table -AutoSize
+            
             do {
                 try {
                     $var = $true
-                    [int]$WimSelected = Read-Host 'Enter the number of the driver WIM to install'
+                    [int]$DriverSelected = Read-Host 'Enter the number of the driver source to install'
-                    $WimSelected = $WimSelected - 1
+                    $DriverSelected = $DriverSelected - 1
                 } catch {
                     Write-Host 'Input was not in correct format. Please enter a valid number.'
                     $var = $false
                 }
-            } until (($WimSelected -ge 0) -and ($WimSelected -le ($WimFilesCount -1)) -and $var)
+            } until (($DriverSelected -ge 0) -and ($DriverSelected -lt $DriverSourcesCount) -and $var)
             
-            $WimToInstall = $array[$WimSelected].WimFile
+            $DriverSourcePath = $DriverSources[$DriverSelected].Path
-            WriteLog "$WimToInstall was selected."
+            $DriverSourceType = $DriverSources[$DriverSelected].Type
+            WriteLog "User selected Type: $DriverSourceType, Path: $DriverSourcePath"
         }
     } else {
-        WriteLog "No driver WIMs found. Searching for driver folders."
+        WriteLog "No driver WIMs or folders found in Drivers directory."
-        # Fallback to folder logic
-        $DriverFolders = Get-ChildItem -Path $DriversPath -Directory
-        $DriverFoldersCount = $DriverFolders.Count
-        If ($DriverFoldersCount -gt 1)
-        {
-            WriteLog "Found $DriverFoldersCount driver folders"
-            $array = @()
-            for($i=0; $i -le $DriverFoldersCount -1; $i++){
-                $Properties = [ordered]@{Number = $i + 1; Drivers = $DriverFolders[$i].FullName}
-                $array += New-Object PSObject -Property $Properties
-            }
-            $array | Format-Table -AutoSize -Property Number, Drivers
-            do {
-                try {
-                    $var = $true
-                    [int]$DriversSelected = Read-Host 'Enter the set of drivers to install'
-                    $DriversSelected = $DriversSelected - 1
-                } catch {
-                    Write-Host 'Input was not in correct format. Please enter a valid driver folder number'
-                    $var = $false
-                }
-            } until (($DriversSelected -ge 0) -and ($DriversSelected -le ($DriverFoldersCount -1)) -and $var) 
-
-            $DriverFolderToInstall = $array[$DriversSelected].Drivers
-            WriteLog "$DriverFolderToInstall was selected"
-        }
-        elseif ($DriverFoldersCount -eq 1) {
-            WriteLog "Found $DriverFoldersCount driver folder"
-            $DriverFolderToInstall = $DriverFolders[0].FullName
-            WriteLog "$DriverFolderToInstall will be installed"
-        } 
-        else {
-            Writelog 'No driver WIMs or folders found in Drivers directory.'
-        }
     }
 } else {
     WriteLog "Drivers folder not found at $DriversPath. Skipping driver installation."
@@ -613,8 +607,9 @@ If ($computername){
 }
 
 #Add Drivers
-if ($null -ne $WimToInstall) {
+if ($null -ne $DriverSourcePath) {
-    WriteLog "Installing drivers from WIM: $WimToInstall"
+    if ($DriverSourceType -eq 'WIM') {
+        WriteLog "Installing drivers from WIM: $DriverSourcePath"
         $TempDriverDir = "W:\TempDrivers"
         try {
             WriteLog "Creating temporary directory for drivers at $TempDriverDir"
@@ -622,7 +617,7 @@ if ($null -ne $WimToInstall) {
             
             WriteLog "Extracting WIM contents to $TempDriverDir"
             Write-Warning 'Extracting Drivers from WIM - dism will pop a window with no progress. This can take a few minutes to complete. Please be patient.'
-        Invoke-Process dism.exe "/Apply-Image /ImageFile:""$WimToInstall"" /Index:1 /ApplyDir:""$TempDriverDir"""
+            Invoke-Process dism.exe "/Apply-Image /ImageFile:""$DriverSourcePath"" /Index:1 /ApplyDir:""$TempDriverDir"""
             WriteLog "WIM extraction successful."
 
             WriteLog "Injecting drivers from $TempDriverDir"
@@ -642,21 +637,20 @@ if ($null -ne $WimToInstall) {
                 WriteLog "Cleanup successful."
             }
         }
-} elseif ($null -ne $DriverFolderToInstall) {
+    } elseif ($DriverSourceType -eq 'Folder') {
-    WriteLog "Injecting drivers from folder: $DriverFolderToInstall"
+        WriteLog "Injecting drivers from folder: $DriverSourcePath"
         Write-Warning 'Copying Drivers - dism will pop a window with no progress. This can take a few minutes to complete. This is done so drivers are logged to the scriptlog.txt file. Please be patient.'
-    Invoke-Process dism.exe "/image:W:\ /Add-Driver /Driver:""$DriverFolderToInstall"" /Recurse"
+        Invoke-Process dism.exe "/image:W:\ /Add-Driver /Driver:""$DriverSourcePath"" /Recurse"
         WriteLog "Driver injection from folder succeeded."
+    }
 } else {
     WriteLog "No drivers to install."
 }
 
 WriteLog "Setting Windows Boot Manager to be first in the firmware display order."
 Invoke-Process bcdedit.exe "/set {fwbootmgr} displayorder {bootmgr} /addfirst"
-WriteLog "Windows Boot Manager has been set to be first in the firmware display order."
 WriteLog "Setting Windows Boot Manager to be first in the default display order."
 Invoke-Process bcdedit.exe "/set {bootmgr} displayorder {default} /addfirst"
-WriteLog "Windows Boot Manager has been set to be first in the default display order."
 #Copy DISM log to USBDrive
 WriteLog "Copying dism log to $USBDrive"
 invoke-process xcopy "X:\Windows\logs\dism\dism.log $USBDrive /Y"