From 6ba1850678e3ddd54a92d5019760b7956fa0446f Mon Sep 17 00:00:00 2001
From: Michael Hoang <enzime@users.noreply.github.com>
Date: Tue, 21 Jan 2025 16:27:49 +1100
Subject: [PATCH] flash: don't allow partitioning time secrets

---
 checks/flash/flake-module.nix         | 37 ++++++++++++++++++++-------
 pkgs/clan-cli/clan_cli/flash/flash.py | 10 ++++++--
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/checks/flash/flake-module.nix b/checks/flash/flake-module.nix
index 89a652c69..fb0c4610c 100644
--- a/checks/flash/flake-module.nix
+++ b/checks/flash/flake-module.nix
@@ -1,5 +1,25 @@
-{ self, ... }:
+{ self, lib, ... }:
 {
+  clan.machines.test-flash-machine = {
+    clan.core.networking.targetHost = "test-flash-machine";
+    fileSystems."/".device = lib.mkDefault "/dev/vda";
+    boot.loader.grub.device = lib.mkDefault "/dev/vda";
+
+    imports = [ self.nixosModules.test-flash-machine ];
+  };
+
+  flake.nixosModules = {
+    test-flash-machine =
+      { lib, ... }:
+      {
+        imports = [ self.nixosModules.test-install-machine ];
+
+        clan.core.vars.generators.test = lib.mkForce { };
+
+        disko.devices.disk.main.preCreateHook = lib.mkForce "";
+      };
+  };
+
   perSystem =
     {
       nodes,
@@ -10,19 +30,18 @@
     let
       dependencies = [
         pkgs.disko
-        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-install-machine.pkgs.perlPackages.ConfigIniFiles
-        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-install-machine.pkgs.perlPackages.FileSlurp
+        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-flash-machine.pkgs.perlPackages.ConfigIniFiles
+        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-flash-machine.pkgs.perlPackages.FileSlurp
 
-        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-install-machine.config.system.build.toplevel
-        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-install-machine.config.system.build.diskoScript
-        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-install-machine.config.system.build.diskoScript.drvPath
-        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-install-machine.config.system.clan.deployment.file
+        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-flash-machine.config.system.build.toplevel
+        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-flash-machine.config.system.build.diskoScript
+        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-flash-machine.config.system.build.diskoScript.drvPath
+        self.clanInternals.machines.${pkgs.hostPlatform.system}.test-flash-machine.config.system.clan.deployment.file
 
       ] ++ builtins.map (i: i.outPath) (builtins.attrValues self.inputs);
       closureInfo = pkgs.closureInfo { rootPaths = dependencies; };
     in
     {
-      # Currently disabled...
       checks = pkgs.lib.mkIf (pkgs.stdenv.isLinux) {
         flash = (import ../lib/test-base.nix) {
           name = "flash";
@@ -46,7 +65,7 @@
           testScript = ''
             start_all()
 
-            machine.succeed("clan flash write --debug --flake ${../..} --yes --disk main /dev/vdb test-install-machine")
+            machine.succeed("clan flash write --debug --flake ${../..} --yes --disk main /dev/vdb test-flash-machine")
           '';
         } { inherit pkgs self; };
       };
diff --git a/pkgs/clan-cli/clan_cli/flash/flash.py b/pkgs/clan-cli/clan_cli/flash/flash.py
index 0ae009a1d..66d948c9a 100644
--- a/pkgs/clan-cli/clan_cli/flash/flash.py
+++ b/pkgs/clan-cli/clan_cli/flash/flash.py
@@ -14,7 +14,7 @@ from clan_cli.facts.generate import generate_facts
 from clan_cli.facts.secret_modules import SecretStoreBase
 from clan_cli.machines.machines import Machine
 from clan_cli.nix import nix_shell
-from clan_cli.vars.generate import generate_vars_for_machine
+from clan_cli.vars.generate import generate_vars
 
 from .automount import pause_automounting
 from .list import list_possible_keymaps, list_possible_languages
@@ -54,8 +54,8 @@ def flash_machine(
             extra_args = []
         system_config_nix: dict[str, Any] = {}
 
-        generate_vars_for_machine(machine, generator_name=None, regenerate=False)
         generate_facts([machine])
+        generate_vars([machine])
 
         if system_config.language:
             if system_config.language not in list_possible_languages():
@@ -90,6 +90,12 @@ def flash_machine(
                 "users": {"root": {"openssh": {"authorizedKeys": {"keys": root_keys}}}}
             }
 
+        for generator in machine.vars_generators:
+            for file in generator.files:
+                if file.needed_for == "partitioning":
+                    msg = f"Partitioning time secrets are not supported with `clan flash write`: clan.core.vars.generators.{generator.name}.files.{file.name}"
+                    raise ClanError(msg)
+
         secret_facts_module = importlib.import_module(machine.secret_facts_module)
         secret_facts_store: SecretStoreBase = secret_facts_module.SecretStore(
             machine=machine