From 5914bc35700be8472e97d011f437f414e7fa7cad Mon Sep 17 00:00:00 2001
From: Manuel Thalmann <m@nuth.ch>
Date: Tue, 27 Aug 2024 04:24:29 +0200
Subject: [PATCH] Initialize operation for all users

---
 scripts/Common/Scripts/Config.ps1     |  1 -
 scripts/Common/Scripts/Operations.ps1 | 14 +++++++++-----
 scripts/Windows/OS/Install.ps1        |  8 ++++----
 scripts/Windows/Scripts/WSL.ps1       | 11 +++++++++++
 4 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/scripts/Common/Scripts/Config.ps1 b/scripts/Common/Scripts/Config.ps1
index e448cea5..af9d51d9 100644
--- a/scripts/Common/Scripts/Config.ps1
+++ b/scripts/Common/Scripts/Config.ps1
@@ -4,7 +4,6 @@ using namespace System.Security.Principal;
 
 enum SetupStage {
     Idle
-    Initialize
     OneShot
     Configure
     Install
diff --git a/scripts/Common/Scripts/Operations.ps1 b/scripts/Common/Scripts/Operations.ps1
index 89a1b889..be0991cc 100644
--- a/scripts/Common/Scripts/Operations.ps1
+++ b/scripts/Common/Scripts/Operations.ps1
@@ -64,10 +64,10 @@ $null = New-Module {
                 };
             }
 
-            Set-Stage ((Get-Stage) ?? ([SetupStage]::Initialize));
-
             & {
-                while ((Get-Stage) -eq ([SetupStage]::Initialize)) {
+                $initialized = $false;
+
+                while (-not $initialized) {
                     if ($IsWindows) {
                         if (-not ((Test-Command "choco") -and (Test-Command "refreshenv"))) {
                             Invoke-Hook "Install-Chocolatey" -Fallback {
@@ -174,7 +174,11 @@ $null = New-Module {
                         }
 
                         if (-not (Test-WslDistributions)) {
-                            Install-WslDistribution;
+                            if (-not (Test-Path (Get-WslDistributionDisk))) {
+                                Install-WslDistribution;
+                            }
+
+                            Register-WslDistribution;
                             continue;
                         }
 
@@ -219,7 +223,7 @@ $null = New-Module {
 
                         Install-ChocoPackage selenium-gecko-driver firefox;
                         Install-WingetPackage AutoHotkey.AutoHotkey;
-                        Set-Stage ([SetupStage]::Configure);
+                        $initialized = $true;
                     }
                 }
             }
diff --git a/scripts/Windows/OS/Install.ps1 b/scripts/Windows/OS/Install.ps1
index 96ddcff5..418c1d26 100644
--- a/scripts/Windows/OS/Install.ps1
+++ b/scripts/Windows/OS/Install.ps1
@@ -38,11 +38,11 @@ $null = New-Module {
     #>
     function Start-InstallationLoop {
         while (-not (Get-IsFinished)) {
-            if (-not (Test-Wsl)) {
-                Register-WslDistribution;
-            }
-
             switch (Get-Stage) {
+                ($null) {
+                    Set-Stage ([SetupStage]::Configure);
+                    break;
+                }
                 ([SetupStage]::OneShot) {
                     Write-Host "Running OneShot task ``$(Get-OneShotTask)``";
 
diff --git a/scripts/Windows/Scripts/WSL.ps1 b/scripts/Windows/Scripts/WSL.ps1
index 8f7b3b88..1273d397 100644
--- a/scripts/Windows/Scripts/WSL.ps1
+++ b/scripts/Windows/Scripts/WSL.ps1
@@ -55,6 +55,12 @@ function Install-Wsl {
 function Install-WslDistribution {
     $dir = Get-WslDistributionPath;
     $root = Split-Path -Parent $dir;
+    $registryPath = "HKCU:/Software/Microsoft/Windows/CurrentVersion/Lxss";
+    $key = Get-Item $registryPath;
+
+    if ($key) {
+        $key = $key | Rename-Item -NewName "$(Split-Path -Leaf $key)_" -PassThru;
+    }
 
     if (-not (Test-Path $root)) {
         $null = New-Item -ItemType Directory $root;
@@ -64,6 +70,11 @@ function Install-WslDistribution {
     Set-UserPermissions $dir;
     & "$dir\ubuntu.exe" install --root;
     wsl --shutdown;
+    Remove-Item -Recurse -Force $registryPath;
+
+    if ($key) {
+        Move-Item $key.PSPath $registryPath;
+    }
 }
 
 <#