clan-cli: add update command

pkgs/clan-cli/clan_cli/ssh/__init__.py

@@ -156,6 +156,7 @@ class Host:
         host_key_check: HostKeyCheck = HostKeyCheck.STRICT,
         meta: Dict[str, Any] = {},
         verbose_ssh: bool = False,
+        ssh_options: dict[str, str] = {},
     ) -> None:
         """
         Creates a Host
@@ -179,6 +180,7 @@ class Host:
         self.host_key_check = host_key_check
         self.meta = meta
         self.verbose_ssh = verbose_ssh
+        self.ssh_options = ssh_options
 
     def _prefix_output(
         self,
@@ -451,6 +453,10 @@ class Host:
             ssh_target = self.host
 
         ssh_opts = ["-A"] if self.forward_agent else []
+
+        for k, v in self.ssh_options.items():
+            ssh_opts.extend(["-o", f"{k}={shlex.quote(v)}"])
+
         if self.port:
             ssh_opts.extend(["-p", str(self.port)])
         if self.key:

pkgs/clan-cli/clan_cli/update.py

@@ -0,0 +1,105 @@
+import argparse
+import json
+import subprocess
+
+from .ssh import Host, HostGroup, HostKeyCheck
+
+
+def deploy_nixos(hosts: HostGroup) -> None:
+    """
+    Deploy to all hosts in parallel
+    """
+
+    flake_store_paths = {}
+    for h in hosts.hosts:
+        flake_uri = str(h.meta.get("flake_uri", ".#"))
+        if flake_uri not in flake_store_paths:
+            res = subprocess.run(
+                [
+                    "nix",
+                    "--extra-experimental-features",
+                    "nix-command flakes",
+                    "flake",
+                    "metadata",
+                    "--json",
+                    flake_uri,
+                ],
+                check=True,
+                text=True,
+                stdout=subprocess.PIPE,
+            )
+            data = json.loads(res.stdout)
+            flake_store_paths[flake_uri] = data["path"]
+
+    def deploy(h: Host) -> None:
+        target = f"{h.user or 'root'}@{h.host}"
+        flake_store_path = flake_store_paths[str(h.meta.get("flake_uri", ".#"))]
+        flake_path = str(h.meta.get("flake_path", "/etc/nixos"))
+        ssh_arg = f"-p {h.port}" if h.port else ""
+
+        if h.host_key_check != HostKeyCheck.STRICT:
+            ssh_arg += " -o StrictHostKeyChecking=no"
+        if h.host_key_check == HostKeyCheck.NONE:
+            ssh_arg += " -o UserKnownHostsFile=/dev/null"
+
+        ssh_arg += " -i " + h.key if h.key else ""
+
+        h.run_local(
+            f"rsync --checksum -vaF --delete -e 'ssh {ssh_arg}' {flake_store_path}/ {target}:{flake_path}"
+        )
+
+        flake_attr = h.meta.get("flake_attr", "")
+        if flake_attr:
+            flake_attr = "#" + flake_attr
+        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", "switch"]
+            + extra_args
+            + [
+                "--fast",
+                "--option",
+                "keep-going",
+                "true",
+                "--option",
+                "accept-flake-config",
+                "true",
+                "--build-host",
+                "",
+                "--flake",
+                f"{flake_path}{flake_attr}",
+            ]
+        )
+        if target_host:
+            cmd.extend(["--target-host", target_host])
+        ret = h.run(cmd, check=False)
+        # re-retry switch if the first time fails
+        if ret.returncode != 0:
+            ret = h.run(cmd)
+
+    hosts.run_function(deploy)
+
+
+# FIXME: we want some kind of inventory here.
+def update(args: argparse.Namespace) -> None:
+    deploy_nixos(
+        HostGroup(
+            [Host(args.host, user=args.user, meta=dict(flake_attr=args.flake_attr))]
+        )
+    )
+
+
+def register_parser(parser: argparse.ArgumentParser) -> None:
+    parser.add_mutually_exclusive_group(required=True)
+    # TODO pass all args we don't parse into ssh_args, currently it fails if arg starts with -
+    parser.add_argument("--flake-uri", type=str, default=".#", desc="nix flake uri")
+    parser.add_argument(
+        "--flake-attr", type=str, description="nixos configuration in the flake"
+    )
+    parser.add_argument("--user", type=str, default="root")
+    parser.add_argument("host", type=str)
+    parser.set_defaults(func=update)

pkgs/clan-cli/default.nix

@@ -19,6 +19,7 @@
 , stdenv
 , wheel
 , zerotierone
+, rsync
 }:
 let
   dependencies = [ argcomplete jsonschema ];
@@ -63,12 +64,12 @@ python3.pkgs.buildPythonPackage {
     '';
     clan-pytest = runCommand "clan-tests"
       {
-        nativeBuildInputs = [ age zerotierone bubblewrap sops nix openssh stdenv.cc ];
+        nativeBuildInputs = [ age zerotierone bubblewrap sops nix openssh rsync stdenv.cc ];
       } ''
       cp -r ${source} ./src
       chmod +w -R ./src
       cd ./src
-      ${checkPython}/bin/python -m pytest ./tests
+      NIX_STATE_DIR=$TMPDIR/nix ${checkPython}/bin/python -m pytest -s ./tests
       touch $out
     '';
   };

pkgs/clan-cli/flake-module.nix

@@ -20,6 +20,7 @@
         zbar
         tor
         age
+        rsync
         sops;
       # Override license so that we can build zerotierone without
       # having to re-import nixpkgs.

pkgs/clan-cli/tests/clan_flake.py

@@ -15,7 +15,7 @@ def clan_flake(temporary_dir: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator
 {
     description = "A flake for testing clan";
     inputs = {};
-    outputs = { self, nixpkgs }: {};
+    outputs = { self }: {};
 }
 """
     )

