From 879c4bc5d43e7b0058c3cd21351e4fc27082cf75 Mon Sep 17 00:00:00 2001
From: Manuel Thalmann <m@nuth.ch>
Date: Tue, 3 Sep 2024 22:15:15 +0200
Subject: [PATCH] Add code for cleaning up after installations

---
 scripts/Common/Scripts/Config.ps1     |  1 +
 scripts/Common/Scripts/Operations.ps1 | 51 ++++++++++++++++++---------
 scripts/Windows/OS/Manage.ps1         | 32 +++++++++++++++++
 3 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/scripts/Common/Scripts/Config.ps1 b/scripts/Common/Scripts/Config.ps1
index aefedcd3..2c552bbd 100644
--- a/scripts/Common/Scripts/Config.ps1
+++ b/scripts/Common/Scripts/Config.ps1
@@ -5,6 +5,7 @@ using namespace System.Security.Principal;
 enum WindowsInstallerStage {
     Initialize
     Run
+    Cleanup
     Completed
 }
 
diff --git a/scripts/Common/Scripts/Operations.ps1 b/scripts/Common/Scripts/Operations.ps1
index a5bf40d5..6958c8be 100644
--- a/scripts/Common/Scripts/Operations.ps1
+++ b/scripts/Common/Scripts/Operations.ps1
@@ -27,6 +27,33 @@ $null = New-Module {
         Set-SetupOption $taskOption ([string]$Task);
     };
 
+    <#
+        .SYNOPSIS
+        Gets the PowerShell modules required for operating.
+    #>
+    function Get-RequiredModules {
+        $modules = @(
+            @("PSScriptAnalyzer")
+        ) + (& {
+            if (-not $IsWindows) {
+                @()
+            } else {
+                @(
+                    @("KnownFolders"),
+                    @("PSWindowsUpdate"),
+                    @("LocalAccounts", $true),
+                    @("NuGet")
+                )
+            }
+        });
+
+        for ($i = 0; $i -lt $modules.Count; $i++) {
+            if ($modules[$i] -is [string]) {
+                $modules[$i] = @($modules[$i]);
+            }
+        }
+    }
+
     function Start-Operation {
         param(
             [switch] $NonInteractive,
@@ -215,22 +242,7 @@ $null = New-Module {
                     }
 
                     Invoke-Hook "Install-PSModules" -Fallback {
-                        $modules = @(
-                            @("PSScriptAnalyzer")
-                        ) + (& {
-                            if (-not $IsWindows) {
-                                @()
-                            } else {
-                                @(
-                                    @("KnownFolders"),
-                                    @("PSWindowsUpdate"),
-                                    @("LocalAccounts", $true),
-                                    @("NuGet")
-                                )
-                            }
-                        });
-
-                        foreach ($module in $modules) {
+                        foreach ($module in (Get-RequiredModules)) {
                             $parameters = @{ };
 
                             if ($module -is [string]) {
@@ -403,6 +415,13 @@ $null = New-Module {
     function Clear-OperationResources {
         if ($IsWindows) {
             Uninstall-WslDistribution;
+            $null = Uninstall-Package Selenium.WebDriver -ErrorAction Continue;
+            Uninstall-ChocoPackage 7zip.portable gsudo selenium-gecko-driver yq;
+            Uninstall-WingetPackage AutoHotkey.AutoHotkey;
+        }
+
+        foreach ($module in (Get-RequiredModules)) {
+            Uninstall-Module -Force -Name $module[0];
         }
     }
 };
diff --git a/scripts/Windows/OS/Manage.ps1 b/scripts/Windows/OS/Manage.ps1
index 8f26ff2a..1b3e3946 100644
--- a/scripts/Windows/OS/Manage.ps1
+++ b/scripts/Windows/OS/Manage.ps1
@@ -20,6 +20,7 @@ $null = New-Module {
     . "$PSScriptRoot/../Types/WindowsInstallerAction.ps1";
     . "$PSScriptRoot/../../Common/Scripts/Config.ps1";
     . "$PSScriptRoot/../../Common/Scripts/Operations.ps1";
+    . "$PSScriptRoot/../../Common/Scripts/Scripting.ps1";
     . "$PSScriptRoot/../../Common/Scripts/Software.ps1";
     . "$PSScriptRoot/../../Common/Scripts/SoftwareManagement.ps1";
     . "$PSScriptRoot/../../Common/Types/InstallerAction.ps1";
@@ -119,6 +120,8 @@ $null = New-Module {
                                         Write-Host "Entering install phase";
                                         Deploy-SoftwareAction;
                                         Set-SetupStage ([SetupStage]::CreateUser);
+                                        Restart-Intermediate;
+                                        return;
                                     }
                                     ([SetupStage]::CreateUser) {
                                         Install-ValhallaUsers;
@@ -129,7 +132,36 @@ $null = New-Module {
                         }
                     }
 
+                    Set-Stage ([WindowsInstallerStage]::Cleanup);
+                    break;
+                }
+                ([WindowsInstallerStage]::Cleanup) {
+                    $taskName = "PortValhalla Cleaner";
+                    $setupUser = Get-SetupUser;
+                    Clear-OperationResources;
+                    Remove-Item -Recurse -Force "C:\ProgramData\PortValhalla";
+                    Get-SetupConfigKey | Remove-Item -Recurse -Force;
+                    Disable-Autologin;
+                    Disable-LocalUser $setupUser;
+
+                    $script = {
+                        param(
+                            $TaskName,
+                            $UserName
+                        )
+
+                        $user = Get-LocalUser $UserName;
+                        [string] $sid = $user.SID;
+                        Remove-LocalUser $user;
+                        Get-CimInstance Win32_UserProfile | Where-Object { $_.SID -eq $sid } | Remove-CimInstance;
+                        Unregister-ScheduledTask -Confirm:$false $TaskName;
+                    };
+
+                    $trigger = New-ScheduledTaskTrigger -AtStartup;
+                    $task = New-ScheduledTaskAction -Execute "pwsh" -Argument "-Command & { $script } $(ConvertTo-Injection $taskName) $(ConvertTo-Injection $setupUser)";
+                    $null = Register-ScheduledTask -Force $taskName -Action $task -Trigger $trigger -RunLevel Highest -User "SYSTEM";
                     Set-Stage ([WindowsInstallerStage]::Completed);
+                    Restart-Intermediate -NoRegister;
                     break;
                 }
             }