From f93032a6e46aee659536e567b27a1aeeb1808027 Mon Sep 17 00:00:00 2001
From: Manuel Thalmann <m@nuth.ch>
Date: Fri, 9 Aug 2024 14:44:57 +0200
Subject: [PATCH] Allow registering reboot for the default user

---
 scripts/Windows/Scripts/PowerManagement.ps1 | 56 +++++++++++++++++++--
 1 file changed, 52 insertions(+), 4 deletions(-)

diff --git a/scripts/Windows/Scripts/PowerManagement.ps1 b/scripts/Windows/Scripts/PowerManagement.ps1
index 59b262aa..3559211e 100644
--- a/scripts/Windows/Scripts/PowerManagement.ps1
+++ b/scripts/Windows/Scripts/PowerManagement.ps1
@@ -1,10 +1,12 @@
 using namespace Microsoft.Win32;
 
 $null = New-Module {
+    . "$PSScriptRoot/../Scripts/Registry.ps1";
     . "$PSScriptRoot/../../Common/Scripts/Config.ps1";
     . "$PSScriptRoot/../../Common/Scripts/Scripting.ps1";
     [RegistryKey] $key = $null;
-    $runOncePath = "HKLM:\Software\Microsoft\Windows\CurrentVersion\RunOnce";
+    $runOncePath = "Software\Microsoft\Windows\CurrentVersion\RunOnce";
+    $systemRunOncePath = "HKLM:\$runOncePath";
     $logonPath = "HKLM:\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Winlogon"
     $runOnceName = "PortValhalla";
     $autologinOption = "AutoAdminLogon";
@@ -17,7 +19,23 @@ $null = New-Module {
         Gets the reghistry key containing the `RunOnce` commands.
     #>
     function Get-RunOnceKey {
-        Get-Item $runOncePath;
+        param(
+            [RegistryKey] $UserKey
+        )
+
+        [string] $path = $null;
+
+        if ($UserKey) {
+            $path = "$($UserKey.PSPath)\$runOncePath";
+        } else {
+            $path = $systemRunOncePath;
+        }
+
+        if (-not (Test-Path $path)) {
+            New-Item $path;
+        } else {
+            Get-Item $path;
+        }
     }
 
     <#
@@ -29,10 +47,23 @@ $null = New-Module {
     #>
     function Register-Setup {
         param(
+            [Parameter(ParameterSetName="System")]
+            [switch] $System,
+            [Parameter(ParameterSetName="User",Mandatory)]
+            [switch] $User,
+            [Parameter(Mandatory)]
             [RegistryKey] $UserKey
         )
 
-        $key = Get-RunOnceKey $UserKey;
+        if ($User.IsPresent) {
+            if (-not $UserKey) {
+                $UserKey = Get-Item "HKCU:\";
+            }
+
+            $key = Get-RunOnceKey $UserKey;
+        } else {
+            $key = Get-RunOnceKey;
+        }
 
         Set-ItemProperty -Path $key.PSPath -Name $runOnceName -Type "ExpandString" -Value (
             "pwsh -Command " +
@@ -106,7 +137,24 @@ $null = New-Module {
         Reboots the machine intermediately and restarts the setup after the next login.
     #>
     function Restart-Intermediate {
-        Register-Setup;
+        param(
+            [switch] $DefaultUser
+        )
+
+        $register = { param($Key) Register-Setup -UserKey $Key; };
+
+        if ($DefaultUser) {
+            Edit-DefaultUserKey {
+                param(
+                    [RegistryKey] $Key
+                )
+
+                & $register -Key $Key;
+            }
+        }
+
+        & $register;
+        Read-Host "Press enter to reboot";
         Restart-Computer -Force;
     }
 }