From ed4fa1a33d087f0e999424974e282fba0913df7f Mon Sep 17 00:00:00 2001
From: Manuel Thalmann <m@nuth.ch>
Date: Sat, 7 Dec 2024 16:12:04 +0100
Subject: [PATCH] Split up disk script into multiple scripts

---
 lib/modules/partition/disks.nix | 115 +++++++++++++++++++-------------
 1 file changed, 68 insertions(+), 47 deletions(-)

diff --git a/lib/modules/partition/disks.nix b/lib/modules/partition/disks.nix
index 2a2c580c..17b5561f 100644
--- a/lib/modules/partition/disks.nix
+++ b/lib/modules/partition/disks.nix
@@ -298,9 +298,50 @@ in
             default = { };
           };
 
+          scripts = {
+            init = mkOption {
+              type = types.str;
+              description = "The script for initializing the disk partitioning script.";
+            };
+
+            partition = mkOption {
+              type = types.str;
+              description = "The script for partitioning the disks.";
+            };
+
+            mount = mkOption {
+              type = types.str;
+              description = "The script for mounting the partitioned disks.";
+            };
+          };
+
           script = mkOption {
             type = types.str;
             description = "The script for partitioning the system's disks.";
+            default =
+              let
+                cfg = config.valhalla.fileSystems.diskSetup;
+                inherit (cfg) scripts;
+                disks = (builtins.attrValues cfg.disks);
+              in
+              ''
+                #!/bin/bash
+                set -o errexit
+                ${scripts.init}
+                ${lib.strings.concatLines (lib.optionals ((builtins.length disks) > 0) [
+                  ''echo "$(tput setaf 3)==== WARNING ====$(tput sgr0)"''
+                  (''echo "Continuing this script will alter the partitions of ''
+                    + (lib.strings.concatStringsSep ", " (builtins.map (_: "${_.deviceVariable}") (lib.lists.init disks)))
+                    + (if (builtins.length disks) > 1 then " and " else "") + (lib.lists.last disks).deviceVariable + ''"'')
+                  ''
+                    if ! fish ${./confirm.fish} "Are you sure you want to continue?" "n"; then
+                      exit 1
+                    fi
+                  ''
+                ])}
+                ${scripts.partition}
+                ${scripts.mount}
+              '';
           };
         };
       };
@@ -310,7 +351,7 @@ in
   config = {
     valhalla = {
       fileSystems.diskSetup = {
-        script = lib.mkDefault (
+        scripts =
           let
             cfg = config.valhalla.fileSystems;
             inherit (cfg) rootDir;
@@ -318,58 +359,38 @@ in
             partPath = part: "/dev/disk/by-label/${part.label}";
             disks = ((builtins.attrValues cfg.diskSetup.disks));
             partitions = (builtins.concatMap (_: (builtins.attrValues _.partitions)) disks);
-
-            mountScript = lib.strings.concatLines (builtins.concatMap
-              (
-                _: [
-                  probeScript
-                  (builtins.concatStringsSep " " ([
-                    "sudo"
-                    "mount"
-                    "--mkdir"
-                  ] ++ (lib.optionals (_.format == "ntfs") [
-                    "-t"
-                    "ntfs3"
-                  ]) ++ [
-                    (builtins.concatStringsSep " " (builtins.map (_: "-o ${_}") _.mountOptions))
-                    (partPath _)
-                    (normalizePath "/${rootDir}/${_.mountPoint}")
-                  ]))
-                ]
-              )
-              (lib.lists.sortOn
-                (_: normalizePath "/${_.mountPoint}")
-                (builtins.filter (_: _.mountPoint != null) partitions)));
-
-            swapScript = lib.strings.concatLines (builtins.map
+          in
+          {
+            init = lib.strings.concatLines (builtins.map (_: _.deviceScript) disks);
+            partition = lib.strings.concatLines (builtins.map (_: _.script) disks);
+            mount = lib.strings.concatLines (
+              (builtins.concatMap
+                (
+                  _: [
+                    probeScript
+                    (builtins.concatStringsSep " " (
+                      [ "sudo" "mount" "--mkdir" ] ++
+                        (lib.optionals (_.format == "ntfs") [ "-t" "ntfs3" ]) ++
+                        [
+                          (builtins.concatStringsSep " " (builtins.map (_: "-o ${_}") _.mountOptions))
+                          (partPath _)
+                          (normalizePath "/${rootDir}/${_.mountPoint}")
+                        ]
+                    ))
+                  ]
+                )
+                (lib.lists.sortOn
+                  (_: normalizePath "/${_.mountPoint}")
+                  (builtins.filter (_: _.mountPoint != null) partitions))) ++
+            (builtins.map
               (
                 _: ''
                   ${probeScript}
                   sudo swapon ${partPath _}
                 ''
               )
-              (builtins.filter (_: _.useSwap) partitions));
-          in
-          lib.strings.concatLines ([
-            "#!/bin/bash"
-            "set -o errexit"
-          ]
-          ++ (builtins.map (_: _.deviceScript) disks)
-          ++ lib.optionals ((builtins.length disks) > 0) [
-            ''echo "$(tput setaf 3)==== WARNING ====$(tput sgr0)"''
-            (''echo "Continuing this script will alter the partitions of ''
-              + (lib.strings.concatStringsSep ", " (builtins.map (_: "${_.deviceVariable}") (lib.lists.init disks)))
-              + (if (builtins.length disks) > 1 then " and " else "") + (lib.lists.last disks).deviceVariable + ''"'')
-            ''
-              if ! fish ${./confirm.fish} "Are you sure you want to continue?" "n"; then
-                exit 1
-              fi
-            ''
-          ] ++ (builtins.map (_: _.script) disks) ++ [
-            mountScript
-            swapScript
-          ])
-        );
+              (builtins.filter (_: _.useSwap) partitions)));
+          };
       };
     };
   };