From 01c070f3a2e824d55a86efc9adc579327f4de139 Mon Sep 17 00:00:00 2001
From: CodingWonders <101426328+CodingWonders@users.noreply.github.com>
Date: Thu, 25 Apr 2024 16:03:12 +0200
Subject: [PATCH] Update Invoke-WPFMicrowin.ps1

Use the driver injection directory if it is specified and if the directory exists
---
 functions/public/Invoke-WPFMicrowin.ps1 | 65 ++++++++++++++++---------
 1 file changed, 41 insertions(+), 24 deletions(-)

diff --git a/functions/public/Invoke-WPFMicrowin.ps1 b/functions/public/Invoke-WPFMicrowin.ps1
index 58d0f35f..564a836a 100644
--- a/functions/public/Invoke-WPFMicrowin.ps1
+++ b/functions/public/Invoke-WPFMicrowin.ps1
@@ -112,6 +112,47 @@ public class PowerManagement {
             return
         }
 
+		if ($importDrivers)
+		{
+			Write-Host "Exporting drivers from active installation..."
+			if (Test-Path "$env:TEMP\DRV_EXPORT")
+			{
+				Remove-Item "$env:TEMP\DRV_EXPORT" -Recurse -Force
+			}
+			if (($injectDrivers -and (Test-Path $sync.MicrowinDriverLocation.Text)))
+			{
+				Write-Host "Using specified driver source..."
+				dism /english /online /export-driver /destination="$($sync.MicrowinDriverLocation.Text)" | Out-Host
+				if ($?)
+				{
+					# Don't add exported drivers yet, that is run later
+					Write-Host "Drivers have been exported successfully."
+				}
+				else
+				{
+					Write-Host "Failed to export drivers."
+				}
+			}
+			else
+			{
+				New-Item -Path "$env:TEMP\DRV_EXPORT" -ItemType Directory -Force
+				dism /english /online /export-driver /destination="$env:TEMP\DRV_EXPORT" | Out-Host
+				if ($?)
+				{
+					Write-Host "Adding exported drivers..."
+					dism /english /image="$scratchDir" /add-driver /driver="$env:TEMP\DRV_EXPORT" /recurse | Out-Host
+				}
+				else
+				{
+					Write-Host "Failed to export drivers. Continuing without importing them..."
+				}
+				if (Test-Path "$env:TEMP\DRV_EXPORT")
+				{
+					Remove-Item "$env:TEMP\DRV_EXPORT" -Recurse -Force
+				}				
+			}
+		}
+
 		if ($injectDrivers)
 		{
 			$driverPath = $sync.MicrowinDriverLocation.Text
@@ -126,30 +167,6 @@ public class PowerManagement {
 			}
 		}
 
-		if ($importDrivers)
-		{
-			Write-Host "Exporting drivers from active installation..."
-			if (Test-Path "$env:TEMP\DRV_EXPORT")
-			{
-				Remove-Item "$env:TEMP\DRV_EXPORT" -Recurse -Force
-			}
-			New-Item -Path "$env:TEMP\DRV_EXPORT" -ItemType Directory -Force
-			dism /english /online /export-driver /destination="$env:TEMP\DRV_EXPORT" | Out-Host
-			if ($?)
-			{
-				Write-Host "Adding exported drivers..."
-				dism /english /image="$scratchDir" /add-driver /driver="$env:TEMP\DRV_EXPORT" /recurse | Out-Host
-			}
-			else
-			{
-				Write-Host "Failed to export drivers. Continuing without importing them..."
-			}
-			if (Test-Path "$env:TEMP\DRV_EXPORT")
-			{
-				Remove-Item "$env:TEMP\DRV_EXPORT" -Recurse -Force
-			}
-		}
-
 		Write-Host "Remove Features from the image"
 		Remove-Features -keepDefender:$keepDefender
 		Write-Host "Removing features complete!"