From d9682ef4c8318965a1eb5c2bc813132eafb09df5 Mon Sep 17 00:00:00 2001
From: rad <80015713+RadNotRed@users.noreply.github.com>
Date: Mon, 17 Feb 2025 18:37:31 -0800
Subject: [PATCH] feat: GUI & Rewrite

---
 .../2. Drivers/Run Update Drivers.cmd         |  45 +++----
 .../Scripts/ScriptWrappers/UpdateDrivers.ps1  | 110 +++++++++++++++---
 2 files changed, 116 insertions(+), 39 deletions(-)

diff --git a/src/playbook/Executables/AtlasDesktop/2. Drivers/Run Update Drivers.cmd b/src/playbook/Executables/AtlasDesktop/2. Drivers/Run Update Drivers.cmd
index 3ad12bf6f2..e189badc02 100644
--- a/src/playbook/Executables/AtlasDesktop/2. Drivers/Run Update Drivers.cmd	
+++ b/src/playbook/Executables/AtlasDesktop/2. Drivers/Run Update Drivers.cmd	
@@ -1,20 +1,25 @@
-@echo off
-set "script=%windir%\AtlasModules\Scripts\ScriptWrappers\UpdateDrivers.ps1"
-
-if not exist "%script%" (
-    echo Script not found: "%script%"
-    pause
-    exit /b 1
-)
-
-whoami /user | find /i "S-1-5-18" > nul 2>&1 || (
-    call RunAsTI.cmd "%~f0" %*
-    exit /b
-)
-
-@REM The "echo Y" is to bypass the NuGet installation user prompt.
-@REM If there's a better way, feel free to PR.
-echo Y | powershell -ExecutionPolicy Bypass -NoProfile -File "%script%"
-
-pause
-exit /b 0
+@echo off
+set "script=C:\Users\rad\Downloads\UpdateDrivers.ps1"
+
+set "___args="%~f0" %*"
+fltmc > nul 2>&1 || (
+    echo Administrator privileges are required.
+    powershell -c "Start-Process -Verb RunAs -FilePath 'cmd' -ArgumentList """/c $env:___args"""" 2> nul || (
+        echo You must run this script as admin.
+        if "%*"=="" pause
+        exit /b 1
+    )
+    exit /b
+)
+
+if not exist "%script%" (
+    echo Script not found: "%script%"
+    pause
+    exit /b 1
+)
+
+powershell -ExecutionPolicy Bypass -NoProfile -File "%script%"
+
+echo.
+pause > null
+exit /b 0
diff --git a/src/playbook/Executables/AtlasModules/Scripts/ScriptWrappers/UpdateDrivers.ps1 b/src/playbook/Executables/AtlasModules/Scripts/ScriptWrappers/UpdateDrivers.ps1
index 0bb1fd1e0f..cbfb0c0c6e 100644
--- a/src/playbook/Executables/AtlasModules/Scripts/ScriptWrappers/UpdateDrivers.ps1
+++ b/src/playbook/Executables/AtlasModules/Scripts/ScriptWrappers/UpdateDrivers.ps1
@@ -3,18 +3,26 @@ param (
     [switch]$RestartAfterUpdate
 )
 
-function Install-PSWindowsUpdateModule {
-    # Make sure that PowerShellGet is updated just in case 
-    if (-not (Get-Module -ListAvailable -Name PowerShellGet)) {
-        Install-Module -Name PowerShellGet -Force -AllowClobber -Confirm:$false -Scope AllUsers
+$script:SelectedUpdates = @()
+
+function Test-Admin {
+    $currentUser = [System.Security.Principal.WindowsIdentity]::GetCurrent()
+    $principal   = New-Object System.Security.Principal.WindowsPrincipal($currentUser)
+    if (-not $principal.IsInRole([System.Security.Principal.WindowsBuiltInRole]::Administrator)) {
+        Write-Host "Restarting script with administrator privileges..."
+        Start-Process -FilePath "powershell" -ArgumentList "-ExecutionPolicy Bypass -NoProfile -File `"$PSCommandPath`"" -Verb RunAs
+        exit
     }
+}
+Test-Admin
 
-    # Making it trusted now will avoid further user prompts in the future
+function Install-PSWindowsUpdateModule {
+    if (-not (Get-PackageProvider -ListAvailable | Where-Object Name -eq "NuGet")) {
+        Install-PackageProvider -Name NuGet -Force -Confirm:$false
+    }
     if ((Get-PSRepository -Name PSGallery).InstallationPolicy -ne "Trusted") {
         Set-PSRepository -Name PSGallery -InstallationPolicy Trusted
     }
-
-    # Install PSWindowsUpdate
     if (-not (Get-Module -ListAvailable -Name PSWindowsUpdate)) {
         Install-Module -Name PSWindowsUpdate -Force -Scope CurrentUser -Confirm:$false
     }
@@ -25,23 +33,84 @@ function Enable-MicrosoftUpdate {
     Add-WUServiceManager -ServiceID "7971f918-a847-4430-9279-4a52d1efe18d" -AddServiceFlag 7 -Confirm:$false | Out-Null
 }
 
+function Show-DriverSelection {
+    param (
+        [array]$Updates
+    )
+    if ($Updates.Count -eq 0) {
+        return @()
+    }
+    
+    Add-Type -AssemblyName PresentationFramework
+
+    $Window = New-Object System.Windows.Window
+    $Window.Title = "Select drivers to install"
+    $Window.Width = 500
+    $Window.Height = 400
+    $Window.WindowStartupLocation = "CenterScreen"
+
+    $StackPanel = New-Object System.Windows.Controls.StackPanel
+
+    $ListBox = New-Object System.Windows.Controls.ListBox
+    $ListBox.SelectionMode = "Extended"
+    foreach ($update in $Updates) {
+        $item = New-Object System.Windows.Controls.ListBoxItem
+        $item.Content = $update.Title.ToString().Trim()
+        $ListBox.Items.Add($item) | Out-Null
+    }
+    $StackPanel.Children.Add($ListBox) | Out-Null
+
+    $OKButton = New-Object System.Windows.Controls.Button
+    $OKButton.Content = "OK"
+    $OKButton.Margin = "10,10,10,10"
+    $OKButton.Add_Click({
+        $Window.Tag = $ListBox.SelectedItems
+        $Window.Close()
+    })
+    $StackPanel.Children.Add($OKButton) | Out-Null
+
+    $Window.Content = $StackPanel
+    $Window.ShowDialog() | Out-Null
+
+    $selectedUpdates = @()
+    foreach ($selected in $Window.Tag) {
+        $title = $selected.Content
+        $update = $Updates | Where-Object { $_.Title.ToString().Trim() -eq $title }
+        if ($update) {
+            $selectedUpdates += $update
+        }
+    }
+    return $selectedUpdates
+}
+
 function Update-Drivers {
     Write-Host "Checking for driver updates..."
     $Updates = Get-WUList -MicrosoftUpdate -Category "Drivers"
-
+    
     if ($Updates.Count -gt 0) {
-        Write-Host "Installing available driver updates..."
-        $Updates | Format-Table ComputerName, Status, KB, Size, Title -AutoSize
-        Get-WUInstall -MicrosoftUpdate -Category "Drivers" -AcceptAll -IgnoreReboot -Confirm:$false | Out-Null
-        Write-Host "Driver updates installed successfully!"
+        Write-Host "Available driver updates:"
+        $selection = Show-DriverSelection -Updates $Updates
+
+        $selection = @($selection)
         
-        # Restart if needed
-        if ($RestartAfterUpdate) {
-            Write-Host "Restarting the system in 10 seconds..."
-            Start-Sleep -Seconds 10
-            Restart-Computer -Force
+        if ($selection.Count -gt 0) {
+            Write-Host "Installing selected driver updates..."
+            $selection | Format-Table ComputerName, Status, KB, Size, Title -AutoSize
+            $selection | Get-WUInstall -AcceptAll -IgnoreReboot -Confirm:$false | Out-Null
+            Write-Host "Driver updates installed successfully!"
+
+            $restartChoice = Read-Host "Do you want to restart now? (Y/N)"
+            if ($restartChoice -match "^[Yy]$") {
+                Write-Host "Restarting the system in 10 seconds..."
+                Start-Sleep -Seconds 10
+                Restart-Computer -Force
+            }
         }
-    } else {
+        else {
+            Write-Host "No drivers were selected for update."
+        }
+    }
+    else {
         Write-Host "No driver updates found."
     }
 }
@@ -50,4 +119,7 @@ Install-PSWindowsUpdateModule
 Enable-MicrosoftUpdate
 Update-Drivers
 
-Write-Host "Driver update process completed."
+if (-not $PSCmdlet.MyInvocation.BoundParameters.ContainsKey('Silent')) {
+    Write-Host "Press any key to exit..."
+    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
+}