From dbc3eeef4230ccac861e042446bff39f336c26cc Mon Sep 17 00:00:00 2001
From: Manuel Thalmann <m@nuth.ch>
Date: Mon, 3 Jul 2023 13:44:51 +0200
Subject: [PATCH] Move UAC handling to context

---
 scripts/Windows/OS/User.ps1         | 20 ++------------------
 scripts/Windows/Scripts/Context.ps1 | 13 +++++++++++++
 2 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/scripts/Windows/OS/User.ps1 b/scripts/Windows/OS/User.ps1
index 9eb734dd..59b67eef 100644
--- a/scripts/Windows/OS/User.ps1
+++ b/scripts/Windows/OS/User.ps1
@@ -64,30 +64,14 @@ function New-PersonalUser([Context] $context)
     }
 }
 
-function Get-SystemPolicyKey() {
-    [OutputType([Microsoft.Win32.RegistryKey])]
-    param()
-    $keyPath = "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\System";
-    return Get-Item "$keyPath";
-}
-
-function Get-UACState() {
-    return Get-ItemPropertyValue -Path (Get-SystemPolicyKey).PSPath -Name "EnableLUA";
-}
-
-function Set-UACState([bool] $value) {
-    $null = Set-ItemProperty -Path (Get-SystemPolicyKey).PSPath -Name "EnableLUA" -Value [int]$value;
-}
-
 function Enable-UACNextLogin([Context] $context) {
-    Set-UACState $true;
+    $context.SetUACState($true);
     $tempTask = "PortValhalla Temp";
     $autoLoginName = "PortValhalla AutoLogin Setup";
     $uacDisablerName = "PortValhalla UAC Disabler";
     $autoLoginTrigger = Get-Random -Maximum 65535;
     $uacDisablerTrigger = Get-Random -Maximum 65535;
 
-    $key = Get-SystemPolicyKey;
     $context.Set($autoLoginTriggerProperty, $autoLoginTrigger, "DWord");
     $context.Set($uacDisablerTriggerProperty, $uacDisablerTrigger, "DWord");
 
@@ -107,11 +91,11 @@ function Enable-UACNextLogin([Context] $context) {
             $uacDisablerTrigger,
             "PortValhalla UAC Disabler",
             @(
-                "Set-ItemProperty '$($key.PSPath)' -Name 'EnableLUA' -Value 0 -Type DWord;",
                 "Unregister-ScheduledTask -Confirm:`$false '$autoLoginName';",
                 "Unregister-ScheduledTask -Confirm:`$false '$uacDisablerName';",
                 ". '$PSScriptRoot/../Scripts/Context.ps1';",
                 "`$context = [Context]::new();",
+                "`$context.SetUACState(`$false);",
                 "`$context.Remove('$autoLoginTriggerProperty');",
                 "`$context.Remove('$uacDisablerTriggerProperty');",
                 "`$context.DeregisterNewUserReboot();",
diff --git a/scripts/Windows/Scripts/Context.ps1 b/scripts/Windows/Scripts/Context.ps1
index fccc7928..1f4a28a5 100644
--- a/scripts/Windows/Scripts/Context.ps1
+++ b/scripts/Windows/Scripts/Context.ps1
@@ -148,6 +148,11 @@ class Context {
         $action.Invoke($key);
     }
 
+    [Microsoft.Win32.RegistryKey] GetSystemPolicyKey() {
+        $keyPath = "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\System";
+        return Get-Item "$keyPath";
+    }
+
     [Microsoft.Win32.RegistryKey] GetRunOnceKey() {
         return $this.GetRunOnceKey($null);
     }
@@ -169,6 +174,14 @@ class Context {
         return $result;
     }
 
+    [bool] GetUACState() {
+        return [bool](Get-ItemPropertyValue -Path ($this.GetSystemPolicyKey().PSPath) -Name "EnableLUA");
+    }
+
+    [void] SetUACState([bool] $value) {
+        $null = Set-ItemProperty -Path ($this.GetSystemPolicyKey().PSPath) -Name "EnableLUA" -Value [int]$value;
+    }
+
     [void] RegisterReboot() {
         $this.RegisterReboot($null);
     }