From ebf89c0f3194e03247010d0dec8daadf149e60a6 Mon Sep 17 00:00:00 2001
From: Manuel Thalmann <m@nuth.ch>
Date: Mon, 29 Jul 2024 01:44:12 +0200
Subject: [PATCH] Load dualboot sizes from nix

---
 profiles/DerGeret/Windows/Setup.ps1 | 135 +--------------------
 scripts/Windows/OS/Setup.ps1        | 179 +++++++++++++++++++++++++++-
 2 files changed, 175 insertions(+), 139 deletions(-)

diff --git a/profiles/DerGeret/Windows/Setup.ps1 b/profiles/DerGeret/Windows/Setup.ps1
index 55b1a368..e7d5d62b 100644
--- a/profiles/DerGeret/Windows/Setup.ps1
+++ b/profiles/DerGeret/Windows/Setup.ps1
@@ -1,139 +1,6 @@
 #!/bin/pwsh
 $env:WIN_COMPUTER_NAME = "DerGeret";
 $env:SETUP_SCRIPT_NAME = "$PSScriptRoot/Restore.ps1";
-
-function Initialize-SetupConfig() {
-    param(
-        [xml] $config,
-        [System.Xml.XmlNamespaceManager] $namespace
-    );
-
-    $setupComponent = $config.SelectSingleNode(
-        "/ua:unattend/ua:settings[@pass='windowsPE']/ua:component[@name='Microsoft-Windows-Setup']",
-        $namespace);
-
-    $diskConfig = $setupComponent.SelectSingleNode("./ua:DiskConfiguration/ua:Disk", $namespace);
-
-    $partitionCreationContainer = $diskConfig.SelectSingleNode("./ua:CreatePartitions", $namespace);
-    $partitionCreations = $partitionCreationContainer.SelectNodes("./ua:CreatePartition", $namespace);
-
-    $partitionModificationContainer = $diskConfig.SelectSingleNode("./ua:ModifyPartitions", $namespace);
-    $partitionModifications = $partitionModificationContainer.SelectNodes("./ua:ModifyPartition", $namespace);
-
-    <#
-    .SYNOPSIS
-        Gets the XML element describing the installation partition ID.
-    #>
-    function Get-InstallationPartition {
-        $setupComponent.SelectSingleNode("./ua:ImageInstall/ua:OSImage/ua:InstallTo/ua:PartitionID", $namespace)
-    }
-
-    <#
-    .SYNOPSIS
-        Increases the ID of all partitions in the specified range by 1.
-    #>
-    function Move-PartitionRange {
-        param (
-            [int]$From = 0,
-            [System.Nullable[int]]$To = $null,
-            [int]$By = 1
-        )
-
-        # Update installation partition ID if necessary
-        $installationPartition = Get-InstallationPartition;
-        $installPartitionID = [int]$installationPartition.InnerText;
-
-        if (($installPartitionID -ge $From) -and (($null -eq $To) -or ($installPartitionID -le $To))) {
-            $installationPartition.InnerText = "$($installPartitionID + $By)";
-        }
-
-        # Update IDs of all partition creations
-        foreach ($partition in $partitionCreations) {
-            $orderNode = $partition.SelectSingleNode("./ua:Order", $namespace);
-            $order = [int]$orderNode.InnerText;
-            $newOrder = $order;
-
-            if (($newOrder -ge $From) -and (($null -eq $To) -or ($newOrder -le $To))) {
-                $newOrder += $By;
-            }
-
-            if ($order -ne $newOrder) {
-                $orderNode.InnerText = "$newOrder";
-            }
-        }
-
-        # Update IDs of all partition modifications
-        foreach ($partition in $partitionModifications) {
-            $partitionNode = $partition.SelectSingleNode("./ua:PartitionID", $namespace);
-            $partitionID = [int]$partitionNode.InnerText;
-            $newID = $partitionID;
-
-            if (($newID -ge $From) -and (($null -eq $To) -or ($newID -le $To))) {
-                $newID += $By;
-            }
-
-            if ($partitionID -ne $newID) {
-                $partitionNode.InnerText = "$newID";
-                $partition.SelectSingleNode("./ua:Order", $namespace).InnerText = "$newID";
-            }
-        }
-    }
-
-    function Add-Partition {
-        param (
-            [int]$Index,
-            [int]$Size,
-            [string]$Type = "Primary"
-        )
-
-        Move-PartitionRange -From $Index -By 1;
-
-        $newPartition = $partitionCreations[0].CloneNode($true);
-        $newPartition.SelectSingleNode("./ua:Order", $namespace).InnerText = "$Index";
-        $newPartition.SelectSingleNode("./ua:Type", $namespace).InnerText = "$Type";
-        $newPartition.SelectSingleNode("./ua:Size", $namespace).InnerText = "$Size";
-        $null = $partitionCreationContainer.AppendChild($newPartition);
-
-        $newModification = $partitionModifications[2].CloneNode($true);
-        $newModification.SelectSingleNode("./ua:Order", $namespace).InnerText = "$Index";
-        $newModification.SelectSingleNode("./ua:PartitionID", $namespace).InnerText = "$Index";
-        $null = $partitionModificationContainer.AppendChild($newModification);
-    }
-
-    <#
-    .SYNOPSIS
-        Relocates the partition with the specified `$From` ID to the specified `$To` ID.
-    #>
-    function Invoke-PartitionRelocation {
-        param (
-            [int]$From,
-            [int]$To
-        )
-
-        Move-PartitionRange $From $From (-1 * ($From + 1))
-
-        if ($From -gt $To) {
-            Move-PartitionRange $To ($From - 1);
-        }
-        elseif ($From -lt $To) {
-            Move-PartitionRange ($From + 1) $To -1;
-        }
-
-        Move-PartitionRange -1 -1 ($To + 1)
-    }
-
-    # Resize EFI partition to 1GB
-    $partitionCreations[1].SelectSingleNode("./ua:Size", $namespace).InnerText = "$(1024)";
-
-    # Swap Windows RE partition (partition #1)  and boot partition (partition #2)
-    Invoke-PartitionRelocation 2 1;
-
-    # Add space before Windows installation... wha-!? For Linux, ofc! I use Arch Linux, btw.
-    $swapSpacer = 100;
-    Add-Partition 2 $swapSpacer;
-
-    # Add a 1.2 TB partition for Linux
-    Add-Partition 3 ((1.2 * 1024 * 1024) - 1024 - $swapSpacer);
-}
+$env:CONFIG_MODULE = "$PSScriptRoot/../config.nix";
 
 . "$PSScriptRoot/../../../scripts/Windows/OS/Setup.ps1";
