add sudo_askpass_proxy

pkgs/clan-cli/clan_cli/machines/hardware.py

@@ -105,8 +105,8 @@ def generate_machine_hardware_info(opts: HardwareGenerateOptions) -> HardwareCon
 
     host = opts.machine.target_host()
 
-    with host.ssh_control_master() as ssh:
-        out = ssh.run(config_command, become_root=True, opts=RunOpts(check=False))
+    with host.ssh_control_master() as ssh, ssh.become_root() as sudo_ssh:
+        out = sudo_ssh.run(config_command, opts=RunOpts(check=False))
     if out.returncode != 0:
         if "nixos-facter" in out.stderr and "not found" in out.stderr:
             machine.error(str(out.stderr))

pkgs/clan-cli/clan_cli/machines/update.py

@@ -5,6 +5,7 @@ import os
 import re
 import shlex
 import sys
+from contextlib import ExitStack
 
 from clan_lib.api import API
 from clan_lib.async_run import AsyncContext, AsyncOpts, AsyncRuntime, is_async_cancelled
@@ -109,94 +110,98 @@ def upload_sources(machine: Machine, host: Remote) -> str:
 
 @API.register
 def deploy_machine(machine: Machine) -> None:
-    target_host = machine.target_host()
-    build_host = machine.build_host()
+    with ExitStack() as stack:
+        target_host = stack.enter_context(machine.target_host().ssh_control_master())
+        build_host = machine.build_host()
+        if build_host is not None:
+            build_host = stack.enter_context(build_host.ssh_control_master())
 
-    host = build_host or target_host
+        host = build_host or target_host
 
-    generate_facts([machine], service=None, regenerate=False)
-    generate_vars([machine], generator_name=None, regenerate=False)
+        sudo_host = stack.enter_context(target_host.become_root())
 
-    upload_secrets(machine)
-    upload_secret_vars(machine, target_host)
+        generate_facts([machine], service=None, regenerate=False)
+        generate_vars([machine], generator_name=None, regenerate=False)
 
-    path = upload_sources(machine, host)
+        upload_secrets(machine, sudo_host)
+        upload_secret_vars(machine, sudo_host)
 
-    nix_options = [
-        "--show-trace",
-        "--option",
-        "keep-going",
-        "true",
-        "--option",
-        "accept-flake-config",
-        "true",
-        "-L",
-        *machine.nix_options,
-        "--flake",
-        f"{path}#{machine.name}",
-    ]
+        path = upload_sources(machine, sudo_host)
 
