2 files changed, 26 insertions, 4 deletions

diff --git a/system-profiles/rebuild-machines/default.nix b/system-profiles/rebuild-machines/default.nix
index 68788023..fe9a5f14 100644
--- a/system-profiles/rebuild-machines/default.nix
+++ b/system-profiles/rebuild-machines/default.nix
@@ -26,7 +26,9 @@ let
     phases = [ "buildPhase" "installPhase" ];
 
     inherit (pkgs) zsh coreutils openssh;
-    inherit (cfg) flake scriptName;
+    inherit (cfg) scriptName;
+    inherit (cfg.flake) flakeOutput;
+    flake = cfg.flake.name;
     nixosRebuild = config.system.build.nixos-rebuild;
     inherit (config.security) wrapperDir;
     inherit sshConfig;
@@ -52,8 +54,21 @@ in {
       };
       
       flake = mkOption {
-        type = types.nullOr types.str;
-        default = "git+ssh://${cfg.repoHost}/nixos?ref=flakes#${hostName}";
+        type = types.submodule {
+          options = {
+            flake = mkOption {
+              type = types.attrs;
+            };
+            flakeOutput = mkOption {
+              type = types.str;
+            };
+            name = mkOption {
+              type = types.str;
+              default = "machines";
+            };
+          };
+        };
+        default = { flake.url = "git+ssh://${cfg.repoHost}/nixos?ref=flakes"; flakeOutput = hostName; };
         description = ''
           The Flake URI of the NixOS configuration to build.
         '';
@@ -136,5 +151,12 @@ in {
         Persistent = true;
       };
     };
+
+    nix.registry = lib.mkIf (cfg.flake != null) {
+      ${cfg.flake.name} = {
+        exact = false;
+        to = cfg.flake.flake;
+      };
+    };
   };
 }
diff --git a/system-profiles/rebuild-machines/rebuild-machine.zsh b/system-profiles/rebuild-machines/rebuild-machine.zsh
index c06ef0e4..e018e7a7 100644
--- a/system-profiles/rebuild-machines/rebuild-machine.zsh
+++ b/system-profiles/rebuild-machines/rebuild-machine.zsh
@@ -7,4 +7,4 @@ fi
 export NIX_SSHOPTS="-F @sshConfig@"
 export GIT_SSH_COMMAND="@openssh@/bin/ssh -F @sshConfig@"
 export GIT_CONFIG_COUNT=1 GIT_CONFIG_KEY_0="init.defaultBranch" GIT_CONFIG_VALUE_0=main
-exec -- @nixosRebuild@/bin/nixos-rebuild --refresh --flake '@flake@' ${@:-switch}
+exec -- @nixosRebuild@/bin/nixos-rebuild --refresh --flake '@flake@#@flakeOutput@' ${@:-switch}