Feat: Automate driver selection during FFU deployment

Adds a `DriverMapping.json` file to automate driver injection during image deployment.

Driver download tasks now generate or update this mapping file with the relative path for each successfully downloaded driver package.

The deployment script now uses this file to automatically detect and select the correct drivers for the target hardware, removing the need for manual selection. The manual driver selection prompt is retained as a fallback.

FFUDevelopment/BuildFFUVM.ps1

@@ -3976,31 +3976,110 @@ if ($driversJsonPath -and (Test-Path $driversJsonPath) -and ($InstallDrivers -or
             -ListViewControl $null `
             -MainThreadLogPath $LogFile
 
-        # Log results from Invoke-ParallelProcessing
+        # After processing, update the driver mapping file
+        $successfullyDownloaded = [System.Collections.Generic.List[PSCustomObject]]::new()
         if ($null -ne $parallelResults) {
-            foreach ($result in $parallelResults) {
+            # Create a lookup table from the original items to get the 'Make'
-                if ($null -ne $result) {
+            $makeLookup = @{}
-                    # The $result here is the direct output from the Save-*DriversTask
+            $driversToProcess | ForEach-Object { $makeLookup[$_.Model] = $_.Make }
-                    # It should be a PSCustomObject with Identifier/Model, Status, Success
+
-                    $identifier = if ($result.PSObject.Properties.Name -contains 'Identifier') { $result.Identifier } elseif ($result.PSObject.Properties.Name -contains 'Model') { $result.Model } else { "UnknownItem" }
+            # Filter for objects that could be results, avoiding stray log strings
-                    $status = if ($result.PSObject.Properties.Name -contains 'Status') { $result.Status } else { "UnknownStatus" }
+            foreach ($result in ($parallelResults | Where-Object { $_ -is [hashtable] })) {
-                    $success = if ($result.PSObject.Properties.Name -contains 'Success') { $result.Success } else { $false }
+                if ($null -eq $result) { continue }
-                    
+
-                    $logMessage = "Driver task for '$identifier': Status: $status, Success: $success"
+                # The result from Invoke-ParallelProcessing is a hashtable.
-                    WriteLog $logMessage
+                # Access properties using their keys.
-                    if (-not $success) {
+                $modelName = $result['Identifier']
-                        Write-Warning $logMessage
+                $resultCode = $result['ResultCode']
+                $driverPath = $result['DriverPath']
+
+                if ([string]::IsNullOrWhiteSpace($modelName)) {
+                    WriteLog "Could not determine model name from result object: $($result | ConvertTo-Json -Compress -Depth 3)"
+                    continue
+                }
+
+                if ($resultCode -eq 0 -and -not [string]::IsNullOrWhiteSpace($driverPath)) {
+                    # The task was successful and returned a driver path.
+                    $make = $makeLookup[$modelName]
+                    if ($make) {
+                        $successfullyDownloaded.Add([PSCustomObject]@{
+                            Make       = $make
+                            Model      = $modelName
+                            DriverPath = $driverPath
+                        })
+                    }
+                    else {
+                        WriteLog "Warning: Could not find 'Make' for successful download of model '$modelName'. Skipping from DriverMapping.json."
                     }
                 }
                 else {
-                    WriteLog "A parallel driver task processed by Invoke-ParallelProcessing returned a null result."
+                    $logMessage = "Driver download failed or did not return a path for model '$modelName'. Status: $($result['Status'])"
+                    WriteLog $logMessage
+                    Write-Warning $logMessage
                 }
             }
         }
         else {
             WriteLog "Invoke-ParallelProcessing returned null or no results."
         }
+
+        # Update the driver mapping JSON if there are any successful downloads
+        if ($successfullyDownloaded.Count -gt 0) {
+            try {
+                WriteLog "Updating DriverMapping.json with $($successfullyDownloaded.Count) successfully downloaded drivers."
+                Update-DriverMappingJson -DownloadedDrivers $successfullyDownloaded -DriversFolder $DriversFolder
+            }
+            catch {
+                WriteLog "Warning: Failed to update DriverMapping.json: $($_.Exception.Message)"
+                # This is not a fatal error for the build process itself, so just show a warning.
+                Write-Warning "The driver download process completed, but failed to update the DriverMapping.json file. Please check the log for details."
+            }
+        }
         WriteLog "Finished processing drivers from $driversJsonPath."
+        
+        # After processing, update the driver mapping file
+        $successfullyDownloaded = [System.Collections.Generic.List[PSCustomObject]]::new()
+        if ($null -ne $parallelResults) {
+            # Create a lookup table from the original items to get the 'Make'
+            $makeLookup = @{}
+            $driversToProcess | ForEach-Object { $makeLookup[$_.Model] = $_.Make }
+        
+            foreach ($result in $parallelResults) {
+                if ($null -ne $result) {
+                    # Collect successful results for driver mapping
+                    if ($result.PSObject.Properties['Success'] -and $result.Success -and $result.PSObject.Properties['DriverPath'] -and -not [string]::IsNullOrWhiteSpace($result.DriverPath)) {
+                        $modelName = if ($result.PSObject.Properties.Name -contains 'Identifier') { $result.Identifier } else { $result.Model }
+                                
+                        # Find the 'Make' from the original list
+                        $make = $makeLookup[$modelName]
+        
+                        if ($make) {
+                            $successfullyDownloaded.Add([PSCustomObject]@{
+                                    Make       = $make
+                                    Model      = $modelName
+                                    DriverPath = $result.DriverPath
+                                })
+                        }
+                        else {
+                            WriteLog "Warning: Could not find 'Make' for successful download of model '$modelName'. Skipping from DriverMapping.json."
+                        }
+                    }
+                }
+            }
+        }
+        
+        # Update the driver mapping JSON if there are any successful downloads
+        if ($successfullyDownloaded.Count -gt 0) {
+            try {
+                WriteLog "Updating DriverMapping.json with $($successfullyDownloaded.Count) successfully downloaded drivers."
+                Update-DriverMappingJson -DownloadedDrivers $successfullyDownloaded -DriversFolder $DriversFolder
+            }
+            catch {
+                WriteLog "Warning: Failed to update DriverMapping.json: $($_.Exception.Message)"
+                # This is not a fatal error for the build process itself, so just show a warning.
+                Write-Warning "The driver download process completed, but failed to update the DriverMapping.json file. Please check the log for details."
+            }
+        }
     }
 }
 # Existing single-model driver download logic

FFUDevelopment/FFU.Common/FFU.Common.Drivers.psm1

@@ -80,8 +80,95 @@ function Compress-DriverFolderToWim {
     }
 }
 
+# --------------------------------------------------------------------------
+# SECTION: Driver Mapping Function
+# --------------------------------------------------------------------------
+
+function Update-DriverMappingJson {
+    [CmdletBinding()]
+    param(
+        [Parameter(Mandatory = $true)]
+        [array]$DownloadedDrivers, # Array of PSCustomObjects with Make, Model, DriverPath
+
+        [Parameter(Mandatory = $true)]
+        [string]$DriversFolder # Base drivers folder (e.g., C:\FFUDevelopment\Drivers)
+    )
+
+    $mappingFilePath = Join-Path -Path $DriversFolder -ChildPath "DriverMapping.json"
+    WriteLog "Updating driver mapping file at: $mappingFilePath"
+
+    # Load existing mapping file or create a new list
+    $mappingList = [System.Collections.Generic.List[PSCustomObject]]::new()
+    if (Test-Path -Path $mappingFilePath -PathType Leaf) {
+        try {
+            $existingJson = Get-Content -Path $mappingFilePath -Raw | ConvertFrom-Json
+            # Ensure it's a collection before adding to the list
+            if ($existingJson -is [array]) {
+                $mappingList.AddRange($existingJson)
+            }
+            else {
+                $mappingList.Add($existingJson)
+            }
+            WriteLog "Loaded $($mappingList.Count) existing entries from $mappingFilePath"
+        }
+        catch {
+            WriteLog "Warning: Could not read or parse existing DriverMapping.json. A new file will be created. Error: $($_.Exception.Message)"
+        }
+    }
+
+    $updatedCount = 0
+    $addedCount = 0
+
+    foreach ($driver in $DownloadedDrivers) {
+        # Skip if any required property is missing or null
+        if (-not $driver.PSObject.Properties['Make'] -or -not $driver.PSObject.Properties['Model'] -or -not $driver.PSObject.Properties['DriverPath'] -or [string]::IsNullOrWhiteSpace($driver.DriverPath)) {
+            WriteLog "Skipping driver entry due to missing or empty Make, Model, or DriverPath. Details: $(($driver | ConvertTo-Json -Compress -Depth 3))"
+            continue
+        }
+
+        # Find existing entry
+        $existingEntry = $mappingList | Where-Object { $_.Manufacturer -eq $driver.Make -and $_.Model -eq $driver.Model } | Select-Object -First 1
+
+        if ($null -ne $existingEntry) {
+            # Update existing entry if the path is different
+            if ($existingEntry.DriverPath -ne $driver.DriverPath) {
+                WriteLog "Updating driver path for '$($driver.Make) - $($driver.Model)' from '$($existingEntry.DriverPath)' to '$($driver.DriverPath)'."
+                $existingEntry.DriverPath = $driver.DriverPath
+                $updatedCount++
+            }
+        }
+        else {
+            # Add new entry
+            $newEntry = [PSCustomObject]@{
+                Manufacturer = $driver.Make
+                Model        = $driver.Model
+                DriverPath   = $driver.DriverPath
+            }
+            $mappingList.Add($newEntry)
+            WriteLog "Adding new mapping for '$($driver.Make) - $($driver.Model)' with path '$($driver.DriverPath)'."
+            $addedCount++
+        }
+    }
+
+    if ($updatedCount -gt 0 -or $addedCount -gt 0) {
+        try {
+            # Sort the list for consistency before saving
+            $sortedList = $mappingList | Sort-Object -Property Manufacturer, Model
+            $sortedList | ConvertTo-Json -Depth 5 | Set-Content -Path $mappingFilePath -Encoding UTF8
+            WriteLog "Successfully saved DriverMapping.json with $addedCount new entries and $updatedCount updated entries."
+        }
+        catch {
+            WriteLog "Error saving updated DriverMapping.json: $($_.Exception.Message)"
+            throw "Failed to save driver mapping file."
+        }
+    }
+    else {
+        WriteLog "No changes needed for DriverMapping.json."
+    }
+}
+
 # --------------------------------------------------------------------------
 # SECTION: Module Export
 # --------------------------------------------------------------------------
 
-Export-ModuleMember -Function Compress-DriverFolderToWim
+Export-ModuleMember -Function Compress-DriverFolderToWim, Update-DriverMappingJson

FFUDevelopment/FFU.Common/FFU.Common.Parallel.psm1

@@ -276,11 +276,17 @@ function Invoke-ParallelProcessing {
                 $localProgressQueue.Enqueue(@{ Identifier = $resultIdentifier; Status = $resultStatus })
             }
 
+            $driverPathValue = $null
+            if ($null -ne $taskResult -and $taskResult.PSObject.Properties.Name -contains 'DriverPath') {
+                $driverPathValue = $taskResult.DriverPath
+            }
+
             # Return a consistent hashtable structure (final result)
             return @{
                 Identifier = $resultIdentifier
                 Status     = $resultStatus # Return the final status
                 ResultCode = $resultCode
+                DriverPath = $driverPathValue
             }
 
         } -ThrottleLimit 5 -AsJob
@@ -383,10 +389,9 @@ function Invoke-ParallelProcessing {
                             $finalStatus = "$ErrorStatusPrefix Invalid Result Format"
                             $processedCount++ # Count as processed to avoid loop issues
                         }
-                        # Add the received result (even if format was unexpected, for logging)
+                            # Add the received result (even if format was unexpected, for logging)
-                        if ($null -ne $result) { $resultsCollection.Add($result) }
+                            if ($null -ne $result) { $resultsCollection.Add($result) }
-                        break # Only process first result from this job
+                        }
-                    }
                 }
                 else {
                     # Job completed but had no data

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.Dell.psm1

@@ -176,6 +176,7 @@ function Save-DellDriversTask {
     
     $makeDriversPath = Join-Path -Path $DriversFolder -ChildPath $Make
     $modelPath = Join-Path -Path $makeDriversPath -ChildPath $modelName
+    $driverRelativePath = Join-Path -Path $make -ChildPath $modelName # Relative path for the driver folder
 
     try {
         # Define paths for Dell catalog. The catalog is assumed to be prepared by the calling function.
@@ -190,7 +191,7 @@ function Save-DellDriversTask {
                 $status = "Already downloaded"
                 WriteLog "Drivers for '$modelName' already exist in '$modelPath'."
                 if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status $status }
-                return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $true }
+                return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $true; DriverPath = $driverRelativePath }
             }
             else {
                 WriteLog "Driver folder '$modelPath' for '$modelName' exists but is empty/small. Re-downloading."
@@ -635,6 +636,7 @@ function Save-DellDriversTask {
             if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status $status }
             $wimFileName = "$($modelName).wim"
             $destinationWimPath = Join-Path -Path $makeDriversPath -ChildPath $wimFileName
+            $driverRelativePath = Join-Path -Path $make -ChildPath $wimFileName # Update relative path to the WIM file
             WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
             try {
                 $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -ErrorAction Stop
@@ -667,14 +669,14 @@ function Save-DellDriversTask {
         # Enqueue the error status before returning
         if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status $status }
         # Ensure return object is created even on error
-        return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $success }
+        return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $success; DriverPath = $null }
     }
 
     # Enqueue the final status (success or error) before returning
     if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status $status }
 
     # Return the final status
-    return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $success }
+    return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $success; DriverPath = $driverRelativePath }
 }
 
 Export-ModuleMember -Function *

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.HP.psm1

@@ -115,6 +115,7 @@ function Save-HPDriversTask {
     $hpDriversBaseFolder = Join-Path -Path $DriversFolder -ChildPath $make # Changed variable name for clarity
     $platformListXml = Join-Path -Path $hpDriversBaseFolder -ChildPath "PlatformList.xml"
     $modelSpecificFolder = Join-Path -Path $hpDriversBaseFolder -ChildPath ($modelName -replace '[\\/:"*?<>|]', '_') # Sanitize model name for folder path
+    $driverRelativePath = Join-Path -Path $make -ChildPath ($modelName -replace '[\\/:"*?<>|]', '_') # Relative path for the driver folder
     $finalStatus = "" # Initialize final status
     $successState = $true # Assume success unless an operation fails
     
@@ -130,7 +131,7 @@ function Save-HPDriversTask {
             $errMsg = "Failed to create base HP driver folder '$hpDriversBaseFolder': $($_.Exception.Message)"
             WriteLog $errMsg
             if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status "Error: Create HP dir failed" }
-            return [PSCustomObject]@{ Identifier = $identifier; Status = "Error: Create HP dir failed"; Success = $false }
+            return [PSCustomObject]@{ Identifier = $identifier; Status = "Error: Create HP dir failed"; Success = $false; DriverPath = $null }
         }
     }
 
@@ -167,10 +168,14 @@ function Save-HPDriversTask {
             $finalStatus = "Already downloaded"
         }
         if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status $finalStatus }
+        if ($CompressToWim) {
+            $driverRelativePath = Join-Path -Path $make -ChildPath "$($identifier).wim"
+        }
         return [PSCustomObject]@{
             Identifier = $identifier
             Status     = $finalStatus
-            Success    = $successState 
+            Success    = $successState
+            DriverPath = $driverRelativePath
         }
     }
 
@@ -376,6 +381,7 @@ function Save-HPDriversTask {
                 Compress-DriverFolderToWim -SourceFolderPath $modelSpecificFolder -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -ErrorAction Stop
                 WriteLog "Compression successful for '$identifier'."
                 $finalStatus = "Completed & Compressed"
+                $driverRelativePath = Join-Path -Path $make -ChildPath "$($identifier).wim" # Update relative path to the WIM
             }
             catch {
                 WriteLog "Error during compression for '$identifier': $($_.Exception.Message)"
@@ -389,6 +395,7 @@ function Save-HPDriversTask {
         WriteLog $errorMessage
         $finalStatus = "Error: $($_.Exception.Message.Split([Environment]::NewLine)[0])"
         $successState = $false
+        $driverRelativePath = $null # Ensure path is null on error
         if (Test-Path -Path $modelSpecificFolder -PathType Container) {
             WriteLog "Attempting to remove partially created folder $modelSpecificFolder due to error."
             Remove-Item -Path $modelSpecificFolder -Recurse -Force -ErrorAction SilentlyContinue
@@ -396,7 +403,7 @@ function Save-HPDriversTask {
     }
             
     if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status $finalStatus }
-    return [PSCustomObject]@{ Identifier = $identifier; Status = $finalStatus; Success = $successState }
+    return [PSCustomObject]@{ Identifier = $identifier; Status = $finalStatus; Success = $successState; DriverPath = $driverRelativePath }
 }
 
 Export-ModuleMember -Function *

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.Lenovo.psm1

@@ -89,7 +89,8 @@ function Save-LenovoDriversTask {
     # Define paths
     $makeDriversPath = Join-Path -Path $DriversFolder -ChildPath $Make
     # Use the identifier (which contains the model name and machine type) and sanitize it for the path
-    $modelPath = Join-Path -Path $makeDriversPath -ChildPath ($identifier -replace '[\\/:"*?<>|]', '_') 
+    $modelPath = Join-Path -Path $makeDriversPath -ChildPath ($identifier -replace '[\\/:"*?<>|]', '_')
+    $driverRelativePath = Join-Path -Path $make -ChildPath ($identifier -replace '[\\/:"*?<>|]', '_') # Relative path for the driver folder
     $tempDownloadPath = Join-Path -Path $makeDriversPath -ChildPath "_TEMP_$($machineType)_$($PID)" # Temp folder for catalog/package XMLs
     
     if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status "Checking..." }
@@ -102,7 +103,7 @@ function Save-LenovoDriversTask {
                 $status = "Already downloaded"
                 WriteLog "Drivers for '$identifier' already exist in '$modelPath'."
                 if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status $status }
-                return [PSCustomObject]@{ Identifier = $identifier; Status = $status; Success = $true }
+                return [PSCustomObject]@{ Identifier = $identifier; Status = $status; Success = $true; DriverPath = $driverRelativePath }
             }
             else {
                 WriteLog "Driver folder '$modelPath' for '$identifier' exists but is empty/small. Re-downloading."
@@ -380,6 +381,7 @@ function Save-LenovoDriversTask {
             if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status $status }
             $wimFileName = "$($identifier).wim" # Use sanitized identifier for filename
             $destinationWimPath = Join-Path -Path $makeDriversPath -ChildPath $wimFileName
+            $driverRelativePath = Join-Path -Path $make -ChildPath $wimFileName # Update relative path to the WIM file
             WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
             try {
                 $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $identifier -WimDescription $identifier -ErrorAction Stop
@@ -412,7 +414,7 @@ function Save-LenovoDriversTask {
         # Enqueue the error status before returning
         if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status $status }
         # Ensure return object is created even on error
-        return [PSCustomObject]@{ Identifier = $identifier; Status = $status; Success = $success }
+        return [PSCustomObject]@{ Identifier = $identifier; Status = $status; Success = $success; DriverPath = $null }
     }
     finally {
         # Clean up the main catalog XML and temp folder
@@ -424,7 +426,7 @@ function Save-LenovoDriversTask {
     if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status $status }
 
     # Return the final status
-    return [PSCustomObject]@{ Identifier = $identifier; Status = $status; Success = $success }
+    return [PSCustomObject]@{ Identifier = $identifier; Status = $status; Success = $success; DriverPath = $driverRelativePath }
 }
 
 Export-ModuleMember -Function *

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.Microsoft.psm1

@@ -94,6 +94,7 @@ function Save-MicrosoftDriversTask {
     $modelName = $DriverItemData.Model
     $modelLink = $DriverItemData.Link
     $make = $DriverItemData.Make
+    $driverRelativePath = Join-Path -Path $make -ChildPath $modelName # Relative path for the driver folder
     $status = "Getting download link..." # Initial local status
     $success = $false
     
@@ -112,7 +113,7 @@ function Save-MicrosoftDriversTask {
                 # Enqueue this status before returning
                 if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status $status }
                 # Return success immediately
-                return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $true }
+                return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $true; DriverPath = $driverRelativePath }
             }
             else {
                 # Status is not set to error here, just log and continue
@@ -351,7 +352,8 @@ function Save-MicrosoftDriversTask {
                     if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status $status }
                     $wimFileName = "$($modelName).wim"
                     # Corrected WIM path: WIM file should be next to the model folder, not inside it.
-                    $destinationWimPath = Join-Path -Path $makeDriversPath -ChildPath $wimFileName 
+                    $destinationWimPath = Join-Path -Path $makeDriversPath -ChildPath $wimFileName
+                    $driverRelativePath = Join-Path -Path $make -ChildPath $wimFileName # Update relative path to the WIM file
                     WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
                     try {
                         # Use the function from the imported common module
@@ -398,14 +400,14 @@ function Save-MicrosoftDriversTask {
         # Enqueue the error status before returning
         if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status $status }
         # Ensure return object is created even on error
-        return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $success }
+        return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $success; DriverPath = $null }
     }
     
     # Enqueue the final status (success or error) before returning
     if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status $status }
     
     # Return the final status (this is still used by Receive-Job for final confirmation)
-    return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $success }
+    return [PSCustomObject]@{ Model = $modelName; Status = $status; Success = $success; DriverPath = $driverRelativePath }
 }
 
 Export-ModuleMember -Function *

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.psm1

@@ -624,7 +624,7 @@ function Invoke-DownloadSelectedDrivers {
         CompressToWim  = $compressDrivers
     }
 
-    Invoke-ParallelProcessing -ItemsToProcess $selectedDrivers `
+    $parallelResults = Invoke-ParallelProcessing -ItemsToProcess $selectedDrivers `
         -ListViewControl $State.Controls.lstDriverModels `
         -IdentifierProperty 'Model' `
         -StatusProperty 'DownloadStatus' `
@@ -636,14 +636,62 @@ function Invoke-DownloadSelectedDrivers {
         -MainThreadLogPath $State.LogFilePath
 
     $overallSuccess = $true
-    # Check if any item has an error status after processing
+    $successfullyDownloaded = [System.Collections.Generic.List[PSCustomObject]]::new()
-    # We iterate over $State.Controls.lstDriverModels.Items because their DownloadStatus property was updated by Invoke-ParallelProcessing
+
-    foreach ($item in ($State.Controls.lstDriverModels.Items | Where-Object { $_.IsSelected })) {
+    # Check the results from the parallel processing tasks
-        # Check only originally selected items
+    if ($null -ne $parallelResults) {
-        if ($item.DownloadStatus -like 'Error:*') {
+        # Create a lookup from the original selected drivers to get the 'Make' property,
-            $overallSuccess = $false
+        # as the result object might only have 'Identifier' or 'Model'.
-            WriteLog "Error detected for model $($item.Model) (Make: $($item.Make)): $($item.DownloadStatus)"
+        $makeLookup = @{}
-            # No break here, log all errors
+        $selectedDrivers | ForEach-Object { $makeLookup[$_.Model] = $_.Make }
+
+        # Filter for objects that could be results, avoiding stray log strings
+        foreach ($result in ($parallelResults | Where-Object { $_ -is [hashtable] })) {
+            if ($null -eq $result) { continue }
+
+            # The result from Invoke-ParallelProcessing is a hashtable.
+            # Access properties using their keys.
+            $modelName = $result['Identifier']
+            $resultCode = $result['ResultCode']
+            $driverPath = $result['DriverPath']
+
+            if ([string]::IsNullOrWhiteSpace($modelName)) {
+                WriteLog "Could not determine model name from result object: $($result | ConvertTo-Json -Compress -Depth 3)"
+                $overallSuccess = $false
+                continue
+            }
+
+            if ($resultCode -ne 0) {
+                $overallSuccess = $false
+                WriteLog "Error detected for model $modelName."
+            }
+            elseif (-not [string]::IsNullOrWhiteSpace($driverPath)) {
+                # The task was successful and returned a driver path.
+                $make = $makeLookup[$modelName]
+                if ($make) {
+                    $successfullyDownloaded.Add([PSCustomObject]@{
+                        Make       = $make
+                        Model      = $modelName
+                        DriverPath = $driverPath
+                    })
+                }
+                else {
+                    WriteLog "Warning: Could not find 'Make' for successful download of model '$modelName'. Skipping from DriverMapping.json."
+                }
+            }
+        }
+    }
+
+    # Update the driver mapping JSON if there are any successful downloads
+    if ($successfullyDownloaded.Count -gt 0) {
+        try {
+            WriteLog "Updating DriverMapping.json with $($successfullyDownloaded.Count) successfully downloaded drivers."
+            Update-DriverMappingJson -DownloadedDrivers $successfullyDownloaded -DriversFolder $localDriversFolder
+        }
+        catch {
+            WriteLog "Failed to update DriverMapping.json: $($_.Exception.Message)"
+            # This is not a fatal error for the download process itself, so just show a warning.
+            [System.Windows.MessageBox]::Show("The driver download process completed, but failed to update the DriverMapping.json file. Please check the log for details.", "Driver Mapping Error", "OK", "Warning")
         }
     }
 

FFUDevelopment/WinPEDeployFFUFiles/ApplyFFU.ps1

@@ -1,10 +1,10 @@
-function Get-USBDrive(){
+function Get-USBDrive() {
-    $USBDriveLetter = (Get-Volume | Where-Object {$_.DriveType -eq 'Removable' -and $_.FileSystemType -eq 'NTFS'}).DriveLetter
+    $USBDriveLetter = (Get-Volume | Where-Object { $_.DriveType -eq 'Removable' -and $_.FileSystemType -eq 'NTFS' }).DriveLetter
-    if ($null -eq $USBDriveLetter){
+    if ($null -eq $USBDriveLetter) {
         #Must be using a fixed USB drive - difficult to grab drive letter from win32_diskdrive. Assume user followed instructions and used Deploy as the friendly name for partition
-        $USBDriveLetter = (Get-Volume | Where-Object {$_.DriveType -eq 'Fixed' -and $_.FileSystemType -eq 'NTFS' -and $_.FileSystemLabel -eq 'Deploy'}).DriveLetter
+        $USBDriveLetter = (Get-Volume | Where-Object { $_.DriveType -eq 'Fixed' -and $_.FileSystemType -eq 'NTFS' -and $_.FileSystemLabel -eq 'Deploy' }).DriveLetter
         #If we didn't get the drive letter, stop the script.
-        if ($null -eq $USBDriveLetter){
+        if ($null -eq $USBDriveLetter) {
             WriteLog 'Cannot find USB drive letter - most likely using a fixed USB drive. Name the 2nd partition with the FFU files as Deploy so the script can grab the drive letter. Exiting'
             Exit
         }
@@ -14,46 +14,46 @@ function Get-USBDrive(){
     return $USBDriveLetter
 }
 
-function Get-HardDrive(){
+function Get-HardDrive() {
     $SystemInfo = Get-WmiObject -Class 'Win32_ComputerSystem'
     $Manufacturer = $SystemInfo.Manufacturer
     $Model = $SystemInfo.Model
     WriteLog "Device Manufacturer: $Manufacturer"
     WriteLog "Device Model: $Model"
     WriteLog 'Getting Hard Drive info'
-    if ($Manufacturer -eq 'Microsoft Corporation' -and $Model -eq 'Virtual Machine'){
+    if ($Manufacturer -eq 'Microsoft Corporation' -and $Model -eq 'Virtual Machine') {
         WriteLog 'Running in a Hyper-V VM. Getting virtual disk on Index 0 and SCSILogicalUnit 0'
-        $DiskDrive = Get-WmiObject -Class 'Win32_DiskDrive' | Where-Object {$_.MediaType -eq 'Fixed hard disk media' `
+        $DiskDrive = Get-WmiObject -Class 'Win32_DiskDrive' | Where-Object { $_.MediaType -eq 'Fixed hard disk media' `
-        -and $_.Model -eq 'Microsoft Virtual Disk' `
+                -and $_.Model -eq 'Microsoft Virtual Disk' `
-        -and $_.Index -eq 0 `
+                -and $_.Index -eq 0 `
-        -and $_.SCSILogicalUnit -eq 0
+                -and $_.SCSILogicalUnit -eq 0
         }
     }
-    else{
+    else {
         WriteLog 'Not running in a VM. Getting physical disk drive'
-        $DiskDrive = Get-WmiObject -Class 'Win32_DiskDrive' | Where-Object {$_.MediaType -eq 'Fixed hard disk media' -and $_.Model -ne 'Microsoft Virtual Disk'}
+        $DiskDrive = Get-WmiObject -Class 'Win32_DiskDrive' | Where-Object { $_.MediaType -eq 'Fixed hard disk media' -and $_.Model -ne 'Microsoft Virtual Disk' }
     }
     $DeviceID = $DiskDrive.DeviceID
     $BytesPerSector = $Diskdrive.BytesPerSector
 
     # Create a custom object to return both values
     $result = New-Object PSObject -Property @{
-        DeviceID = $DeviceID
+        DeviceID       = $DeviceID
         BytesPerSector = $BytesPerSector
     }
 
     return $result
 }
 
-function WriteLog($LogText){ 
+function WriteLog($LogText) { 
     Add-Content -path $LogFile -value "$((Get-Date).ToString()) $LogText"
 }
 
-function Set-DiskpartAnswerFiles($DiskpartFile,$DiskID){
+function Set-DiskpartAnswerFiles($DiskpartFile, $DiskID) {
     (Get-Content $DiskpartFile).Replace('disk 0', "disk $DiskID") | Set-Content -Path $DiskpartFile
 }
 
-function Set-Computername($computername){
+function Set-Computername($computername) {
     [xml]$xml = Get-Content $UnattendFile
     $components = $xml.unattend.settings.component
     $found = $false
@@ -73,60 +73,63 @@ function Set-Computername($computername){
 }
 
 function Invoke-Process {
-	[CmdletBinding(SupportsShouldProcess)]
+    [CmdletBinding(SupportsShouldProcess)]
-	param
+    param
-	(
+    (
-		[Parameter(Mandatory)]
+        [Parameter(Mandatory)]
-		[ValidateNotNullOrEmpty()]
+        [ValidateNotNullOrEmpty()]
-		[string]$FilePath,
+        [string]$FilePath,
 
-		[Parameter()]
+        [Parameter()]
-		[ValidateNotNullOrEmpty()]
+        [ValidateNotNullOrEmpty()]
-		[string]$ArgumentList
+        [string]$ArgumentList
-	)
+    )
 
-	$ErrorActionPreference = 'Stop'
+    $ErrorActionPreference = 'Stop'
 
-	try {
+    try {
-		$stdOutTempFile = "$env:TEMP\$((New-Guid).Guid)"
+        $stdOutTempFile = "$env:TEMP\$((New-Guid).Guid)"
-		$stdErrTempFile = "$env:TEMP\$((New-Guid).Guid)"
+        $stdErrTempFile = "$env:TEMP\$((New-Guid).Guid)"
 
-		$startProcessParams = @{
+        $startProcessParams = @{
-			FilePath               = $FilePath
+            FilePath               = $FilePath
-			ArgumentList           = $ArgumentList
+            ArgumentList           = $ArgumentList
-			RedirectStandardError  = $stdErrTempFile
+            RedirectStandardError  = $stdErrTempFile
-			RedirectStandardOutput = $stdOutTempFile
+            RedirectStandardOutput = $stdOutTempFile
-			Wait                   = $true;
+            Wait                   = $true;
-			PassThru               = $true;
+            PassThru               = $true;
-			NoNewWindow            = $false;
+            NoNewWindow            = $false;
-		}
+        }
-		if ($PSCmdlet.ShouldProcess("Process [$($FilePath)]", "Run with args: [$($ArgumentList)]")) {
+        if ($PSCmdlet.ShouldProcess("Process [$($FilePath)]", "Run with args: [$($ArgumentList)]")) {
-			$cmd = Start-Process @startProcessParams
+            $cmd = Start-Process @startProcessParams
-			$cmdOutput = Get-Content -Path $stdOutTempFile -Raw
+            $cmdOutput = Get-Content -Path $stdOutTempFile -Raw
-			$cmdError = Get-Content -Path $stdErrTempFile -Raw
+            $cmdError = Get-Content -Path $stdErrTempFile -Raw
-			if ($cmd.ExitCode -ne 0) {
+            if ($cmd.ExitCode -ne 0) {
-				if ($cmdError) {
+                if ($cmdError) {
-					throw $cmdError.Trim()
+                    throw $cmdError.Trim()
-				}
+                }
-				if ($cmdOutput) {
+                if ($cmdOutput) {
-					throw $cmdOutput.Trim()
+                    throw $cmdOutput.Trim()
-				}
+                }
-			} else {
+            }
-				if ([string]::IsNullOrEmpty($cmdOutput) -eq $false) {
+            else {
-					WriteLog $cmdOutput
+                if ([string]::IsNullOrEmpty($cmdOutput) -eq $false) {
-				}
+                    WriteLog $cmdOutput
-			}
+                }
-		}
+            }
-	} catch {
+        }
-		#$PSCmdlet.ThrowTerminatingError($_)
+    }
-		WriteLog $_
+    catch {
+        #$PSCmdlet.ThrowTerminatingError($_)
+        WriteLog $_
         Write-Host 'Script failed - check scriptlog.txt on the USB drive for more info'
-		throw $_
+        throw $_
 
-	} finally {
+    }
-		Remove-Item -Path $stdOutTempFile, $stdErrTempFile -Force -ErrorAction Ignore
+    finally {
+        Remove-Item -Path $stdOutTempFile, $stdErrTempFile -Force -ErrorAction Ignore
 		
-	}
+    }
 	
 }
 
@@ -142,7 +145,7 @@ WriteLog "Script version: $version"
 #Find PhysicalDrive
 # $PhysicalDeviceID = Get-HardDrive
 $hardDrive = Get-HardDrive
-if($null -eq $hardDrive){
+if ($null -eq $hardDrive) {
     WriteLog 'No hard drive found. Exiting'
     WriteLog 'Try adding storage drivers to the PE boot image (you can re-create your FFU and USB drive and add the PE drivers to the PEDrivers folder and add -CopyPEDrivers $true to the command line, or manually add them via DISM)'
     Exit
@@ -153,7 +156,7 @@ WriteLog "Physical BytesPerSector is $BytesPerSector"
 WriteLog "Physical DeviceID is $PhysicalDeviceID"
 
 #Parse DiskID Number
-$DiskID = $PhysicalDeviceID.substring($PhysicalDeviceID.length - 1,1)
+$DiskID = $PhysicalDeviceID.substring($PhysicalDeviceID.length - 1, 1)
 WriteLog "DiskID is $DiskID"
 
 #Find FFU Files
@@ -165,8 +168,8 @@ If ($FFUCount -gt 1) {
     WriteLog "Found $FFUCount FFU Files"
     $array = @()
 
-    for($i=0;$i -le $FFUCount -1;$i++){
+    for ($i = 0; $i -le $FFUCount - 1; $i++) {
-        $Properties = [ordered]@{Number = $i + 1 ; FFUFile = $FFUFiles[$i].FullName}
+        $Properties = [ordered]@{Number = $i + 1 ; FFUFile = $FFUFiles[$i].FullName }
         $array += New-Object PSObject -Property $Properties
     }
     $array | Format-Table -AutoSize -Property Number, FFUFile
@@ -174,14 +177,14 @@ If ($FFUCount -gt 1) {
         try {
             $var = $true
             [int]$FFUSelected = Read-Host 'Enter the FFU number to install'
-            $FFUSelected = $FFUSelected -1
+            $FFUSelected = $FFUSelected - 1
         }
 
         catch {
             Write-Host 'Input was not in correct format. Please enter a valid FFU number'
             $var = $false
         }
-    } until (($FFUSelected -le $FFUCount -1) -and $var) 
+    } until (($FFUSelected -le $FFUCount - 1) -and $var) 
 
     $FFUFileToInstall = $array[$FFUSelected].FFUFile
     WriteLog "$FFUFileToInstall was selected"
@@ -199,22 +202,22 @@ else {
 
 #FindAP
 $APFolder = $USBDrive + "Autopilot\"
-If (Test-Path -Path $APFolder){
+If (Test-Path -Path $APFolder) {
     [array]$APFiles = @(Get-ChildItem -Path $APFolder*.json)
     $APFilesCount = $APFiles.Count
-    if ($APFilesCount -ge 1){
+    if ($APFilesCount -ge 1) {
-    $autopilot = $true
+        $autopilot = $true
     }
 }
 
 
 #FindPPKG
 $PPKGFolder = $USBDrive + "PPKG\"
-if (Test-Path -Path $PPKGFolder){
+if (Test-Path -Path $PPKGFolder) {
     [array]$PPKGFiles = @(Get-ChildItem -Path $PPKGFolder*.ppkg)
     $PPKGFilesCount = $PPKGFiles.Count
-    if ($PPKGFilesCount -ge 1){
+    if ($PPKGFilesCount -ge 1) {
-    $PPKG = $true
+        $PPKG = $true
     }
 }
 
@@ -223,35 +226,35 @@ $UnattendFolder = $USBDrive + "unattend\"
 $UnattendFilePath = $UnattendFolder + "unattend.xml"
 $UnattendPrefixPath = $UnattendFolder + "prefixes.txt"
 $UnattendComputerNamePath = $UnattendFolder + "SerialComputerNames.csv"
-If (Test-Path -Path $UnattendFilePath){
+If (Test-Path -Path $UnattendFilePath) {
     $UnattendFile = Get-ChildItem -Path $UnattendFilePath
-    If ($UnattendFile){
+    If ($UnattendFile) {
         $Unattend = $true
     }
 }
-If (Test-Path -Path $UnattendPrefixPath){
+If (Test-Path -Path $UnattendPrefixPath) {
     $UnattendPrefixFile = Get-ChildItem -Path $UnattendPrefixPath
-    If ($UnattendPrefixFile){
+    If ($UnattendPrefixFile) {
         $UnattendPrefix = $true
     }
 }
-If (Test-Path -Path $UnattendComputerNamePath){
+If (Test-Path -Path $UnattendComputerNamePath) {
     $UnattendComputerNameFile = Get-ChildItem -Path $UnattendComputerNamePath
-    If ($UnattendComputerNameFile){
+    If ($UnattendComputerNameFile) {
         $UnattendComputerName = $true
     }
 }
 
 #Ask for device name if unattend exists
-if ($Unattend -and $UnattendPrefix){
+if ($Unattend -and $UnattendPrefix) {
     Writelog 'Unattend file found with prefixes.txt. Getting prefixes.'
     $UnattendPrefixes = @(Get-content $UnattendPrefixFile)
     $UnattendPrefixCount = $UnattendPrefixes.Count
     If ($UnattendPrefixCount -gt 1) {
         WriteLog "Found $UnattendPrefixCount Prefixes"
         $array = @()
-        for($i=0;$i -le $UnattendPrefixCount -1;$i++){
+        for ($i = 0; $i -le $UnattendPrefixCount - 1; $i++) {
-            $Properties = [ordered]@{Number = $i + 1 ; DeviceNamePrefix = $UnattendPrefixes[$i]}
+            $Properties = [ordered]@{Number = $i + 1 ; DeviceNamePrefix = $UnattendPrefixes[$i] }
             $array += New-Object PSObject -Property $Properties
         }
         $array | Format-Table -AutoSize -Property Number, DeviceNamePrefix
@@ -259,33 +262,33 @@ if ($Unattend -and $UnattendPrefix){
             try {
                 $var = $true
                 [int]$PrefixSelected = Read-Host 'Enter the prefix number to use for the device name'
-                $PrefixSelected = $PrefixSelected -1
+                $PrefixSelected = $PrefixSelected - 1
             }
             catch {
                 Write-Host 'Input was not in correct format. Please enter a valid prefix number'
                 $var = $false
             }
-        } until (($PrefixSelected -le $UnattendPrefixCount -1) -and $var) 
+        } until (($PrefixSelected -le $UnattendPrefixCount - 1) -and $var) 
         $PrefixToUse = $array[$PrefixSelected].DeviceNamePrefix
         WriteLog "$PrefixToUse was selected"
     }
     elseif ($UnattendPrefixCount -eq 1) {
-    WriteLog "Found $UnattendPrefixCount Prefix"
+        WriteLog "Found $UnattendPrefixCount Prefix"
-    $PrefixToUse = $UnattendPrefixes[0]
+        $PrefixToUse = $UnattendPrefixes[0]
-    WriteLog "Will use $PrefixToUse as device name prefix"
+        WriteLog "Will use $PrefixToUse as device name prefix"
     }
     #Get serial number to append. This can make names longer than 15 characters. Trim any leading or trailing whitespace
     $serial = (Get-CimInstance -ClassName win32_bios).SerialNumber.Trim()
     #Combine prefix with serial
-    $computername = ($PrefixToUse + $serial) -replace "\s","" # Remove spaces because windows does not support spaces in the computer names
+    $computername = ($PrefixToUse + $serial) -replace "\s", "" # Remove spaces because windows does not support spaces in the computer names
     #If computername is longer than 15 characters, reduce to 15. Sysprep/unattend doesn't like ComputerName being longer than 15 characters even though Windows accepts it
-    If ($computername.Length -gt 15){
+    If ($computername.Length -gt 15) {
-       $computername = $computername.substring(0,15)
+        $computername = $computername.substring(0, 15)
     }
     $computername = Set-Computername($computername)
     Writelog "Computer name set to $computername"
 }
-elseif($Unattend -and $UnattendComputerName){
+elseif ($Unattend -and $UnattendComputerName) {
     Writelog 'Unattend file found with SerialComputerNames.csv. Getting name for current computer.'
     $SerialComputerNames = Import-Csv -Path $UnattendComputerNameFile.FullName -Delimiter ","
 
@@ -296,14 +299,15 @@ elseif($Unattend -and $UnattendComputerName){
         [string]$computername = $SCName.ComputerName
         $computername = Set-Computername($computername)
         Writelog "Computer name set to $computername"
-    } else {
+    }
+    else {
         Writelog 'No matching serial number found in SerialComputerNames.csv. Setting random computer name to complete setup.'
-        [string]$computername = ("FFU-" + (-join ((48..57) + (65..90) + (97..122) | Get-Random -Count 11 | ForEach-Object { [char]$_ })))
+        [string]$computername = ("FFU-" + ( -join ((48..57) + (65..90) + (97..122) | Get-Random -Count 11 | ForEach-Object { [char]$_ })))
         $computername = Set-Computername($computername)
         Writelog "Computer name set to $computername"
     }
 }
-elseif($Unattend) {
+elseif ($Unattend) {
     Writelog 'Unattend file found with no prefixes.txt, asking for name'
     [string]$computername = Read-Host 'Enter device name'
     Set-Computername($computername)
@@ -314,7 +318,7 @@ else {
 }
 
 #If both AP and PPKG folder found with files, ask which to use.
-If($autopilot -eq $true -and $PPKG -eq $true){
+If ($autopilot -eq $true -and $PPKG -eq $true) {
     WriteLog 'Both PPKG and Autopilot json files found'
     Write-Host 'Both Autopilot JSON files and Provisioning packages were found.'
     do {
@@ -328,10 +332,10 @@ If($autopilot -eq $true -and $PPKG -eq $true){
             $var = $false
         }
     } until (($APorPPKG -gt 0 -and $APorPPKG -lt 3) -and $var)
-    If ($APorPPKG -eq 1){
+    If ($APorPPKG -eq 1) {
         $PPKG = $false
     }
-    else{
+    else {
         $autopilot = $false
     } 
 }
@@ -341,8 +345,8 @@ If ($APFilesCount -gt 1 -and $autopilot -eq $true) {
     WriteLog "Found $APFilesCount Autopilot json Files"
     $array = @()
 
-    for($i=0;$i -le $APFilesCount -1;$i++){
+    for ($i = 0; $i -le $APFilesCount - 1; $i++) {
-        $Properties = [ordered]@{Number = $i + 1 ; APFile = $APFiles[$i].FullName; APFileName = $APFiles[$i].Name}
+        $Properties = [ordered]@{Number = $i + 1 ; APFile = $APFiles[$i].FullName; APFileName = $APFiles[$i].Name }
         $array += New-Object PSObject -Property $Properties
     }
     $array | Format-Table -AutoSize -Property Number, APFileName
@@ -357,7 +361,7 @@ If ($APFilesCount -gt 1 -and $autopilot -eq $true) {
             Write-Host 'Input was not in correct format. Please enter a valid AP json file number'
             $var = $false
         }
-    } until (($APFileSelected -le $APFilesCount -1) -and $var) 
+    } until (($APFileSelected -le $APFilesCount - 1) -and $var) 
 
     $APFileToInstall = $array[$APFileSelected].APFile
     $APFileName = $array[$APFileSelected].APFileName
@@ -378,8 +382,8 @@ If ($PPKGFilesCount -gt 1 -and $PPKG -eq $true) {
     WriteLog "Found $PPKGFilesCount PPKG Files"
     $array = @()
 
-    for($i=0;$i -le $PPKGFilesCount -1;$i++){
+    for ($i = 0; $i -le $PPKGFilesCount - 1; $i++) {
-        $Properties = [ordered]@{Number = $i + 1 ; PPKGFile = $PPKGFiles[$i].FullName; PPKGFileName = $PPKGFiles[$i].Name}
+        $Properties = [ordered]@{Number = $i + 1 ; PPKGFile = $PPKGFiles[$i].FullName; PPKGFileName = $PPKGFiles[$i].Name }
         $array += New-Object PSObject -Property $Properties
     }
     $array | Format-Table -AutoSize -Property Number, PPKGFileName
@@ -394,7 +398,7 @@ If ($PPKGFilesCount -gt 1 -and $PPKG -eq $true) {
             Write-Host 'Input was not in correct format. Please enter a valid PPKG file number'
             $var = $false
         }
-    } until (($PPKGFileSelected -le $PPKGFilesCount -1) -and $var) 
+    } until (($PPKGFileSelected -le $PPKGFilesCount - 1) -and $var) 
 
     $PPKGFileToInstall = $array[$PPKGFileSelected].PPKGFile
     WriteLog "$PPKGFileToInstall was selected"
@@ -412,73 +416,139 @@ else {
 $DriversPath = $USBDrive + "Drivers"
 $DriverSourcePath = $null
 $DriverSourceType = $null # Will be 'WIM' or 'Folder'
+$driverMappingPath = Join-Path -Path $DriversPath -ChildPath "DriverMapping.json"
 
-If (Test-Path -Path $DriversPath)
+# --- Automatic Driver Detection using DriverMapping.json ---
-{
+if (Test-Path -Path $driverMappingPath -PathType Leaf) {
-    WriteLog "Searching for driver WIMs and folders in $DriversPath"
+    WriteLog "DriverMapping.json found at $driverMappingPath. Attempting automatic driver selection."
-    
+    try {
-    # Get all WIM files
+        # Get system information
-    $WimFiles = Get-ChildItem -Path $DriversPath -Filter *.wim -Recurse
+        $systemManufacturer = (Get-CimInstance -Class Win32_ComputerSystem).Manufacturer
-    
+        # Lenovo uses a different property for the model name
-    # Get all top-level driver folders
+        $systemModel = if ($systemManufacturer -like '*LENOVO*') {
-    $DriverFolders = Get-ChildItem -Path $DriversPath -Directory
+            (Get-CimInstance -Class Win32_ComputerSystemProduct).Version
-
-    # Create a combined list
-    $DriverSources = @()
-    $WimFiles | ForEach-Object {
-        $DriverSources += [PSCustomObject]@{
-            Type = 'WIM'
-            Path = $_.FullName
         }
-    }
+        else {
-    $DriverFolders | ForEach-Object {
+            (Get-CimInstance -Class Win32_ComputerSystem).Model
-        $DriverSources += [PSCustomObject]@{
-            Type = 'Folder'
-            Path = $_.FullName
         }
-    }
+        WriteLog "Detected System: Manufacturer='$systemManufacturer', Model='$systemModel'"
 
-    $DriverSourcesCount = $DriverSources.Count
+        # Load and parse the mapping file
+        $driverMappings = Get-Content -Path $driverMappingPath -Raw | ConvertFrom-Json
 
-    if ($DriverSourcesCount -gt 0) {
+        # Find a matching rule
-        WriteLog "Found $DriverSourcesCount total driver sources (WIMs and folders)."
+        $matchedRule = $null
-        if ($DriverSourcesCount -eq 1) {
+        foreach ($rule in $driverMappings) {
-            $DriverSourcePath = $DriverSources[0].Path
+            # Use -like for wildcard matching
-            $DriverSourceType = $DriverSources[0].Type
+            if ($systemManufacturer -like "$($rule.Manufacturer)*" -and $systemModel -like "$($rule.Model)*") {
-            WriteLog "Single driver source found. Type: $DriverSourceType, Path: $DriverSourcePath"
+                $matchedRule = $rule
-        } else {
+                break
-            # 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]$DriverSelected = Read-Host 'Enter the number of the driver source to install'
-                    $DriverSelected = $DriverSelected - 1
-                } catch {
-                    Write-Host 'Input was not in correct format. Please enter a valid number.'
-                    $var = $false
-                }
-            } until (($DriverSelected -ge 0) -and ($DriverSelected -lt $DriverSourcesCount) -and $var)
-            
-            $DriverSourcePath = $DriverSources[$DriverSelected].Path
-            $DriverSourceType = $DriverSources[$DriverSelected].Type
-            WriteLog "User selected Type: $DriverSourceType, Path: $DriverSourcePath"
         }
-    } else {
+
-        WriteLog "No driver WIMs or folders found in Drivers directory."
+        if ($null -ne $matchedRule) {
+            WriteLog "Automatic match found: Manufacturer='$($matchedRule.Manufacturer)', Model='$($matchedRule.Model)'"
+            $potentialDriverPath = Join-Path -Path $DriversPath -ChildPath $matchedRule.DriverPath
+            
+            if (Test-Path -Path $potentialDriverPath) {
+                $DriverSourcePath = $potentialDriverPath
+                # Determine if it's a WIM or a Folder
+                if ($DriverSourcePath -like '*.wim') {
+                    $DriverSourceType = 'WIM'
+                }
+                else {
+                    $DriverSourceType = 'Folder'
+                }
+                WriteLog "Automatically selected driver source. Type: $DriverSourceType, Path: $DriverSourcePath"
+            }
+            else {
+                WriteLog "Matched driver path '$potentialDriverPath' not found. Falling back to manual selection."
+            }
+        }
+        else {
+            WriteLog "No matching driver rule found in DriverMapping.json for this system. Falling back to manual selection."
+        }
+    }
+    catch {
+        WriteLog "An error occurred during automatic driver detection: $($_.Exception.Message). Falling back to manual selection."
+    }
+}
+else {
+    WriteLog "DriverMapping.json not found. Proceeding with manual driver selection."
+}
+
+# --- Manual Driver Selection (Fallback) ---
+if ($null -eq $DriverSourcePath) {
+    If (Test-Path -Path $DriversPath) {
+        WriteLog "Searching for driver WIMs and folders in $DriversPath"
+    
+        # Get all WIM files
+        $WimFiles = Get-ChildItem -Path $DriversPath -Filter *.wim -Recurse
+        
+        # Get all top-level driver folders
+        $DriverFolders = Get-ChildItem -Path $DriversPath -Directory
+
+        # Create a combined list
+        $DriverSources = @()
+        $WimFiles | ForEach-Object {
+            $DriverSources += [PSCustomObject]@{
+                Type = 'WIM'
+                Path = $_.FullName
+            }
+        }
+        $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]$DriverSelected = Read-Host 'Enter the number of the driver source to install'
+                        $DriverSelected = $DriverSelected - 1
+                    }
+                    catch {
+                        Write-Host 'Input was not in correct format. Please enter a valid number.'
+                        $var = $false
+                    }
+                } until (($DriverSelected -ge 0) -and ($DriverSelected -lt $DriverSourcesCount) -and $var)
+                
+                $DriverSourcePath = $DriverSources[$DriverSelected].Path
+                $DriverSourceType = $DriverSources[$DriverSelected].Type
+                WriteLog "User selected Type: $DriverSourceType, Path: $DriverSourcePath"
+            }
+        }
+        else {
+            WriteLog "No driver WIMs or folders found in Drivers directory."
+        }
+    }
+    else {
+        WriteLog "Drivers folder not found at $DriversPath. Skipping driver installation."
     }
-} else {
-    WriteLog "Drivers folder not found at $DriversPath. Skipping driver installation."
 }
 #Partition drive
 Writelog 'Clean Disk'
@@ -512,24 +582,24 @@ if ($recoveryPartition) {
     $diskpartScript | diskpart.exe | Out-Null
     WriteLog 'Setting recovery partition attributes complete'
 }
-if($LASTEXITCODE -eq 0){
+if ($LASTEXITCODE -eq 0) {
     WriteLog 'Successfully applied FFU'
 }
-elseif($LASTEXITCODE -eq 1393){
+elseif ($LASTEXITCODE -eq 1393) {
     WriteLog "Failed to apply FFU - LastExitCode = $LastExitCode"
     WriteLog "This is likely due to a mismatched LogicalSectorByteSize"
     WriteLog "BytesPerSector value from Win32_Diskdrive is $BytesPerSector"
-    if ($BytesPerSector -eq 4096){
+    if ($BytesPerSector -eq 4096) {
         WriteLog "The FFU build process by default uses a 512 LogicalSectorByteSize. Rebuild the FFU by adding -LogicalSectorByteSize 4096 to the command line"
     }
-    elseif($BytesPerSector -eq 512){
+    elseif ($BytesPerSector -eq 512) {
         WriteLog "This FFU was likely built with a LogicalSectorByteSize of 4096. Rebuild the FFU by adding -LogicalSectorByteSize 512 to the command line"
     }
     #Copy DISM log to USBDrive
     invoke-process xcopy.exe "X:\Windows\logs\dism\dism.log $USBDrive /Y"
     exit
 }
-else{
+else {
     Writelog "Failed to apply FFU - LastExitCode = $LASTEXITCODE also check dism.log on the USB drive for more info"
     #Copy DISM log to USBDrive
     invoke-process xcopy.exe "X:\Windows\logs\dism\dism.log $USBDrive /Y"
@@ -539,8 +609,7 @@ Get-Disk | Where-Object Number -eq $DiskID | Get-Partition | Where-Object Partit
 
 #Copy modified WinRE if folder exists, else copy inbox WinRE
 $WinRE = $USBDrive + "WinRE\winre.wim"
-If (Test-Path -Path $WinRE)
+If (Test-Path -Path $WinRE) {
-{
     WriteLog 'Copying modified WinRE to Recovery directory'
     Get-Disk | Where-Object Number -eq $DiskID | Get-Partition | Where-Object Type -eq Recovery | Set-Partition -NewDriveLetter R
     Invoke-Process xcopy.exe "/h $WinRE R:\Recovery\WindowsRE\ /Y"
@@ -551,7 +620,7 @@ If (Test-Path -Path $WinRE)
     WriteLog 'Registering location of recovery tools succeeded'
 }
 #Autopilot JSON
-If ($APFileToInstall){
+If ($APFileToInstall) {
     WriteLog "Copying $APFileToInstall to W:\windows\provisioning\autopilot"
     Invoke-process xcopy.exe "$APFileToInstall W:\Windows\provisioning\autopilot\"
     WriteLog "Copying $APFileToInstall to W:\windows\provisioning\autopilot succeeded"
@@ -561,13 +630,13 @@ If ($APFileToInstall){
         WriteLog "Renamed W:\Windows\Provisioning\Autopilot\$APFilename to W:\Windows\Provisioning\Autopilot\AutoPilotConfigurationFile.json"
     }
     
-    catch{
+    catch {
         Writelog "Copying $APFileToInstall to W:\windows\provisioning\autopilot failed with error: $_"
         throw $_
     }
 }
 #Apply PPKG
-If ($PPKGFileToInstall){
+If ($PPKGFileToInstall) {
     try {
         #Make sure to delete any existing PPKG on the USB drive
         Get-Childitem -Path $USBDrive\*.ppkg | ForEach-Object {
@@ -578,21 +647,21 @@ If ($PPKGFileToInstall){
         WriteLog "Copying $PPKGFileToInstall to $USBDrive succeeded"
     }
 
-    catch{
+    catch {
         Writelog "Copying $PPKGFileToInstall to $USBDrive failed with error: $_"
         throw $_
     }
 }
 #Set DeviceName
-If ($computername){
+If ($computername) {
-    try{
+    try {
         $PantherDir = 'w:\windows\panther'
-        If (Test-Path -Path $PantherDir){
+        If (Test-Path -Path $PantherDir) {
             Writelog "Copying $UnattendFile to $PantherDir"
             Invoke-process xcopy "$UnattendFile $PantherDir /Y"
             WriteLog "Copying $UnattendFile to $PantherDir succeeded"
         }
-        else{
+        else {
             Writelog "$PantherDir doesn't exist, creating it"
             New-Item -Path $PantherDir -ItemType Directory -Force
             Writelog "Copying $UnattendFile to $PantherDir"
@@ -600,7 +669,7 @@ If ($computername){
             WriteLog "Copying $UnattendFile to $PantherDir succeeded"
         }
     }
-    catch{
+    catch {
         WriteLog "Copying Unattend.xml to name device failed"
         throw $_
     }   
@@ -625,25 +694,29 @@ if ($null -ne $DriverSourcePath) {
             Invoke-Process dism.exe "/image:W:\ /Add-Driver /Driver:""$TempDriverDir"" /Recurse"
             WriteLog "Driver injection from WIM succeeded."
 
-        } catch {
+        }
+        catch {
             WriteLog "An error occurred during WIM driver installation: $_"
             # Copy DISM log to USBDrive for debugging
             invoke-process xcopy.exe "X:\Windows\logs\dism\dism.log $USBDrive /Y"
             throw $_
-        } finally {
+        }
+        finally {
             if (Test-Path -Path $TempDriverDir) {
                 WriteLog "Cleaning up temporary driver directory: $TempDriverDir"
                 Remove-Item -Path $TempDriverDir -Recurse -Force
                 WriteLog "Cleanup successful."
             }
         }
-    } elseif ($DriverSourceType -eq 'Folder') {
+    }
+    elseif ($DriverSourceType -eq 'Folder') {
         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:""$DriverSourcePath"" /Recurse"
         WriteLog "Driver injection from folder succeeded."
     }
-} else {
+}
+else {
     WriteLog "No drivers to install."
 }