diff --git a/scripts/Windows/OS/Setup.ps1 b/scripts/Windows/OS/Setup.ps1
index f10efc0a..71d64cfa 100644
--- a/scripts/Windows/OS/Setup.ps1
+++ b/scripts/Windows/OS/Setup.ps1
@@ -6,7 +6,7 @@ function Start-Setup {
     $null = $env:SETUP_SCRIPT_NAME;
     $null = $env:CONFIG_MODULE;
 
-    $config = ConvertFrom-Json (Get-Content "$env:CONFIG_MODULE.json");
+    $valhallaConfig = ConvertFrom-Json (Get-Content "$env:CONFIG_MODULE.json");
 
     [xml]$unattendedConfig = [xml]::new();
     $unattendedConfig.PreserveWhitespace = $true;
@@ -19,6 +19,10 @@ function Start-Setup {
     $namespace = New-Object -TypeName "Xml.XmlNamespaceManager" -ArgumentList $unattendedConfig.NameTable;
     $namespace.AddNamespace("ua", $unattendedConfig.DocumentElement.NamespaceURI);
 
+    function Get-InstallDisk {
+        [int]$installTarget.Disk.InnerText;
+    }
+
     function Get-PassSettings {
         [OutputType([xml])]
         param(
@@ -28,6 +32,22 @@ function Start-Setup {
         return $unattendedConfig.SelectSingleNode("/ua:unattend/ua:settings[@pass='$passName']", $namespace);
     }
 
+    function Get-Component {
+        [OutputType([xml])]
+        param(
+            [System.Xml.XmlNode] $pass,
+            [string] $componentName = $null
+        );
+
+        $Local:selector = "./ua:component";
+
+        if ($null -ne $componentName) {
+            $Local:selector += "[@name='$componentName']";
+        }
+
+        return $pass.SelectSingleNode($Local:selector, $namespace);
+    }
+
     function Get-RemoteScriptPath($path) {
         $relativePath = [System.IO.Path]::GetRelativePath("$PSScriptRoot/../../..", $path);
         Join-Path $env:REMOTE_PROJECT_PATH $relativePath;
@@ -43,9 +63,144 @@ function Start-Setup {
         "(Join-Path `$env:SystemDrive $(Get-Injection (Get-RemoteScriptPath $path)))"
     }
 
+    function Get-DiskConfig {
+        param(
+            [int] $Disk
+        )
+
+        $node = $disks | Where-Object { $_.SelectSingleNode("./ua:DiskID", $namespace).InnerText -eq $Disk };
+        $creations = $node.SelectSingleNode("./ua:CreatePartitions", $namespace);
+        $modifications = $node.SelectSingleNode("./ua:ModifyPartitions", $namespace);
+
+        @{
+            PartitionCreationContainer = $creations;
+            PartitionCreations = $creations.SelectNodes("./ua:CreatePartition", $namespace);
+            PartitionModificationContainer = $modifications;
+            PartitionModifications = $modifications.SelectNodes("./ua:ModifyPartition", $namespace);
+        };
+    }
+
+    function Move-PartitionRange {
+        param (
+            [int] $Disk = (Get-InstallDisk),
+            [Parameter(Position=0)]
+            [int] $From = 0,
+            [Parameter(Position=1)]
+            [System.Nullable[int]] $To = $null,
+            [Parameter(Position=2)]
+            [int] $By = 1
+        )
+
+        $diskInfo = Get-DiskConfig $Disk;
+
+        if ((Get-InstallDisk) -eq $Disk) {
+            $partition = [int]$installTarget.Partition.InnerText;
+
+            if (($partition -ge $From) -and (($null -eq $To) -or ($partition -le $To))) {
+                $installTarget.Partition.InnerText = "$($partition + $By)";
+            }
+        }
+
+        foreach ($config in @(@($diskInfo.PartitionCreations, @("Order")),
+                              @($diskInfo.PartitionModifications, @("Order", "PartitionID")))) {
+            foreach ($partition in $config[0]) {
+                foreach ($property in $config[1]) {
+                    $partitionNode = $partition.SelectSingleNode("./ua:$property", $namespace);
+                    $partitionID = [int]$partitionNode.InnerText;
+                    $newID = $partitionID;
+
+                    if (($newID -ge $From) -and (($null -eq $To) -or ($newID -le $To))) {
+                        $newID += $By;
+                    }
+
+                    if ($partitionID -ne $newID) {
+                        $partitionNode.InnerText = "$newID";
+                    }
+                }
+            }
+        }
+    }
+
+    function Add-Partition {
+        param (
+            [int] $Disk = (Get-InstallDisk),
+            [Parameter(Position=0)]
+            [int] $Index,
+            [Parameter(Position=1)]
+            [int] $Size,
+            [Parameter(Position=2)]
+            [string] $Type = "Primary"
+        )
+
+        $diskInfo = Get-DiskConfig $Disk;
+        Move-PartitionRange -Disk $Disk -From $Index -By 1;
+
+        $configs = @(
+            @(
+                $diskInfo.PartitionCreations,
+                0,
+                @(
+                    @("Order", "$Index"),
+                    @("Type", "$Type"),
+                    @("Size", "$Size"))),
+            @(
+                $diskInfo.PartitionModifications,
+                2,
+                @(
+                    @("Order", "$Index"),
+                    @("PartitionID", "$Index")))
+        );
+
+        foreach ($config in $configs) {
+            $partition = $config[0][$config[1]];
+            $newPartition = $partition.CloneNode($true);
+
+            foreach ($entry in $config[2]) {
+                $newPartition.SelectSingleNode("./ua:$($entry[0])", $namespace).InnerText = $entry[1];
+            }
+
+            $null = $partition.ParentNode.AppendChild($newPartition);
+        }
+    }
+
+    function Move-Partition {
+        param (
+            [int] $Disk = (Get-InstallDisk),
+            [Parameter(Position=0)]
+            [int] $From,
+            [Parameter(Position=1)]
+            [int] $To
+        )
+
+        Move-PartitionRange -Disk $Disk $From $From (-1 * ($From + 1));
+
+        if ($From -gt $To) {
+            Move-PartitionRange -Disk $Disk $To ($From - 1);
+        }
+        elseif ($From -lt $To) {
+            Move-PartitionRange -Disk $Disk ($From + 1) $To - 1;
+        }
+
+        Move-PartitionRange -Disk $Disk -1 -1 ($To + 1)
+    }
+
+    # Collect necessary variables
+    $winpePass = Get-PassSettings "windowsPE";
+    $setupConfig = Get-Component $winpePass "Microsoft-Windows-Setup";
+    $disks = $setupConfig.SelectNodes("./ua:DiskConfiguration/ua:Disk", $namespace);
+
+    $installTarget = & {
+        $target = $setupConfig.SelectSingleNode("./ua:ImageInstall/ua:OSImage/ua:InstallTo", $namespace);
+
+        @{
+            Disk = $target.SelectSingleNode("./ua:DiskID", $namespace);
+            Partition = $target.SelectSingleNode("./ua:PartitionID", $namespace);
+        };
+    };
+
     # Adjust unattended settings
-    $specializeSettings = Get-PassSettings "specialize";
-    $specializeSettings.SelectSingleNode("./ua:component[@name='Microsoft-Windows-Shell-Setup']/ua:ComputerName", $namespace).InnerText = "$env:WIN_COMPUTER_NAME";
+    $computerName = (Get-Component (Get-PassSettings "specialize") "Microsoft-Windows-Shell-Setup").SelectSingleNode("./ua:ComputerName", $namespace);
+    $computerName.InnerText = "$env:WIN_COMPUTER_NAME";
 
     # Execute corresponding installer script after startup
     $oobeSystemSettings = Get-PassSettings "oobeSystem";
@@ -62,8 +217,22 @@ function Start-Setup {
     $orderElement.InnerText = ([int]($orderElement.InnerText) + 1);
     $newCommand.SelectSingleNode("./ua:Description", $namespace).InnerText = "Install PowerShell Core and git and run setup script";
 
-    if (Get-Command Initialize-SetupConfig -ErrorAction SilentlyContinue) {
-        Initialize-SetupConfig $unattendedConfig $namespace;
+    if ($valhallaConfig.dualboot.enable) {
+        $diskSize = [long](ConvertFrom-Csv (wmic diskdrive where "Index=$(Get-InstallDisk)" get Size | ForEach-Object { "$_".Trim(); })).Size;
+
+        # Calculate Linux size
+        $linuxSize = ([System.Math]::Floor(($diskSize * ($valhallaConfig.dualboot.linuxPercentage / 100)) / 1024 / 1024 / 1024)) * 1024;
+
+        # Resize EFI partition to 1GB (for GRUB)
+        (Get-DiskConfig (Get-InstallDisk)).PartitionCreations[1].SelectSingleNode("./ua:Size", $namespace).InnerText = "$(1024)";
+
+        # Move boot partition to the beginning
+        Move-Partition 2 1;
+
+        # Add a Swap spacer
+        $swapSpacer = 100;
+        Add-Partition 2 $swapSpacer;
+        Add-Partition 3 ($linuxSize - $swapSpacer);
     }
 
     $unattendedConfigFile = "X:\unattend.xml";