pkgs/clan-cli/tests/sshd.py

@@ -5,7 +5,7 @@ import time
 from pathlib import Path
 from sys import platform
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Iterator, Optional
+from typing import TYPE_CHECKING, Iterator
 
 import pytest
 
@@ -22,8 +22,11 @@ class Sshd:
 
 
 class SshdConfig:
-    def __init__(self, path: str, key: str, preload_lib: Optional[str]) -> None:
+    def __init__(
+        self, path: Path, login_shell: Path, key: str, preload_lib: Path
+    ) -> None:
         self.path = path
+        self.login_shell = login_shell
         self.key = key
         self.preload_lib = preload_lib
 
@@ -53,16 +56,39 @@ def sshd_config(project_root: Path, test_root: Path) -> Iterator[SshdConfig]:
         HostKey {host_key}
         LogLevel DEBUG3
         # In the nix build sandbox we don't get any meaningful PATH after login
-        SetEnv PATH={os.environ.get("PATH", "")}
         MaxStartups 64:30:256
         AuthorizedKeysFile {host_key}.pub
+        AcceptEnv REALPATH
         """
         )
+        login_shell = dir / "shell"
+
+        bash = shutil.which("bash")
+        path = os.environ["PATH"]
+        assert bash is not None
+
+        login_shell.write_text(
+            f"""#!{bash}
+if [[ -f /etc/profile ]]; then
+  source /etc/profile
+fi
+if [[ -n "$REALPATH" ]]; then
+   export PATH="$REALPATH:${path}"
+else
+   export PATH="${path}"
+fi
+exec {bash} -l "${{@}}"
+        """
+        )
+        login_shell.chmod(0o755)
 
         lib_path = None
-        if platform == "linux":
+        assert (
+            platform == "linux"
+        ), "we do not support the ld_preload trick on non-linux just now"
+
         # This enforces a login shell by overriding the login shell of `getpwnam(3)`
-            lib_path = str(dir / "libgetpwnam-preload.so")
+        lib_path = dir / "libgetpwnam-preload.so"
         subprocess.run(
             [
                 os.environ.get("CC", "cc"),
@@ -74,7 +100,7 @@ def sshd_config(project_root: Path, test_root: Path) -> Iterator[SshdConfig]:
             check=True,
         )
 
-        yield SshdConfig(str(sshd_config), str(host_key), lib_path)
+        yield SshdConfig(sshd_config, login_shell, str(host_key), lib_path)
 
 
 @pytest.fixture
@@ -83,12 +109,12 @@ def sshd(sshd_config: SshdConfig, command: "Command", ports: "Ports") -> Iterato
     sshd = shutil.which("sshd")
     assert sshd is not None, "no sshd binary found"
     env = {}
-    if sshd_config.preload_lib is not None:
+    env = dict(
-        bash = shutil.which("bash")
+        LD_PRELOAD=str(sshd_config.preload_lib),
-        assert bash is not None
+        LOGIN_SHELL=str(sshd_config.login_shell),
-        env = dict(LD_PRELOAD=str(sshd_config.preload_lib), LOGIN_SHELL=bash)
+    )
     proc = command.run(
-        [sshd, "-f", sshd_config.path, "-D", "-p", str(port)], extra_env=env
+        [sshd, "-f", str(sshd_config.path), "-D", "-p", str(port)], extra_env=env
     )
 
     while True:

pkgs/clan-cli/tests/test_update.py

@@ -0,0 +1,35 @@
+import os
+import shutil
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from environment import mock_env
+from host_group import HostGroup
+
+from clan_cli.update import deploy_nixos
+
+
+def test_update(clan_flake: Path, host_group: HostGroup) -> None:
+    assert len(host_group.hosts) == 1
+    host = host_group.hosts[0]
+
+    with TemporaryDirectory() as tmpdir:
+        host.meta["flake_uri"] = clan_flake
+        host.meta["flake_path"] = str(Path(tmpdir) / "rsync-target")
+        host.ssh_options["SendEnv"] = "REALPATH"
+        bin = Path(tmpdir).joinpath("bin")
+        bin.mkdir()
+        nixos_rebuild = bin.joinpath("nixos-rebuild")
+        bash = shutil.which("bash")
+        assert bash is not None
+        nixos_rebuild.write_text(
+            f"""#!{bash}
+exit 0
+"""
+        )
+        nixos_rebuild.chmod(0o755)
+        path = f"{tmpdir}/bin:{os.environ['PATH']}"
+        nix_state_dir = Path(tmpdir).joinpath("nix")
+        nix_state_dir.mkdir()
+        with mock_env(REALPATH=path):
+            deploy_nixos(host_group)