Execute OneShot tasks in any operation

scripts/Common/Scripts/Config.ps1

@@ -3,8 +3,6 @@ using namespace System.Security.AccessControl;
 using namespace System.Security.Principal;
 
 enum SetupStage {
-    Idle
-    OneShot
     Configure
     Install
     CreateUser
@@ -389,7 +387,7 @@ $null = New-Module {
         Gets a value indicating whether the setup has finished.
     #>
     function Get-IsFinished {
-        return [bool] (((Get-Stage) -eq ([SetupStage]::Idle)) -or (Get-SetupOption $finishedOption));
+        return [bool](Get-SetupOption $finishedOption);
     }
 
     <#

scripts/Common/Scripts/Operations.ps1

@@ -29,18 +29,24 @@ $null = New-Module {
 
     function Start-Operation {
         param(
-            [switch] $NonInteractive = ((Get-Stage) -eq ([SetupStage]::OneShot)),
+            [switch] $NonInteractive,
             [switch] $NoImplicitCleanup,
             [scriptblock] $Action
         )
 
         $cleanup = { };
+        $taskPending = $false;
 
         if (-not $Global:InOperation) {
             if ($env:DEBUG) {
                 Set-PSDebug -Trace 1;
             }
 
+            if ($null -ne (Get-OneShotTask)) {
+                $taskPending = $true;
+                [switch] $NonInteractive = $true;
+            }
+
             $Global:InOperation = $true;
             $Global:ErrorActionPreference = $NonInteractive.IsPresent ? 'Continue' : 'Inquire';
 
@@ -229,7 +235,12 @@ $null = New-Module {
             }
         }
 
+        if ($taskPending) {
+            Start-OneShot;
+        } else {
             & $Action;
+        }
+
         & $cleanup;
     }
 
@@ -238,7 +249,13 @@ $null = New-Module {
         Gets the current OneShot task.
     #>
     function Get-OneShotTask {
-        [OneShotTask](Get-SetupOption $taskOption);
+        $task = Get-SetupOption $taskOption;
+
+        if ($task) {
+            return [OneShotTask]$task;
+        } else {
+            return $null;
+        }
     }
 
     <#
@@ -291,8 +308,6 @@ $null = New-Module {
             [OneShotTask] $Task
         )
 
-        $currentStage = Get-Stage;
-        Set-Stage ([SetupStage]::OneShot);
         & $taskSetter $Task;
 
         & {
@@ -321,8 +336,6 @@ $null = New-Module {
             }
         };
 
-        Set-Stage $currentStage;
-
         if (Test-Path $errorPath) {
             $errorMessage = Get-Content $errorPath;
             Remove-Item $errorPath;
@@ -340,23 +353,27 @@ $null = New-Module {
         Executes the specified action and notifies the OneShot task executor.
     #>
     function Start-OneShot {
-        param(
-            [scriptblock] $Action
-        )
-
         try {
-            Start-Operation -NonInteractive @PSBoundParameters;
+            Write-Host "Running OneShot task ``$(Get-OneShotTask)``";
+
+            switch (Get-OneShotTask) {
+                ([OneShotTask]::InitializeMSAccount) {
+                    Initialize-UserCreation;
+                }
+                ([OneShotTask]::DisableUAC) {
+                    Disable-UAC;
+                    Register-Setup;
+                }
+            }
         }
         catch {
             Set-Content -Path $errorPath -Value $Error;
             Set-UserPermissions $errorPath;
         }
         finally {
-            Set-Stage ([SetupStage]::Idle);
+            & $taskSetter $null;
             Write-EventLog -LogName $logName -Source $logName -EventId $oneShotTrigger -Message "The OneShot task ``$(Get-OneShotTask)`` finished.";
         }
-
-        exit;
     }
 
     <#

scripts/Windows/OS/Install.ps1

@@ -38,29 +38,6 @@ $null = New-Module {
     #>
     function Start-InstallationLoop {
         while (-not (Get-IsFinished)) {
-            switch (Get-Stage) {
-                ($null) {
-                    Set-Stage ([SetupStage]::Configure);
-                    break;
-                }
-                ([SetupStage]::OneShot) {
-                    Write-Host "Running OneShot task ``$(Get-OneShotTask)``";
-
-                    Start-OneShot {
-                        switch (Get-OneShotTask) {
-                            ([OneShotTask]::InitializeMSAccount) {
-                                Initialize-UserCreation;
-                            }
-                            ([OneShotTask]::DisableUAC) {
-                                Disable-UAC;
-                                Register-Setup;
-                            }
-                        }
-                    };
-
-                    break;
-                }
-                default {
             if (Test-Admin) {
                 $null = Import-Module PSWindowsUpdate;
 
@@ -361,7 +338,11 @@ $null = New-Module {
                 };
             }
 
-                    switch ($_) {
+            switch (Get-Stage) {
+                ($null) {
+                    Set-Stage ([SetupStage]::Configure);
+                    break;
+                }
                 ([SetupStage]::Configure) {
                     if (Get-Config "valhalla.windows.dualboot.enable") {
                         if (-not (Test-Qemu)) {
@@ -482,8 +463,6 @@ $null = New-Module {
             }
         }
     }
-        }
-    }
     
     function Invoke-WindowsInstallation([Context] $context) {
         $Global:InformationPreference = "Continue";