-    become_root = True
-
-    if machine._class_ == "nixos":
-        nix_options += [
-            "--fast",
-            "--build-host",
-            "",
+        nix_options = [
+            "--show-trace",
+            "--option",
+            "keep-going",
+            "true",
+            "--option",
+            "accept-flake-config",
+            "true",
+            "-L",
+            *machine.nix_options,
+            "--flake",
+            f"{path}#{machine.name}",
         ]
 
-        if build_host:
-            become_root = False
-            nix_options += ["--target-host", target_host.target]
+        become_root = True
 
-            if target_host.user != "root":
-                nix_options += ["--use-remote-sudo"]
-        switch_cmd = ["nixos-rebuild", "switch", *nix_options]
-    elif machine._class_ == "darwin":
-        # use absolute path to darwin-rebuild
-        switch_cmd = [
-            "/run/current-system/sw/bin/darwin-rebuild",
-            "switch",
-            *nix_options,
-        ]
+        if machine._class_ == "nixos":
+            nix_options += [
+                "--fast",
+                "--build-host",
+                "",
+            ]
 
-    remote_env = host.nix_ssh_env(control_master=False)
-    ret = host.run(
-        switch_cmd,
-        RunOpts(
-            check=False,
-            log=Log.BOTH,
-            msg_color=MsgColor(stderr=AnsiColor.DEFAULT),
-            needs_user_terminal=True,
-        ),
-        extra_env=remote_env,
-        become_root=become_root,
-        control_master=False,
-    )
+            if build_host:
+                become_root = False
+                nix_options += ["--target-host", target_host.target]
 
-    if is_async_cancelled():
-        return
+                if target_host.user != "root":
+                    nix_options += ["--use-remote-sudo"]
+            switch_cmd = ["nixos-rebuild", "switch", *nix_options]
+        elif machine._class_ == "darwin":
+            # use absolute path to darwin-rebuild
+            switch_cmd = [
+                "/run/current-system/sw/bin/darwin-rebuild",
+                "switch",
+                *nix_options,
+            ]
 
-    # retry nixos-rebuild switch if the first attempt failed
-    if ret.returncode != 0:
-        is_mobile = machine.deployment.get("nixosMobileWorkaround", False)
-        # if the machine is mobile, we retry to deploy with the mobile workaround method
-        if is_mobile:
-            machine.info(
-                "Mobile machine detected, applying workaround deployment method"
-            )
+        if become_root:
+            host = sudo_host
+
+        remote_env = host.nix_ssh_env(control_master=False)
         ret = host.run(
-            ["nixos--rebuild", "test", *nix_options] if is_mobile else switch_cmd,
+            switch_cmd,
             RunOpts(
+                check=False,
                 log=Log.BOTH,
                 msg_color=MsgColor(stderr=AnsiColor.DEFAULT),
                 needs_user_terminal=True,
             ),
             extra_env=remote_env,
-            become_root=become_root,
-            control_master=False,
         )
 
+        if is_async_cancelled():
+            return
+
+        # retry nixos-rebuild switch if the first attempt failed
+        if ret.returncode != 0:
+            is_mobile = machine.deployment.get("nixosMobileWorkaround", False)
+            # if the machine is mobile, we retry to deploy with the mobile workaround method
+            if is_mobile:
+                machine.info(
+                    "Mobile machine detected, applying workaround deployment method"
+                )
+            ret = host.run(
+                ["nixos--rebuild", "test", *nix_options] if is_mobile else switch_cmd,
+                RunOpts(
+                    log=Log.BOTH,
+                    msg_color=MsgColor(stderr=AnsiColor.DEFAULT),
+                    needs_user_terminal=True,
+                ),
+                extra_env=remote_env,
+            )
+
 
 def deploy_machines(machines: list[Machine]) -> None:
     """

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

@@ -98,8 +98,12 @@ def upload(
             raise ClanError(msg)
 
         # TODO accept `input` to be  an IO object instead of bytes so that we don't have to read the tarfile into memory.
-        with tar_path.open("rb") as f, host.ssh_control_master() as ssh:
-            ssh.run(
+        with (
+            tar_path.open("rb") as f,
+            host.ssh_control_master() as ssh,
+            ssh.become_root() as sudo_ssh,
+        ):
+            sudo_ssh.run(
                 [
                     "bash",
                     "-c",
@@ -114,5 +118,4 @@ def upload(
                     prefix=host.command_prefix,
                     needs_user_terminal=True,
                 ),
-                become_root=True,
             )

pkgs/clan-cli/clan_cli/tests/sshd.py

@@ -80,6 +80,7 @@ exec {bash} -l "${{@}}"
 
         fake_sudo.write_text(
             f"""#!{bash}
+shift
 exec "${{@}}"
         """
         )

pkgs/clan-cli/clan_cli/tests/test_ssh_remote.py

@@ -9,6 +9,7 @@ from clan_lib.async_run import AsyncRuntime
 from clan_lib.cmd import ClanCmdTimeoutError, Log, RunOpts
 from clan_lib.errors import ClanError, CmdOut
 from clan_lib.ssh.remote import Remote
+from clan_lib.ssh.sudo_askpass_proxy import SudoAskpassProxy
 
 if sys.platform == "darwin":
     pytest.skip("preload doesn't work on darwin", allow_module_level=True)
@@ -178,6 +179,23 @@ def test_run_no_shell(hosts: list[Remote], runtime: AsyncRuntime) -> None:
     assert proc.wait().result.stdout == "hello\n"
 
 
+def test_sudo_ask_proxy(hosts: list[Remote]) -> None:
+    host = hosts[0]
+    with host.ssh_control_master() as host:
+        proxy = SudoAskpassProxy(host, prompt_command=["bash", "-c", "echo yes"])
+
+        try:
+            askpass_path = proxy.run()
+            out = host.run(
+                ["bash", "-c", askpass_path],
+                opts=RunOpts(check=False, log=Log.BOTH),
+            )
+            assert out.returncode == 0
+            assert out.stdout == "yes\n"
+        finally:
+            proxy.cleanup()
+
+
 def test_run_function(hosts: list[Remote], runtime: AsyncRuntime) -> None:
     def some_func(h: Remote) -> bool:
         with h.ssh_control_master() as ssh: