From be3a75bbd77a015d4bd180489d5e6f42772650ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Thalheim?= <joerg@thalheim.io>
Date: Tue, 6 Feb 2024 16:11:55 +0100
Subject: [PATCH] add support for build machines

---
 pkgs/clan-cli/clan_cli/machines/machines.py | 17 +++++++++++++++++
 pkgs/clan-cli/clan_cli/machines/update.py   | 18 ++++--------------
 pkgs/clan-cli/clan_cli/ssh/__init__.py      | 10 +++-------
 3 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/pkgs/clan-cli/clan_cli/machines/machines.py b/pkgs/clan-cli/clan_cli/machines/machines.py
index d02319e03..07874dc17 100644
--- a/pkgs/clan-cli/clan_cli/machines/machines.py
+++ b/pkgs/clan-cli/clan_cli/machines/machines.py
@@ -97,6 +97,23 @@ class Machine:
             self.name, self.target_host_address, meta={"machine": self}
         )
 
+    @property
+    def build_host(self) -> Host:
+        """
+        The host where the machine is built and deployed from.
+        Can be the same as the target host.
+        """
+        build_host = self.deployment_info.get("buildHost")
+        if build_host is None:
+            return self.target_host
+        # enable ssh agent forwarding to allow the build host to access the target host
+        return parse_deployment_address(
+            self.name,
+            build_host,
+            forward_agent=True,
+            meta={"machine": self, "target_host": self.target_host},
+        )
+
     def eval_nix(self, attr: str, refresh: bool = False) -> str:
         """
         eval a nix attribute of the machine
diff --git a/pkgs/clan-cli/clan_cli/machines/update.py b/pkgs/clan-cli/clan_cli/machines/update.py
index a83b17325..493970106 100644
--- a/pkgs/clan-cli/clan_cli/machines/update.py
+++ b/pkgs/clan-cli/clan_cli/machines/update.py
@@ -110,11 +110,6 @@ def deploy_nixos(hosts: HostGroup) -> None:
         generate_secrets(machine)
         upload_secrets(machine)
 
-        target_host = h.meta.get("target_host")
-        if target_host:
-            target_user = h.meta.get("target_user")
-            if target_user:
-                target_host = f"{target_user}@{target_host}"
         extra_args = h.meta.get("extra_args", [])
         cmd = [
             "nixos-rebuild",
@@ -132,7 +127,8 @@ def deploy_nixos(hosts: HostGroup) -> None:
             "--flake",
             f"{path}#{machine.name}",
         ]
-        if target_host:
+        if target_host := h.meta.get("target_host"):
+            target_host = f"{target_host.user or 'root'}@{target_host.host}"
             cmd.extend(["--target-host", target_host])
         ret = h.run(cmd, check=False)
         # re-retry switch if the first time fails
@@ -157,16 +153,10 @@ def get_all_machines(clan_dir: Path) -> HostGroup:
     for name, machine_data in machines.items():
         machine = Machine(name=name, flake=clan_dir, deployment_info=machine_data)
         try:
-            machine.target_host_address
+            hosts.append(machine.build_host)
         except ClanError:
             ignored_machines.append(name)
             continue
-        host = parse_deployment_address(
-            name,
-            host=machine.target_host_address,
-            meta={"machine": machine},
-        )
-        hosts.append(host)
     if not hosts and ignored_machines != []:
         print(
             "WARNING: No machines to update. The following defined machines were ignored because they do not have `clan.networking.targetHost` nixos option set:",
@@ -182,7 +172,7 @@ def get_selected_machines(machine_names: list[str], flake_dir: Path) -> HostGrou
     hosts = []
     for name in machine_names:
         machine = Machine(name=name, flake=flake_dir)
-        hosts.append(machine.target_host)
+        hosts.append(machine.build_host)
     return HostGroup(hosts)
 
 
diff --git a/pkgs/clan-cli/clan_cli/ssh/__init__.py b/pkgs/clan-cli/clan_cli/ssh/__init__.py
index 0755e6619..b0c3b3f02 100644
--- a/pkgs/clan-cli/clan_cli/ssh/__init__.py
+++ b/pkgs/clan-cli/clan_cli/ssh/__init__.py
@@ -16,12 +16,7 @@ from enum import Enum
 from pathlib import Path
 from shlex import quote
 from threading import Thread
-from typing import (
-    IO,
-    Any,
-    Generic,
-    TypeVar,
-)
+from typing import IO, Any, Generic, TypeVar
 
 # https://no-color.org
 DISABLE_COLOR = not sys.stderr.isatty() or os.environ.get("NO_COLOR", "") != ""
@@ -753,7 +748,7 @@ class HostGroup:
 
 
 def parse_deployment_address(
-    machine_name: str, host: str, meta: dict[str, Any] = {}
+    machine_name: str, host: str, forward_agent: bool = True, meta: dict[str, Any] = {}
 ) -> Host:
     parts = host.split("@")
     user: str | None = None
@@ -780,6 +775,7 @@ def parse_deployment_address(
         user=user,
         port=port,
         command_prefix=machine_name,
+        forward_agent=forward_agent,
         meta=meta,
         ssh_options=options,
     )