From 8835de9f902c9dc15a78fc24c37db0041c97f973 Mon Sep 17 00:00:00 2001
From: romanitho <96626929+Romanitho@users.noreply.github.com>
Date: Mon, 9 Oct 2023 12:52:03 +0200
Subject: [PATCH] Final touches

---
 Winget-AutoUpdate/Winget-Upgrade.ps1          | 12 +++++++++
 .../functions/Invoke-PostUpdateActions.ps1    | 26 +++++++------------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/Winget-AutoUpdate/Winget-Upgrade.ps1 b/Winget-AutoUpdate/Winget-Upgrade.ps1
index 28f1daae..cbe2684c 100644
--- a/Winget-AutoUpdate/Winget-Upgrade.ps1
+++ b/Winget-AutoUpdate/Winget-Upgrade.ps1
@@ -311,6 +311,18 @@ if (Test-Network) {
             #Run WAU in user context if feature is activated
             if ($WAUConfig.WAU_UserContext -eq 1) {
 
+                #Create User context task if not existing
+                $UserContextTask = Get-ScheduledTask -TaskName 'Winget-AutoUpdate-UserContext' -ErrorAction SilentlyContinue
+                if (!$UserContextTask) {
+                    #Create the scheduled task in User context
+                    $taskAction = New-ScheduledTaskAction -Execute "wscript.exe" -Argument "`"$($WingetUpdatePath)\Invisible.vbs`" `"powershell.exe -NoProfile -ExecutionPolicy Bypass -File `"`"`"$($WingetUpdatePath)\winget-upgrade.ps1`"`""
+                    $taskUserPrincipal = New-ScheduledTaskPrincipal -GroupId S-1-5-11
+                    $taskSettings = New-ScheduledTaskSettingsSet -Compatibility Win8 -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -ExecutionTimeLimit 03:00:00
+                    $task = New-ScheduledTask -Action $taskAction -Principal $taskUserPrincipal -Settings $taskSettings
+                    Register-ScheduledTask -TaskName 'Winget-AutoUpdate-UserContext' -TaskPath 'WAU' -InputObject $task -Force | Out-Null
+                    Write-ToLog "-> User Context task created."
+                }
+
                 #User check routine from: https://stackoverflow.com/questions/23219718/powershell-script-to-see-currently-logged-in-users-domain-and-machine-status
                 $explorerprocesses = @(Get-WmiObject -Query "Select * FROM Win32_Process WHERE Name='explorer.exe'" -ErrorAction SilentlyContinue)
                 If ($explorerprocesses.Count -eq 0) {
diff --git a/Winget-AutoUpdate/functions/Invoke-PostUpdateActions.ps1 b/Winget-AutoUpdate/functions/Invoke-PostUpdateActions.ps1
index e0c70cd9..62996c54 100644
--- a/Winget-AutoUpdate/functions/Invoke-PostUpdateActions.ps1
+++ b/Winget-AutoUpdate/functions/Invoke-PostUpdateActions.ps1
@@ -186,26 +186,18 @@ function Invoke-PostUpdateActions {
         }
     }
 
-    #Create User context task if not existing and manage its activation
-    $UserContextTask = Get-ScheduledTask -TaskName 'Winget-AutoUpdate-UserContext' -ErrorAction SilentlyContinue
-    if (!$UserContextTask) {
-        #Create the scheduled task in User context
-        $taskAction = New-ScheduledTaskAction -Execute "wscript.exe" -Argument "`"$($WingetUpdatePath)\Invisible.vbs`" `"powershell.exe -NoProfile -ExecutionPolicy Bypass -File `"`"`"$($WingetUpdatePath)\winget-upgrade.ps1`"`""
-        $taskUserPrincipal = New-ScheduledTaskPrincipal -GroupId S-1-5-11
-        $taskSettings = New-ScheduledTaskSettingsSet -Compatibility Win8 -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -ExecutionTimeLimit 03:00:00
-        $task = New-ScheduledTask -Action $taskAction -Principal $taskUserPrincipal -Settings $taskSettings
-        Register-ScheduledTask -TaskName 'Winget-AutoUpdate-UserContext' -TaskPath 'WAU' -InputObject $task -Force | Out-Null
-
-        #If scheduled task didn't exist, user context was not activated. Set it in registry.
-        New-ItemProperty $regPath -Name WAU_UserContext -Value 0 -PropertyType DWord -Force | Out-Null
-        Write-ToLog "-> User Context task created and set to 'disabled' in registry."
-    }
-    #If scheduled task exists but no WAU_UserContext setting, user context was activated. Set it in registry.
-    elseif (!($WAUConfig.WAU_UserContext)) {
+    #Activate WAU in user context if previously configured (as "Winget-AutoUpdate-UserContext" at root)
+    $UserContextTask = Get-ScheduledTask -TaskName 'Winget-AutoUpdate-UserContext' -TaskPath '\' -ErrorAction SilentlyContinue
+    if ($UserContextTask) {
+        #Remove Winget-AutoUpdate-UserContext at root.
+        Unregister-ScheduledTask $UserContextTask -Confirm:$False
+
+        #Set it in registry as activated.
         New-ItemProperty $regPath -Name WAU_UserContext -Value 1 -PropertyType DWord -Force | Out-Null
-        Write-ToLog "-> User Context task existes and set to 'enabled' in registry."
+        Write-ToLog "-> Old User Context task deleted and set to 'enabled' in registry."
     }
 
+
     ### End of post update actions ###
 
     #Reset WAU_UpdatePostActions Value