nixos-rebuild-ng: implement --target-host (#359097)

pkgs/by-name/ni/nixos-rebuild-ng/README.md

@@ -76,7 +76,7 @@ can run:
 # run program
 python -m nixos_rebuild
 # run tests
-python -m pytest
+pytest
 # check types
 mypy .
 # fix lint issues
@@ -119,7 +119,7 @@ ruff format .
 ## TODO
 
 - [ ] Remote host/builders (via SSH)
-- [ ] Improve nix arguments handling (e.g.: `nixFlags` vs `copyFlags` in the
+- [x] Improve nix arguments handling (e.g.: `nixFlags` vs `copyFlags` in the
   old `nixos-rebuild`)
 - [ ] `_NIXOS_REBUILD_EXEC`
 - [ ] Port `nixos-rebuild.passthru.tests`

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py

@@ -1,74 +1,149 @@
-from __future__ import annotations
-
 import argparse
+import atexit
 import json
 import os
 import sys
+from pathlib import Path
 from subprocess import run
+from tempfile import TemporaryDirectory
 from typing import assert_never
 
+from . import nix
 from .models import Action, Flake, NRError, Profile
-from .nix import (
+from .process import Remote, cleanup_ssh
-    edit,
-    find_file,
-    get_nixpkgs_rev,
-    list_generations,
-    nixos_build,
-    nixos_build_flake,
-    rollback,
-    rollback_temporary_profile,
-    set_profile,
-    switch_to_configuration,
-    upgrade_channels,
-)
 from .utils import info
 
-VERBOSE = False
+VERBOSE = 0
 
 
-def parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
+def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentParser]]:
-    parser = argparse.ArgumentParser(
+    common_flags = argparse.ArgumentParser(add_help=False)
+    common_flags.add_argument("--verbose", "-v", action="count", default=0)
+    common_flags.add_argument("--max-jobs", "-j")
+    common_flags.add_argument("--cores")
+    common_flags.add_argument("--log-format")
+    common_flags.add_argument("--keep-going", "-k", action="store_true")
+    common_flags.add_argument("--keep-failed", "-K", action="store_true")
+    common_flags.add_argument("--fallback", action="store_true")
+    common_flags.add_argument("--repair", action="store_true")
+    common_flags.add_argument("--option", nargs=2)
+
+    common_build_flags = argparse.ArgumentParser(add_help=False)
+    common_build_flags.add_argument("--include", "-I")
+    common_build_flags.add_argument("--quiet", action="store_true")
+    common_build_flags.add_argument("--print-build-logs", "-L", action="store_true")
+    common_build_flags.add_argument("--show-trace", action="store_true")
+
+    flake_build_flags = argparse.ArgumentParser(add_help=False)
+    flake_build_flags.add_argument("--accept-flake-config", action="store_true")
+    flake_build_flags.add_argument("--refresh", action="store_true")
+    flake_build_flags.add_argument("--impure", action="store_true")
+    flake_build_flags.add_argument("--offline", action="store_true")
+    flake_build_flags.add_argument("--no-net", action="store_true")
+    flake_build_flags.add_argument("--recreate-lock-file", action="store_true")
+    flake_build_flags.add_argument("--no-update-lock-file", action="store_true")
+    flake_build_flags.add_argument("--no-write-lock-file", action="store_true")
+    flake_build_flags.add_argument("--no-registries", action="store_true")
+    flake_build_flags.add_argument("--commit-lock-file", action="store_true")
+    flake_build_flags.add_argument("--update-input")
+    flake_build_flags.add_argument("--override-input", nargs=2)
+
+    classic_build_flags = argparse.ArgumentParser(add_help=False)
+    classic_build_flags.add_argument("--no-build-output", "-Q", action="store_true")
+
+    copy_flags = argparse.ArgumentParser(add_help=False)
+    copy_flags.add_argument("--use-substitutes", "-s", action="store_true")
+
+    sub_parsers = {
+        "common_flags": common_flags,
+        "common_build_flags": common_build_flags,
+        "flake_build_flags": flake_build_flags,
+        "classic_build_flags": classic_build_flags,
+        "copy_flags": copy_flags,
+    }
+
+    main_parser = argparse.ArgumentParser(
         prog="nixos-rebuild",
+        parents=list(sub_parsers.values()),
         description="Reconfigure a NixOS machine",
         add_help=False,
         allow_abbrev=False,
     )
-    parser.add_argument("--help", action="store_true")
+    main_parser.add_argument("--help", "-h", action="store_true")
-    parser.add_argument("--file", "-f")
+    main_parser.add_argument("--file", "-f")
-    parser.add_argument("--attr", "-A")
+    main_parser.add_argument("--attr", "-A")
-    parser.add_argument("--flake", nargs="?", const=True)
+    main_parser.add_argument("--flake", nargs="?", const=True)
-    parser.add_argument("--no-flake", dest="flake", action="store_false")
+    main_parser.add_argument("--no-flake", dest="flake", action="store_false")
-    parser.add_argument("--install-bootloader", action="store_true")
+    main_parser.add_argument("--install-bootloader", action="store_true")
-    # TODO: add deprecated=True in Python >=3.13
+    main_parser.add_argument("--install-grub", action="store_true")  # deprecated
-    parser.add_argument("--install-grub", action="store_true")
+    main_parser.add_argument("--profile-name", "-p", default="system")
-    parser.add_argument("--profile-name", "-p", default="system")
+    main_parser.add_argument("--specialisation", "-c")
-    parser.add_argument("--specialisation", "-c")
+    main_parser.add_argument("--rollback", action="store_true")
-    parser.add_argument("--rollback", action="store_true")
+    main_parser.add_argument("--upgrade", action="store_true")
-    parser.add_argument("--upgrade", action="store_true")
+    main_parser.add_argument("--upgrade-all", action="store_true")
-    parser.add_argument("--upgrade-all", action="store_true")
+    main_parser.add_argument("--json", action="store_true")
-    parser.add_argument("--json", action="store_true")
+    main_parser.add_argument("--sudo", action="store_true")
-    parser.add_argument("action", choices=Action.values(), nargs="?")
+    main_parser.add_argument("--ask-sudo-password", action="store_true")
+    main_parser.add_argument("--use-remote-sudo", action="store_true")  # deprecated
+    main_parser.add_argument("--no-ssh-tty", action="store_true")  # deprecated
+    # parser.add_argument("--build-host")  # TODO
+    main_parser.add_argument("--target-host")
+    main_parser.add_argument("action", choices=Action.values(), nargs="?")
 
-    args, remainder = parser.parse_known_args(argv[1:])
+    return main_parser, sub_parsers
+
+
+def parse_args(
+    argv: list[str],
+) -> tuple[argparse.Namespace, dict[str, argparse.Namespace]]:
+    parser, sub_parsers = get_parser()
+    args = parser.parse_args(argv[1:])
+    args_groups = {
+        group: parser.parse_known_args(argv[1:])[0]
+        for group, parser in sub_parsers.items()
+    }
+
+    def parser_warn(msg: str) -> None:
+        info(f"{parser.prog}: warning: {msg}")
 
     global VERBOSE
-    # Manually parse verbose flag since this is a nix flag that also affect
+    # This flag affects both nix and this script
-    # the script
+    VERBOSE = args.verbose
-    VERBOSE = any(v == "--verbose" or v.startswith("-v") for v in remainder)
 
     # https://github.com/NixOS/nixpkgs/blob/master/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh#L56
     if args.action == Action.DRY_RUN.value:
         args.action = Action.DRY_BUILD.value
 
+    if args.ask_sudo_password:
+        args.sudo = True
+
+    # TODO: use deprecated=True in Python >=3.13
     if args.install_grub:
-        info(
+        parser_warn("--install-grub deprecated, use --install-bootloader instead")
-            f"{parser.prog}: warning: --install-grub deprecated, use --install-bootloader instead"
-        )
         args.install_bootloader = True
 
+    # TODO: use deprecated=True in Python >=3.13
+    if args.use_remote_sudo:
+        parser_warn("--use-remote-sudo deprecated, use --sudo instead")
+        args.sudo = True
+
+    # TODO: use deprecated=True in Python >=3.13
+    if args.no_ssh_tty:
+        parser_warn("--no-ssh-tty deprecated, SSH's TTY is never used anymore")
+
     if args.action == Action.EDIT.value and (args.file or args.attr):
         parser.error("--file and --attr are not supported with 'edit'")
 
+    if args.target_host and args.action not in (
+        Action.SWITCH.value,
+        Action.BOOT.value,
+        Action.TEST.value,
+        Action.BUILD.value,
+        Action.DRY_BUILD.value,
+        Action.DRY_ACTIVATE.value,
+    ):
+        parser.error(f"--target-host is not supported with '{args.action}'")
+
     if args.flake and (args.file or args.attr):
         parser.error("--flake cannot be used with --file or --attr")
 
@@ -76,17 +151,29 @@ def parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
         r = run(["man", "8", "nixos-rebuild"], check=False)
         parser.exit(r.returncode)
 
-    return args, remainder
+    return args, args_groups
 
 
 def execute(argv: list[str]) -> None:
-    args, nix_flags = parse_args(argv)
+    args, args_groups = parse_args(argv)
+
+    common_flags = vars(args_groups["common_flags"])
+    common_build_flags = common_flags | vars(args_groups["common_build_flags"])
+    build_flags = common_build_flags | vars(args_groups["classic_build_flags"])
+    flake_build_flags = common_build_flags | vars(args_groups["flake_build_flags"])
+    copy_flags = common_flags | vars(args_groups["copy_flags"])
+
+    # Will be cleaned up on exit automatically.
+    tmpdir = TemporaryDirectory(prefix="nixos-rebuild.")
+    tmpdir_path = Path(tmpdir.name)
+    atexit.register(cleanup_ssh, tmpdir_path)
 
     profile = Profile.from_name(args.profile_name)
-    flake = Flake.from_arg(args.flake)
+    target_host = Remote.from_arg(args.target_host, args.ask_sudo_password, tmpdir_path)
+    flake = Flake.from_arg(args.flake, target_host)
 
     if args.upgrade or args.upgrade_all:
-        upgrade_channels(bool(args.upgrade_all))
+        nix.upgrade_channels(bool(args.upgrade_all))
 
     action = Action(args.action)
     # Only run shell scripts from the Nixpkgs tree if the action is
@@ -96,8 +183,8 @@ def execute(argv: list[str]) -> None:
     # untrusted tree.
     can_run = action in (Action.SWITCH, Action.BOOT, Action.TEST)
     if can_run and not flake:
-        nixpkgs_path = find_file("nixpkgs", nix_flags)
+        nixpkgs_path = nix.find_file("nixpkgs", **build_flags)
-        rev = get_nixpkgs_rev(nixpkgs_path)
+        rev = nix.get_nixpkgs_rev(nixpkgs_path)
         if nixpkgs_path and rev:
             (nixpkgs_path / ".version-suffix").write_text(rev)
 
@@ -105,88 +192,99 @@ def execute(argv: list[str]) -> None:
         case Action.SWITCH | Action.BOOT:
             info("building the system configuration...")
             if args.rollback:
-                path_to_config = rollback(profile)
+                path_to_config = nix.rollback(profile, target_host, sudo=args.sudo)
-            elif flake:
-                path_to_config = nixos_build_flake(
-                    "toplevel",
-                    flake,
-                    nix_flags,
-                    no_link=True,
-                )
-                set_profile(profile, path_to_config)
             else:
-                path_to_config = nixos_build(
+                if flake:
-                    "system",
+                    path_to_config = nix.nixos_build_flake(
-                    args.attr,
+                        "toplevel",
-                    args.file,
+                        flake,
-                    nix_flags,
+                        no_link=True,
-                    no_out_link=True,
+                        **flake_build_flags,
-                )
+                    )
-                set_profile(profile, path_to_config)
+                else:
-            switch_to_configuration(
+                    path_to_config = nix.nixos_build(
+                        "system",
+                        args.attr,
+                        args.file,
+                        no_out_link=True,
+                        **build_flags,
+                    )
+                nix.copy_closure(path_to_config, target_host, **copy_flags)
+                nix.set_profile(profile, path_to_config, target_host, sudo=args.sudo)
+            nix.switch_to_configuration(
                 path_to_config,
                 action,
+                target_host,
+                sudo=args.sudo,
                 specialisation=args.specialisation,
                 install_bootloader=args.install_bootloader,
             )
         case Action.TEST | Action.BUILD | Action.DRY_BUILD | Action.DRY_ACTIVATE:
             info("building the system configuration...")
             dry_run = action == Action.DRY_BUILD
-            if args.rollback and action in (Action.TEST, Action.BUILD):
+            if args.rollback:
-                maybe_path_to_config = rollback_temporary_profile(profile)
+                if action not in (Action.TEST, Action.BUILD):
+                    raise NRError(f"--rollback is incompatible with '{action}'")
+                maybe_path_to_config = nix.rollback_temporary_profile(
+                    profile,
+                    target_host,
+                    sudo=args.sudo,
+                )
                 if maybe_path_to_config:  # kinda silly but this makes mypy happy
                     path_to_config = maybe_path_to_config
                 else:
                     raise NRError("could not find previous generation")
             elif flake:
-                path_to_config = nixos_build_flake(
+                path_to_config = nix.nixos_build_flake(
                     "toplevel",
                     flake,
-                    nix_flags,
                     keep_going=True,
                     dry_run=dry_run,
+                    **flake_build_flags,
                 )
             else:
-                path_to_config = nixos_build(
+                path_to_config = nix.nixos_build(
                     "system",
                     args.attr,
                     args.file,
-                    nix_flags,
                     keep_going=True,
                     dry_run=dry_run,
+                    **build_flags,
                 )
             if action in (Action.TEST, Action.DRY_ACTIVATE):
-                switch_to_configuration(
+                nix.switch_to_configuration(
                     path_to_config,
                     action,
+                    target_host,
+                    sudo=args.sudo,
                     specialisation=args.specialisation,
                 )
         case Action.BUILD_VM | Action.BUILD_VM_WITH_BOOTLOADER:
             info("building the system configuration...")
             attr = "vm" if action == Action.BUILD_VM else "vmWithBootLoader"
             if flake:
-                path_to_config = nixos_build_flake(
+                path_to_config = nix.nixos_build_flake(
                     attr,
                     flake,
-                    nix_flags,
                     keep_going=True,
+                    **flake_build_flags,
                 )
             else:
-                path_to_config = nixos_build(
+                path_to_config = nix.nixos_build(
                     attr,
                     args.attr,
                     args.file,
-                    nix_flags,
                     keep_going=True,
+                    **build_flags,
                 )
             vm_path = next(path_to_config.glob("bin/run-*-vm"), "./result/bin/run-*-vm")
             print(f"Done. The virtual machine can be started by running '{vm_path}'")
         case Action.EDIT:
-            edit(flake, nix_flags)
+            nix.edit(flake, **flake_build_flags)
         case Action.DRY_RUN:
             assert False, "DRY_RUN should be a DRY_BUILD alias"
         case Action.LIST_GENERATIONS:
-            generations = list_generations(profile)
+            generations = nix.list_generations(profile)
             if args.json:
                 print(json.dumps(generations, indent=2))
             else:

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py

@@ -2,10 +2,13 @@ from __future__ import annotations
 
 import platform
 import re
+import subprocess
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Any, ClassVar, TypedDict, override
+from typing import Any, Callable, ClassVar, Self, TypedDict, override
+
+from .process import Remote, run_wrapper
 
 
 class NRError(Exception):
@@ -53,21 +56,37 @@ class Flake:
         return f"{self.path}#{self.attr}"
 
     @classmethod
-    def parse(cls, flake_str: str, hostname: str | None = None) -> Flake:
+    def parse(
+        cls,
+        flake_str: str,
+        hostname_fn: Callable[[], str | None] = lambda: None,
+    ) -> Self:
         m = cls._re.match(flake_str)
         assert m is not None, f"got no matches for {flake_str}"
         attr = m.group("attr")
-        nixos_attr = f"nixosConfigurations.{attr or hostname or "default"}"
+        nixos_attr = f"nixosConfigurations.{attr or hostname_fn() or "default"}"
-        return Flake(Path(m.group("path")), nixos_attr)
+        return cls(Path(m.group("path")), nixos_attr)
 
     @classmethod
-    def from_arg(cls, flake_arg: Any) -> Flake | None:
+    def from_arg(cls, flake_arg: Any, target_host: Remote | None) -> Self | None:
-        hostname = platform.node()
+        def get_hostname() -> str | None:
+            if target_host:
+                try:
+                    return run_wrapper(
+                        ["uname", "-n"],
+                        capture_output=True,
+                        remote=target_host,
+                    ).stdout.strip()
+                except (AttributeError, subprocess.CalledProcessError):
+                    return None
+            else:
+                return platform.node()
+
         match flake_arg:
             case str(s):
-                return cls.parse(s, hostname)
+                return cls.parse(s, get_hostname)
             case True:
-                return cls.parse(".", hostname)
+                return cls.parse(".", get_hostname)
             case False:
                 return None
             case _:
@@ -77,7 +96,7 @@ class Flake:
                     # It can be a symlink to the actual flake.
                     if default_path.is_symlink():
                         default_path = default_path.readlink()
-                    return cls.parse(str(default_path.parent), hostname)
+                    return cls.parse(str(default_path.parent), get_hostname)
                 else:
                     return None
 
@@ -105,12 +124,12 @@ class Profile:
     name: str
     path: Path
 
-    @staticmethod
+    @classmethod
-    def from_name(name: str = "system") -> Profile:
+    def from_name(cls, name: str = "system") -> Self:
         match name:
             case "system":
-                return Profile(name, Path("/nix/var/nix/profiles/system"))
+                return cls(name, Path("/nix/var/nix/profiles/system"))
             case _:
                 path = Path("/nix/var/nix/profiles/system-profiles") / name
                 path.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
-                return Profile(name, path)
+                return cls(name, path)

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py

@@ -1,10 +1,7 @@
-from __future__ import annotations
-
 import os
-import shutil
 from datetime import datetime
 from pathlib import Path
-from subprocess import PIPE, CalledProcessError, run
+from subprocess import PIPE, CalledProcessError
 from typing import Final
 
 from .models import (
@@ -14,27 +11,56 @@ from .models import (
     GenerationJson,
     NRError,
     Profile,
+    Remote,
 )
-from .utils import dict_to_flags, info
+from .process import run_wrapper
+from .utils import Args, dict_to_flags, info
 
 FLAKE_FLAGS: Final = ["--extra-experimental-features", "nix-command flakes"]
 
 
-def edit(flake: Flake | None, nix_flags: list[str] | None = None) -> None:
+def copy_closure(
+    closure: Path,
+    target_host: Remote | None,
+    **copy_flags: Args,
+) -> None:
+    host = target_host
+    if not host:
+        return
+
+    run_wrapper(
+        [
+            "nix-copy-closure",
+            *dict_to_flags(copy_flags),
+            "--to",
+            host.host,
+            closure,
+        ],
+        extra_env={"NIX_SSHOPTS": " ".join(host.opts)},
+    )
+
+
+def edit(flake: Flake | None, **flake_flags: Args) -> None:
     "Try to find and open NixOS configuration file in editor."
     if flake:
-        run(
+        run_wrapper(
-            ["nix", *FLAKE_FLAGS, "edit", *(nix_flags or []), "--", str(flake)],
+            [
+                "nix",
+                *FLAKE_FLAGS,
+                "edit",
+                *dict_to_flags(flake_flags),
+                "--",
+                str(flake),
+            ],
             check=False,
         )
     else:
-        if nix_flags:
+        if flake_flags:
             raise NRError("'edit' does not support extra Nix flags")
         nixos_config = Path(
             os.getenv("NIXOS_CONFIG")
-            or run(
+            or run_wrapper(
                 ["nix-instantiate", "--find-file", "nixos-config"],
-                text=True,
                 stdout=PIPE,
                 check=False,
             ).stdout.strip()
@@ -44,18 +70,17 @@ def edit(flake: Flake | None, nix_flags: list[str] | None = None) -> None:
             nixos_config /= "default.nix"
 
         if nixos_config.exists():
-            run([os.getenv("EDITOR", "nano"), nixos_config], check=False)
+            run_wrapper([os.getenv("EDITOR", "nano"), nixos_config], check=False)
         else:
             raise NRError("cannot find NixOS config file")
 
 
-def find_file(file: str, nix_flags: list[str] | None = None) -> Path | None:
+def find_file(file: str, **nix_flags: Args) -> Path | None:
     "Find classic Nixpkgs location."
-    r = run(
+    r = run_wrapper(
-        ["nix-instantiate", "--find-file", file, *(nix_flags or [])],
+        ["nix-instantiate", "--find-file", file, *dict_to_flags(nix_flags)],
         stdout=PIPE,
         check=False,
-        text=True,
     )
     if r.returncode:
         return None
@@ -69,23 +94,24 @@ def get_nixpkgs_rev(nixpkgs_path: Path | None) -> str | None:
     if not nixpkgs_path:
         return None
 
-    # Git is not included in the closure for nixos-rebuild so we need to check
+    try:
-    if not shutil.which("git"):
+        # Get current revision
+        r = run_wrapper(
+            ["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
+            check=False,
+            stdout=PIPE,
+        )
+    except FileNotFoundError:
+        # Git is not included in the closure so we need to check
         info(f"warning: Git not found; cannot figure out revision of '{nixpkgs_path}'")
         return None
 
-    # Get current revision
+    if rev := r.stdout.strip():
-    r = run(
-        ["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
-        check=False,
-        stdout=PIPE,
-        text=True,
-    )
-    rev = r.stdout.strip()
-
-    if rev:
         # Check if repo is dirty
-        if run(["git", "-C", nixpkgs_path, "diff", "--quiet"], check=False).returncode:
+        if run_wrapper(
+            ["git", "-C", nixpkgs_path, "diff", "--quiet"],
+            check=False,
+        ).returncode:
             rev += "M"
         return f".git.{rev}"
     else:
@@ -120,7 +146,12 @@ def _parse_generation_from_nix_env(line: str) -> Generation:
     )
 
 
-def get_generations(profile: Profile, lock_profile: bool = False) -> list[Generation]:
+def get_generations(
+    profile: Profile,
+    target_host: Remote | None = None,
+    using_nix_env: bool = False,
+    sudo: bool = False,
+) -> list[Generation]:
     """Get all NixOS generations from profile.
 
     Includes generation ID (e.g.: 1, 2), timestamp (e.g.: when it was created)
@@ -132,19 +163,20 @@ def get_generations(profile: Profile, lock_profile: bool = False) -> list[Genera
         raise NRError(f"no profile '{profile.name}' found")
 
     result = []
-    if lock_profile:
+    if using_nix_env:
         # Using `nix-env --list-generations` needs root to lock the profile
         # TODO: do we actually need to lock profile for e.g.: rollback?
         # https://github.com/NixOS/nix/issues/5144
-        r = run(
+        r = run_wrapper(
             ["nix-env", "-p", profile.path, "--list-generations"],
-            text=True,
+            stdout=PIPE,
-            stdout=True,
+            remote=target_host,
-            check=True,
+            sudo=sudo,
         )
         for line in r.stdout.splitlines():
             result.append(_parse_generation_from_nix_env(line))
     else:
+        assert not target_host, "target_host is not supported when using_nix_env=False"
         for p in profile.path.parent.glob("system-*-link"):
             result.append(_parse_generation_from_nix_store(p, profile))
     return sorted(result, key=lambda d: d.id)
@@ -179,11 +211,9 @@ def list_generations(profile: Profile) -> list[GenerationJson]:
             s.name for s in (generation_path / "specialisation").glob("*") if s.is_dir()
         ]
         try:
-            configuration_revision = run(
+            configuration_revision = run_wrapper(
                 [generation_path / "sw/bin/nixos-version", "--configuration-revision"],
                 capture_output=True,
-                check=True,
-                text=True,
             ).stdout.strip()
         except (CalledProcessError, IOError):
             configuration_revision = "Unknown"
@@ -207,8 +237,7 @@ def nixos_build(
     attr: str,
     pre_attr: str | None,
     file: str | None,
-    nix_flags: list[str] | None = None,
+    **nix_flags: Args,
-    **kwargs: bool | str,
 ) -> Path:
     """Build NixOS attribute using classic Nix.
 
@@ -227,16 +256,15 @@ def nixos_build(
         ]
     else:
         run_args = ["nix-build", "<nixpkgs/nixos>", "--attr", attr]
-    run_args += dict_to_flags(kwargs) + (nix_flags or [])
+    run_args += dict_to_flags(nix_flags)
-    r = run(run_args, check=True, text=True, stdout=PIPE)
+    r = run_wrapper(run_args, stdout=PIPE)
     return Path(r.stdout.strip())
 
 
 def nixos_build_flake(
     attr: str,
     flake: Flake,
-    nix_flags: list[str] | None = None,
+    **flake_flags: Args,
-    **kwargs: bool | str,
 ) -> Path:
     """Build NixOS attribute using Flakes.
 
@@ -248,22 +276,35 @@ def nixos_build_flake(
         "build",
         "--print-out-paths",
         f"{flake}.config.system.build.{attr}",
+        *dict_to_flags(flake_flags),
     ]
-    run_args += dict_to_flags(kwargs) + (nix_flags or [])
+    r = run_wrapper(run_args, stdout=PIPE)
-    r = run(run_args, check=True, text=True, stdout=PIPE)
     return Path(r.stdout.strip())
 
 
-def rollback(profile: Profile) -> Path:
+def rollback(profile: Profile, target_host: Remote | None, sudo: bool) -> Path:
     "Rollback Nix profile, like one created by `nixos-rebuild switch`."
-    run(["nix-env", "--rollback", "-p", profile.path], check=True)
+    run_wrapper(
+        ["nix-env", "--rollback", "-p", profile.path],
+        remote=target_host,
+        sudo=sudo,
+    )
     # Rollback config PATH is the own profile
     return profile.path
 
 
-def rollback_temporary_profile(profile: Profile) -> Path | None:
+def rollback_temporary_profile(
+    profile: Profile,
+    target_host: Remote | None,
+    sudo: bool,
+) -> Path | None:
     "Rollback a temporary Nix profile, like one created by `nixos-rebuild test`."
-    generations = get_generations(profile, lock_profile=True)
+    generations = get_generations(
+        profile,
+        target_host=target_host,
+        using_nix_env=True,
+        sudo=sudo,
+    )
     previous_gen_id = None
     for generation in generations:
         if not generation.current:
@@ -275,14 +316,25 @@ def rollback_temporary_profile(profile: Profile) -> Path | None:
         return None
 
 
-def set_profile(profile: Profile, path_to_config: Path) -> None:
+def set_profile(
+    profile: Profile,
+    path_to_config: Path,
+    target_host: Remote | None,
+    sudo: bool,
+) -> None:
     "Set a path as the current active Nix profile."
-    run(["nix-env", "-p", profile.path, "--set", path_to_config], check=True)
+    run_wrapper(
+        ["nix-env", "-p", profile.path, "--set", path_to_config],
+        remote=target_host,
+        sudo=sudo,
+    )
 
 
 def switch_to_configuration(
     path_to_config: Path,
     action: Action,
+    target_host: Remote | None,
+    sudo: bool,
     install_bootloader: bool = False,
     specialisation: str | None = None,
 ) -> None:
@@ -301,13 +353,11 @@ def switch_to_configuration(
         if not path_to_config.exists():
             raise NRError(f"specialisation not found: {specialisation}")
 
-    run(
+    run_wrapper(
         [path_to_config / "bin/switch-to-configuration", str(action)],
-        env={
+        extra_env={"NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0"},
-            "NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0",
+        remote=target_host,
-            "LOCALE_ARCHIVE": os.getenv("LOCALE_ARCHIVE", ""),
+        sudo=sudo,
-        },
-        check=True,
     )
 
 
@@ -323,4 +373,4 @@ def upgrade_channels(all: bool = False) -> None:
             or channel_path.name == "nixos"
             or (channel_path / ".update-on-nixos-rebuild").exists()
         ):
-            run(["nix-channel", "--update", channel_path.name], check=False)
+            run_wrapper(["nix-channel", "--update", channel_path.name], check=False)

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import os
+import subprocess
+from dataclasses import dataclass
+from getpass import getpass
+from pathlib import Path
+from typing import Self, Sequence, TypedDict, Unpack
+
+from .utils import info
+
+
+@dataclass(frozen=True)
+class Remote:
+    host: str
+    opts: list[str]
+    sudo_password: str | None
+
+    @classmethod
+    def from_arg(
+        cls,
+        host: str | None,
+        ask_sudo_password: bool | None,
+        tmp_dir: Path,
+    ) -> Self | None:
+        if not host:
+            return None
+
+        opts = os.getenv("NIX_SSHOPTS", "").split()
+        cls._validate_opts(opts, ask_sudo_password)
+        opts += [
+            # SSH ControlMaster flags, allow for faster re-connection
+            "-o",
+            "ControlMaster=auto",
+            "-o",
+            f"ControlPath={tmp_dir / "ssh-%n"}",
+            "-o",
+            "ControlPersist=60",
+        ]
+        sudo_password = None
+        if ask_sudo_password:
+            sudo_password = getpass(f"[sudo] password for {host}: ")
+        return cls(host, opts, sudo_password)
+
+    @staticmethod
+    def _validate_opts(opts: list[str], ask_sudo_password: bool | None) -> None:
+        for o in opts:
+            if o in ["-t", "-tt", "RequestTTY=yes", "RequestTTY=force"]:
+                info(
+                    f"warning: detected option '{o}' in NIX_SSHOPTS. SSH's TTY "
+                    + "may cause issues, it is recommended to remove this option"
+                )
+                if not ask_sudo_password:
+                    info(
+                        "If you want to prompt for sudo password use "
+                        + "'--ask-sudo-password' option instead"
+                    )
+
+
+# Not exhaustive, but we can always extend it later.
+class RunKwargs(TypedDict, total=False):
+    capture_output: bool
+    stderr: int | None
+    stdout: int | None
+
+
+def cleanup_ssh(tmp_dir: Path) -> None:
+    "Close SSH ControlMaster connection."
+    for ctrl in tmp_dir.glob("ssh-*"):
+        subprocess.run(["ssh", "-o", f"ControlPath={ctrl}", "exit"], check=False)
+
+
+def run_wrapper(
+    args: Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]],
+    *,
+    check: bool = True,
+    extra_env: dict[str, str] | None = None,
+    remote: Remote | None = None,
+    sudo: bool = False,
+    **kwargs: Unpack[RunKwargs],
+) -> subprocess.CompletedProcess[str]:
+    "Wrapper around `subprocess.run` that supports extra functionality."
+    env = None
+    input = None
+    if remote:
+        if extra_env:
+            extra_env_args = [f"{env}={value}" for env, value in extra_env.items()]
+            args = ["env", *extra_env_args, *args]
+        if sudo:
+            if remote.sudo_password:
+                args = ["sudo", "--prompt=", "--stdin", *args]
+                input = remote.sudo_password + "\n"
+            else:
+                args = ["sudo", *args]
+        args = ["ssh", *remote.opts, remote.host, "--", *args]
+    else:
+        if extra_env:
+            env = os.environ | extra_env
+        if sudo:
+            args = ["sudo", *args]
+
+    return subprocess.run(
+        args,
+        check=check,
+        env=env,
+        input=input,
+        # Hope nobody is using NixOS with non-UTF8 encodings, but "surrogateescape"
+        # should still work in those systems.
+        text=True,
+        errors="surrogateescape",
+        **kwargs,
+    )

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py

@@ -1,19 +1,18 @@
-from __future__ import annotations
-
 import sys
 from functools import partial
-from typing import Any
+from typing import TypeAlias
 
 info = partial(print, file=sys.stderr)
+Args: TypeAlias = bool | str | list[str] | int | None
 
 
-def dict_to_flags(d: dict[str, Any]) -> list[str]:
+def dict_to_flags(d: dict[str, Args]) -> list[str]:
     flags = []
     for key, value in d.items():
         flag = f"--{'-'.join(key.split('_'))}"
         match value:
-            case None | False:
+            case None | False | 0 | []:
-                pass
+                continue
             case True:
                 flags.append(flag)
             case int():

pkgs/by-name/ni/nixos-rebuild-ng/src/pyproject.toml

@@ -38,8 +38,6 @@ ignore_missing_imports = true
 extend-select = [
     # ensure imports are sorted
     "I",
-    # require 'from __future__ import annotations'
-    "FA102",
     # require `check` argument for `subprocess.run`
     "PLW1510",
 ]
@@ -48,4 +46,5 @@ extend-select = [
 "tests/" = ["FA102"]
 
 [tool.pytest.ini_options]
+pythonpath = ["."]
 addopts = ["--import-mode=importlib"]

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py

@@ -2,7 +2,7 @@ import textwrap
 from pathlib import Path
 from subprocess import PIPE, CompletedProcess
 from typing import Any
-from unittest.mock import call, patch
+from unittest.mock import ANY, call, patch
 
 import pytest
 
@@ -10,10 +10,12 @@ import nixos_rebuild as nr
 
 from .helpers import get_qualified_name
 
-
+DEFAULT_RUN_KWARGS = {
-@pytest.fixture(autouse=True)
+    "env": ANY,
-def setup(monkeypatch: Any) -> None:
+    "input": None,
-    monkeypatch.setenv("LOCALE_ARCHIVE", "/locale")
+    "text": True,
+    "errors": "surrogateescape",
+}
 
 
 def test_parse_args() -> None:
@@ -29,25 +31,28 @@ def test_parse_args() -> None:
         nr.parse_args(["nixos-rebuild", "edit", "--attr", "attr"])
     assert e.value.code == 2
 
-    r1, remainder = nr.parse_args(
+    r1, g1 = nr.parse_args(
         [
             "nixos-rebuild",
             "switch",
             "--install-grub",
             "--flake",
             "/etc/nixos",
-            "--extra",
+            "--option",
-            "flag",
+            "foo",
+            "bar",
         ]
     )
-    assert remainder == ["--extra", "flag"]
+    assert nr.VERBOSE == 0
     assert r1.flake == "/etc/nixos"
     assert r1.install_bootloader is True
     assert r1.install_grub is True
     assert r1.profile_name == "system"
     assert r1.action == "switch"
+    assert r1.option == ["foo", "bar"]
+    assert g1["common_flags"].option == ["foo", "bar"]
 
-    r2, remainder = nr.parse_args(
+    r2, g2 = nr.parse_args(
         [
             "nixos-rebuild",
             "dry-run",
@@ -57,18 +62,21 @@ def test_parse_args() -> None:
             "foo",
             "--attr",
             "bar",
+            "-vvv",
         ]
     )
-    assert remainder == []
+    assert nr.VERBOSE == 3
+    assert r2.verbose == 3
     assert r2.flake is False
     assert r2.action == "dry-build"
     assert r2.file == "foo"
     assert r2.attr == "bar"
+    assert g2["common_flags"].verbose == 3
 
 
-@patch(get_qualified_name(nr.nix.run, nr.nix), autospec=True)
+@patch.dict(nr.process.os.environ, {}, clear=True)
-@patch(get_qualified_name(nr.nix.shutil.which), autospec=True, return_value="/bin/git")
+@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
-def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> None:
+def test_execute_nix_boot(mock_run: Any, tmp_path: Path) -> None:
     nixpkgs_path = tmp_path / "nixpkgs"
     nixpkgs_path.mkdir()
     config_path = tmp_path / "test"
@@ -88,7 +96,6 @@ def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> Non
 
     nr.execute(["nixos-rebuild", "boot", "--no-flake", "-vvv"])
 
-    assert nr.VERBOSE is True
     assert mock_run.call_count == 6
     mock_run.assert_has_calls(
         [
@@ -96,17 +103,18 @@ def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> Non
                 ["nix-instantiate", "--find-file", "nixpkgs", "-vvv"],
                 stdout=PIPE,
                 check=False,
-                text=True,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
                 ["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
                 check=False,
                 stdout=PIPE,
-                text=True,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
                 ["git", "-C", nixpkgs_path, "diff", "--quiet"],
                 check=False,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
                 [
@@ -118,8 +126,8 @@ def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> Non
                     "-vvv",
                 ],
                 check=True,
-                text=True,
                 stdout=PIPE,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
                 [
@@ -130,17 +138,19 @@ def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> Non
                     config_path,
                 ],
                 check=True,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
                 [config_path / "bin/switch-to-configuration", "boot"],
-                env={"NIXOS_INSTALL_BOOTLOADER": "0", "LOCALE_ARCHIVE": "/locale"},
                 check=True,
+                **(DEFAULT_RUN_KWARGS | {"env": {"NIXOS_INSTALL_BOOTLOADER": "0"}}),
             ),
         ]
     )
 
 
-@patch(get_qualified_name(nr.nix.run, nr.nix), autospec=True)
+@patch.dict(nr.process.os.environ, {}, clear=True)
+@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
 def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
     config_path = tmp_path / "test"
     config_path.touch()
@@ -160,11 +170,11 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
             "--flake",
             "/path/to/config#hostname",
             "--install-bootloader",
+            "--sudo",
             "--verbose",
         ]
     )
 
-    assert nr.VERBOSE is True
     assert mock_run.call_count == 3
     mock_run.assert_has_calls(
         [
@@ -177,14 +187,15 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
                     "--print-out-paths",
                     "/path/to/config#nixosConfigurations.hostname.config.system.build.toplevel",
                     "--no-link",
-                    "--verbose",
+                    "-v",
                 ],
                 check=True,
-                text=True,
                 stdout=PIPE,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
                 [
+                    "sudo",
                     "nix-env",
                     "-p",
                     Path("/nix/var/nix/profiles/system"),
@@ -192,25 +203,126 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
                     config_path,
                 ],
                 check=True,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
-                [config_path / "bin/switch-to-configuration", "switch"],
+                ["sudo", config_path / "bin/switch-to-configuration", "switch"],
-                env={"NIXOS_INSTALL_BOOTLOADER": "1", "LOCALE_ARCHIVE": "/locale"},
                 check=True,
+                **(DEFAULT_RUN_KWARGS | {"env": {"NIXOS_INSTALL_BOOTLOADER": "1"}}),
             ),
         ]
     )
 
 
-@patch(get_qualified_name(nr.nix.run, nr.nix), autospec=True)
+@patch.dict(nr.process.os.environ, {}, clear=True)
+@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
+@patch(get_qualified_name(nr.TemporaryDirectory, nr))  # can't autospec
+def test_execute_nix_switch_flake_remote(
+    mock_tmpdir: Any,
+    mock_run: Any,
+    tmp_path: Path,
+) -> None:
+    config_path = tmp_path / "test"
+    config_path.touch()
+    mock_run.side_effect = [
+        # nixos_build_flake
+        CompletedProcess([], 0, str(config_path)),
+        # set_profile
+        CompletedProcess([], 0),
+        # copy_closure
+        CompletedProcess([], 0),
+        # switch_to_configuration
+        CompletedProcess([], 0),
+    ]
+    mock_tmpdir.return_value.name = "/tmp/test"
+
+    nr.execute(
+        [
+            "nixos-rebuild",
+            "switch",
+            "--flake",
+            "/path/to/config#hostname",
+            "--use-remote-sudo",
+            "--target-host",
+            "user@localhost",
+        ]
+    )
+
+    assert mock_run.call_count == 4
+    mock_run.assert_has_calls(
+        [
+            call(
+                [
+                    "nix",
+                    "--extra-experimental-features",
+                    "nix-command flakes",
+                    "build",
+                    "--print-out-paths",
+                    "/path/to/config#nixosConfigurations.hostname.config.system.build.toplevel",
+                    "--no-link",
+                ],
+                check=True,
+                stdout=PIPE,
+                **DEFAULT_RUN_KWARGS,
+            ),
+            call(
+                ["nix-copy-closure", "--to", "user@localhost", config_path],
+                check=True,
+                **DEFAULT_RUN_KWARGS,
+            ),
+            call(
+                [
+                    "ssh",
+                    "-o",
+                    "ControlMaster=auto",
+                    "-o",
+                    "ControlPath=/tmp/test/ssh-%n",
+                    "-o",
+                    "ControlPersist=60",
+                    "user@localhost",
+                    "--",
+                    "sudo",
+                    "nix-env",
+                    "-p",
+                    Path("/nix/var/nix/profiles/system"),
+                    "--set",
+                    config_path,
+                ],
+                check=True,
+                **DEFAULT_RUN_KWARGS,
+            ),
+            call(
+                [
+                    "ssh",
+                    "-o",
+                    "ControlMaster=auto",
+                    "-o",
+                    "ControlPath=/tmp/test/ssh-%n",
+                    "-o",
+                    "ControlPersist=60",
+                    "user@localhost",
+                    "--",
+                    "sudo",
+                    "env",
+                    "NIXOS_INSTALL_BOOTLOADER=0",
+                    config_path / "bin/switch-to-configuration",
+                    "switch",
+                ],
+                check=True,
+                **DEFAULT_RUN_KWARGS,
+            ),
+        ]
+    )
+
+
+@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
 def test_execute_switch_rollback(mock_run: Any, tmp_path: Path) -> None:
     nixpkgs_path = tmp_path / "nixpkgs"
     nixpkgs_path.touch()
 
     nr.execute(["nixos-rebuild", "switch", "--rollback", "--install-bootloader"])
 
-    assert nr.VERBOSE is False
+    assert mock_run.call_count >= 2
-    assert mock_run.call_count == 3
     # ignoring update_nixpkgs_rev calls
     mock_run.assert_has_calls(
         [
@@ -222,20 +334,21 @@ def test_execute_switch_rollback(mock_run: Any, tmp_path: Path) -> None:
                     Path("/nix/var/nix/profiles/system"),
                 ],
                 check=True,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
                 [
                     Path("/nix/var/nix/profiles/system/bin/switch-to-configuration"),
                     "switch",
                 ],
-                env={"NIXOS_INSTALL_BOOTLOADER": "1", "LOCALE_ARCHIVE": "/locale"},
                 check=True,
+                **DEFAULT_RUN_KWARGS,
             ),
         ]
     )
 
 
-@patch(get_qualified_name(nr.nix.run, nr.nix), autospec=True)
+@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
 @patch(get_qualified_name(nr.nix.Path.exists, nr.nix), autospec=True, return_value=True)
 @patch(get_qualified_name(nr.nix.Path.mkdir, nr.nix), autospec=True)
 def test_execute_test_rollback(
@@ -268,7 +381,6 @@ def test_execute_test_rollback(
         ]
     )
 
-    assert nr.VERBOSE is False
     assert mock_run.call_count == 2
     mock_run.assert_has_calls(
         [
@@ -279,9 +391,9 @@ def test_execute_test_rollback(
                     Path("/nix/var/nix/profiles/system-profiles/foo"),
                     "--list-generations",
                 ],
-                text=True,
-                stdout=True,
                 check=True,
+                stdout=PIPE,
+                **DEFAULT_RUN_KWARGS,
             ),
             call(
                 [
@@ -290,8 +402,8 @@ def test_execute_test_rollback(
                     ),
                     "test",
                 ],
-                env={"NIXOS_INSTALL_BOOTLOADER": "0", "LOCALE_ARCHIVE": "/locale"},
                 check=True,
+                **DEFAULT_RUN_KWARGS,
             ),
         ]
     )

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_models.py

@@ -1,9 +1,10 @@
 import platform
+import subprocess
 from pathlib import Path
 from typing import Any
 from unittest.mock import patch
 
-from nixos_rebuild import models as m
+import nixos_rebuild.models as m
 
 from .helpers import get_qualified_name
 
@@ -12,16 +13,15 @@ def test_flake_parse() -> None:
     assert m.Flake.parse("/path/to/flake#attr") == m.Flake(
         Path("/path/to/flake"), "nixosConfigurations.attr"
     )
-    assert m.Flake.parse("/path/ to /flake", "hostname") == m.Flake(
+    assert m.Flake.parse("/path/ to /flake", lambda: "hostname") == m.Flake(
         Path("/path/ to /flake"), "nixosConfigurations.hostname"
     )
-    assert m.Flake.parse("/path/to/flake", "hostname") == m.Flake(
+    assert m.Flake.parse("/path/to/flake", lambda: "hostname") == m.Flake(
         Path("/path/to/flake"), "nixosConfigurations.hostname"
     )
     assert m.Flake.parse(".#attr") == m.Flake(Path("."), "nixosConfigurations.attr")
     assert m.Flake.parse("#attr") == m.Flake(Path("."), "nixosConfigurations.attr")
-    assert m.Flake.parse(".", None) == m.Flake(Path("."), "nixosConfigurations.default")
+    assert m.Flake.parse(".") == m.Flake(Path("."), "nixosConfigurations.default")
-    assert m.Flake.parse("", "") == m.Flake(Path("."), "nixosConfigurations.default")
 
 
 @patch(get_qualified_name(platform.node), autospec=True)
@@ -29,15 +29,17 @@ def test_flake_from_arg(mock_node: Any) -> None:
     mock_node.return_value = "hostname"
 
     # Flake string
-    assert m.Flake.from_arg("/path/to/flake#attr") == m.Flake(
+    assert m.Flake.from_arg("/path/to/flake#attr", None) == m.Flake(
         Path("/path/to/flake"), "nixosConfigurations.attr"
     )
 
     # False
-    assert m.Flake.from_arg(False) is None
+    assert m.Flake.from_arg(False, None) is None
 
     # True
-    assert m.Flake.from_arg(True) == m.Flake(Path("."), "nixosConfigurations.hostname")
+    assert m.Flake.from_arg(True, None) == m.Flake(
+        Path("."), "nixosConfigurations.hostname"
+    )
 
     # None when we do not have /etc/nixos/flake.nix
     with patch(
@@ -45,7 +47,7 @@ def test_flake_from_arg(mock_node: Any) -> None:
         autospec=True,
         return_value=False,
     ):
-        assert m.Flake.from_arg(None) is None
+        assert m.Flake.from_arg(None, None) is None
 
     # None when we have a file in /etc/nixos/flake.nix
     with (
@@ -60,7 +62,7 @@ def test_flake_from_arg(mock_node: Any) -> None:
             return_value=False,
         ),
     ):
-        assert m.Flake.from_arg(None) == m.Flake(
+        assert m.Flake.from_arg(None, None) == m.Flake(
             Path("/etc/nixos"), "nixosConfigurations.hostname"
         )
 
@@ -81,10 +83,21 @@ def test_flake_from_arg(mock_node: Any) -> None:
             return_value=Path("/path/to/flake.nix"),
         ),
     ):
-        assert m.Flake.from_arg(None) == m.Flake(
+        assert m.Flake.from_arg(None, None) == m.Flake(
             Path("/path/to"), "nixosConfigurations.hostname"
         )
 
+    with (
+        patch(
+            get_qualified_name(m.subprocess.run),
+            autospec=True,
+            return_value=subprocess.CompletedProcess([], 0, "remote-hostname\n"),
+        ),
+    ):
+        assert m.Flake.from_arg("/path/to", m.Remote("user@host", [], None)) == m.Flake(
+            Path("/path/to"), "nixosConfigurations.remote-hostname"
+        )
+
 
 @patch(get_qualified_name(m.Path.mkdir, m), autospec=True)
 def test_profile_from_name(mock_mkdir: Any) -> None:

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py

@@ -6,17 +6,31 @@ from unittest.mock import ANY, call, patch
 
 import pytest
 
+import nixos_rebuild.models as m
 import nixos_rebuild.nix as n
-from nixos_rebuild import models as m
 
 from .helpers import get_qualified_name
 
 
-@patch(get_qualified_name(n.run, n), autospec=True)
+@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
+def test_copy_closure(mock_run: Any) -> None:
+    closure = Path("/path/to/closure")
+    n.copy_closure(closure, None)
+    mock_run.assert_not_called()
+
+    target_host = m.Remote("user@host", ["--ssh", "opt"], None)
+    n.copy_closure(closure, target_host)
+    mock_run.assert_called_with(
+        ["nix-copy-closure", "--to", "user@host", closure],
+        extra_env={"NIX_SSHOPTS": "--ssh opt"},
+    )
+
+
+@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
 def test_edit(mock_run: Any, monkeypatch: Any, tmpdir: Any) -> None:
     # Flake
     flake = m.Flake.parse(".#attr")
-    n.edit(flake, ["--commit-lock-file"])
+    n.edit(flake, commit_lock_file=True)
     mock_run.assert_called_with(
         [
             "nix",
@@ -42,14 +56,13 @@ def test_edit(mock_run: Any, monkeypatch: Any, tmpdir: Any) -> None:
         mock_run.assert_called_with(["editor", default_nix], check=False)
 
 
-@patch(get_qualified_name(n.shutil.which), autospec=True, return_value="/bin/git")
+def test_get_nixpkgs_rev() -> None:
-def test_get_nixpkgs_rev(mock_which: Any) -> None:
     assert n.get_nixpkgs_rev(None) is None
 
     path = Path("/path/to/nix")
 
     with patch(
-        get_qualified_name(n.run, n),
+        get_qualified_name(n.run_wrapper, n),
         autospec=True,
         side_effect=[CompletedProcess([], 0, "")],
     ) as mock_run:
@@ -58,7 +71,6 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
             ["git", "-C", path, "rev-parse", "--short", "HEAD"],
             check=False,
             stdout=PIPE,
-            text=True,
         )
 
     expected_calls = [
@@ -66,7 +78,6 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
             ["git", "-C", path, "rev-parse", "--short", "HEAD"],
             check=False,
             stdout=PIPE,
-            text=True,
         ),
         call(
             ["git", "-C", path, "diff", "--quiet"],
@@ -75,7 +86,7 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
     ]
 
     with patch(
-        get_qualified_name(n.run, n),
+        get_qualified_name(n.run_wrapper, n),
         autospec=True,
         side_effect=[
             CompletedProcess([], 0, "0f7c82403fd6"),
@@ -86,7 +97,7 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
         mock_run.assert_has_calls(expected_calls)
 
     with patch(
-        get_qualified_name(n.run, n),
+        get_qualified_name(n.run_wrapper, n),
         autospec=True,
         side_effect=[
             CompletedProcess([], 0, "0f7c82403fd6"),
@@ -109,7 +120,7 @@ def test_get_generations_from_nix_store(tmp_path: Path) -> None:
 
     assert n.get_generations(
         m.Profile("system", tmp_path / "system"),
-        lock_profile=False,
+        using_nix_env=False,
     ) == [
         m.Generation(id=1, current=False, timestamp=ANY),
         m.Generation(id=2, current=True, timestamp=ANY),
@@ -118,7 +129,7 @@ def test_get_generations_from_nix_store(tmp_path: Path) -> None:
 
 
 @patch(
-    get_qualified_name(n.run, n),
+    get_qualified_name(n.run_wrapper, n),
     autospec=True,
     return_value=CompletedProcess(
         [],
@@ -134,7 +145,7 @@ def test_get_generations_from_nix_env(mock_run: Any, tmp_path: Path) -> None:
     path = tmp_path / "test"
     path.touch()
 
-    assert n.get_generations(m.Profile("system", path), lock_profile=True) == [
+    assert n.get_generations(m.Profile("system", path), using_nix_env=True) == [
         m.Generation(id=2082, current=False, timestamp="2024-11-07 22:58:56"),
         m.Generation(id=2083, current=False, timestamp="2024-11-07 22:59:41"),
         m.Generation(id=2084, current=True, timestamp="2024-11-07 23:54:17"),
@@ -183,7 +194,7 @@ def test_list_generations(mock_get_generations: Any, tmp_path: Path) -> None:
 
 
 @patch(
-    get_qualified_name(n.run, n),
+    get_qualified_name(n.run_wrapper, n),
     autospec=True,
     return_value=CompletedProcess([], 0, stdout=" \n/path/to/file\n "),
 )
@@ -193,8 +204,8 @@ def test_nixos_build_flake(mock_run: Any) -> None:
     assert n.nixos_build_flake(
         "toplevel",
         flake,
-        ["--nix-flag", "foo"],
         no_link=True,
+        nix_flag="foo",
     ) == Path("/path/to/file")
     mock_run.assert_called_with(
         [
@@ -208,62 +219,62 @@ def test_nixos_build_flake(mock_run: Any) -> None:
             "--nix-flag",
             "foo",
         ],
-        check=True,
-        text=True,
         stdout=PIPE,
     )
 
 
 @patch(
-    get_qualified_name(n.run, n),
+    get_qualified_name(n.run_wrapper, n),
     autospec=True,
     return_value=CompletedProcess([], 0, stdout=" \n/path/to/file\n "),
 )
 def test_nixos_build(mock_run: Any, monkeypatch: Any) -> None:
-    assert n.nixos_build("attr", None, None, ["--nix-flag", "foo"]) == Path(
+    assert n.nixos_build("attr", None, None, nix_flag="foo") == Path("/path/to/file")
-        "/path/to/file"
-    )
     mock_run.assert_called_with(
         ["nix-build", "<nixpkgs/nixos>", "--attr", "attr", "--nix-flag", "foo"],
-        check=True,
-        text=True,
         stdout=PIPE,
     )
 
     n.nixos_build("attr", "preAttr", "file")
     mock_run.assert_called_with(
         ["nix-build", "file", "--attr", "preAttr.attr"],
-        check=True,
-        text=True,
         stdout=PIPE,
     )
 
     n.nixos_build("attr", None, "file", no_out_link=True)
     mock_run.assert_called_with(
         ["nix-build", "file", "--attr", "attr", "--no-out-link"],
-        check=True,
-        text=True,
         stdout=PIPE,
     )
 
     n.nixos_build("attr", "preAttr", None, no_out_link=False, keep_going=True)
     mock_run.assert_called_with(
         ["nix-build", "default.nix", "--attr", "preAttr.attr", "--keep-going"],
-        check=True,
-        text=True,
         stdout=PIPE,
     )
 
 
-@patch(get_qualified_name(n.run, n), autospec=True)
+@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
 def test_rollback(mock_run: Any, tmp_path: Path) -> None:
     path = tmp_path / "test"
     path.touch()
 
     profile = m.Profile("system", path)
 
-    assert n.rollback(profile) == profile.path
+    assert n.rollback(profile, None, False) == profile.path
-    mock_run.assert_called_with(["nix-env", "--rollback", "-p", path], check=True)
+    mock_run.assert_called_with(
+        ["nix-env", "--rollback", "-p", path],
+        remote=None,
+        sudo=False,
+    )
+
+    target_host = m.Remote("user@localhost", [], None)
+    assert n.rollback(profile, target_host, True) == profile.path
+    mock_run.assert_called_with(
+        ["nix-env", "--rollback", "-p", path],
+        remote=target_host,
+        sudo=True,
+    )
 
 
 def test_rollback_temporary_profile(tmp_path: Path) -> None:
@@ -271,7 +282,7 @@ def test_rollback_temporary_profile(tmp_path: Path) -> None:
     path.touch()
     profile = m.Profile("system", path)
 
-    with patch(get_qualified_name(n.run, n), autospec=True) as mock_run:
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
         mock_run.return_value = CompletedProcess(
             [],
             0,
@@ -282,31 +293,62 @@ def test_rollback_temporary_profile(tmp_path: Path) -> None:
                 """),
         )
         assert (
-            n.rollback_temporary_profile(m.Profile("system", path))
+            n.rollback_temporary_profile(m.Profile("system", path), None, False)
             == path.parent / "system-2083-link"
         )
-        assert (
+        mock_run.assert_called_with(
-            n.rollback_temporary_profile(m.Profile("foo", path))
+            [
-            == path.parent / "foo-2083-link"
+                "nix-env",
+                "-p",
+                path,
+                "--list-generations",
+            ],
+            stdout=PIPE,
+            remote=None,
+            sudo=False,
         )
 
-    with patch(get_qualified_name(n.run, n), autospec=True) as mock_run:
+        target_host = m.Remote("user@localhost", [], None)
+        assert (
+            n.rollback_temporary_profile(m.Profile("foo", path), target_host, True)
+            == path.parent / "foo-2083-link"
+        )
+        mock_run.assert_called_with(
+            [
+                "nix-env",
+                "-p",
+                path,
+                "--list-generations",
+            ],
+            stdout=PIPE,
+            remote=target_host,
+            sudo=True,
+        )
+
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
         mock_run.return_value = CompletedProcess([], 0, stdout="")
-        assert n.rollback_temporary_profile(profile) is None
+        assert n.rollback_temporary_profile(profile, None, False) is None
 
 
-@patch(get_qualified_name(n.run, n), autospec=True)
+@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
 def test_set_profile(mock_run: Any) -> None:
     profile_path = Path("/path/to/profile")
     config_path = Path("/path/to/config")
-    n.set_profile(m.Profile("system", profile_path), config_path)
+    n.set_profile(
+        m.Profile("system", profile_path),
+        config_path,
+        target_host=None,
+        sudo=False,
+    )
 
     mock_run.assert_called_with(
-        ["nix-env", "-p", profile_path, "--set", config_path], check=True
+        ["nix-env", "-p", profile_path, "--set", config_path],
+        remote=None,
+        sudo=False,
     )
 
 
-@patch(get_qualified_name(n.run, n), autospec=True)
+@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
 def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
     profile_path = Path("/path/to/profile")
     config_path = Path("/path/to/config")
@@ -317,19 +359,24 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
         n.switch_to_configuration(
             profile_path,
             m.Action.SWITCH,
+            sudo=False,
+            target_host=None,
             specialisation=None,
             install_bootloader=False,
         )
     mock_run.assert_called_with(
         [profile_path / "bin/switch-to-configuration", "switch"],
-        env={"NIXOS_INSTALL_BOOTLOADER": "0", "LOCALE_ARCHIVE": ""},
+        extra_env={"NIXOS_INSTALL_BOOTLOADER": "0"},
-        check=True,
+        sudo=False,
+        remote=None,
     )
 
     with pytest.raises(m.NRError) as e:
         n.switch_to_configuration(
             config_path,
             m.Action.BOOT,
+            sudo=False,
+            target_host=None,
             specialisation="special",
         )
     assert (
@@ -337,13 +384,17 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
         == "error: '--specialisation' can only be used with 'switch' and 'test'"
     )
 
+    target_host = m.Remote("user@localhost", [], None)
     with monkeypatch.context() as mp:
         mp.setenv("LOCALE_ARCHIVE", "/path/to/locale")
+        mp.setenv("PATH", "/path/to/bin")
         mp.setattr(Path, Path.exists.__name__, lambda self: True)
 
         n.switch_to_configuration(
             Path("/path/to/config"),
             m.Action.TEST,
+            sudo=True,
+            target_host=target_host,
             install_bootloader=True,
             specialisation="special",
         )
@@ -352,8 +403,9 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
             config_path / "specialisation/special/bin/switch-to-configuration",
             "test",
         ],
-        env={"NIXOS_INSTALL_BOOTLOADER": "1", "LOCALE_ARCHIVE": "/path/to/locale"},
+        extra_env={"NIXOS_INSTALL_BOOTLOADER": "1"},
-        check=True,
+        sudo=True,
+        remote=target_host,
     )
 
 
@@ -367,11 +419,11 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
     ],
 )
 def test_upgrade_channels(mock_glob: Any) -> None:
-    with patch(get_qualified_name(n.run, n), autospec=True) as mock_run:
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
         n.upgrade_channels(False)
     mock_run.assert_called_with(["nix-channel", "--update", "nixos"], check=False)
 
-    with patch(get_qualified_name(n.run, n), autospec=True) as mock_run:
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
         n.upgrade_channels(True)
     mock_run.assert_has_calls(
         [

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_process.py

@@ -0,0 +1,121 @@
+from pathlib import Path
+from typing import Any
+from unittest.mock import patch
+
+import nixos_rebuild.models as m
+import nixos_rebuild.process as p
+
+from .helpers import get_qualified_name
+
+
+@patch(get_qualified_name(p.subprocess.run))
+def test_run(mock_run: Any) -> None:
+    p.run_wrapper(["test", "--with", "flags"], check=True)
+    mock_run.assert_called_with(
+        ["test", "--with", "flags"],
+        check=True,
+        text=True,
+        errors="surrogateescape",
+        env=None,
+        input=None,
+    )
+
+    with patch.dict(p.os.environ, {"PATH": "/path/to/bin"}, clear=True):
+        p.run_wrapper(
+            ["test", "--with", "flags"],
+            check=False,
+            sudo=True,
+            extra_env={"FOO": "bar"},
+        )
+    mock_run.assert_called_with(
+        ["sudo", "test", "--with", "flags"],
+        check=False,
+        text=True,
+        errors="surrogateescape",
+        env={
+            "PATH": "/path/to/bin",
+            "FOO": "bar",
+        },
+        input=None,
+    )
+
+    p.run_wrapper(
+        ["test", "--with", "flags"],
+        check=True,
+        remote=m.Remote("user@localhost", ["--ssh", "opt"], "password"),
+    )
+    mock_run.assert_called_with(
+        ["ssh", "--ssh", "opt", "user@localhost", "--", "test", "--with", "flags"],
+        check=True,
+        text=True,
+        errors="surrogateescape",
+        env=None,
+        input=None,
+    )
+
+    p.run_wrapper(
+        ["test", "--with", "flags"],
+        check=True,
+        sudo=True,
+        extra_env={"FOO": "bar"},
+        remote=m.Remote("user@localhost", ["--ssh", "opt"], "password"),
+    )
+    mock_run.assert_called_with(
+        [
+            "ssh",
+            "--ssh",
+            "opt",
+            "user@localhost",
+            "--",
+            "sudo",
+            "--prompt=",
+            "--stdin",
+            "env",
+            "FOO=bar",
+            "test",
+            "--with",
+            "flags",
+        ],
+        check=True,
+        env=None,
+        text=True,
+        errors="surrogateescape",
+        input="password\n",
+    )
+
+
+def test_remote_from_name(monkeypatch: Any, tmpdir: Path) -> None:
+    monkeypatch.setenv("NIX_SSHOPTS", "")
+    assert m.Remote.from_arg("user@localhost", None, tmpdir) == m.Remote(
+        "user@localhost",
+        opts=[
+            "-o",
+            "ControlMaster=auto",
+            "-o",
+            f"ControlPath={tmpdir / "ssh-%n"}",
+            "-o",
+            "ControlPersist=60",
+        ],
+        sudo_password=None,
+    )
+
+    # get_qualified_name doesn't work because getpass is aliased to another
+    # function
+    with patch(f"{p.__name__}.getpass", return_value="password"):
+        monkeypatch.setenv("NIX_SSHOPTS", "-f foo -b bar")
+        assert m.Remote.from_arg("user@localhost", True, tmpdir) == m.Remote(
+            "user@localhost",
+            opts=[
+                "-f",
+                "foo",
+                "-b",
+                "bar",
+                "-o",
+                "ControlMaster=auto",
+                "-o",
+                f"ControlPath={tmpdir / "ssh-%n"}",
+                "-o",
+                "ControlPersist=60",
+            ],
+            sudo_password="password",
+        )

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py

@@ -1,8 +1,8 @@
-from nixos_rebuild import utils as u
+import nixos_rebuild.utils as u
 
 
 def test_dict_to_flags() -> None:
-    r = u.dict_to_flags(
+    r1 = u.dict_to_flags(
         {
             "test_flag_1": True,
             "test_flag_2": False,
@@ -12,7 +12,7 @@ def test_dict_to_flags() -> None:
             "verbose": 5,
         }
     )
-    assert r == [
+    assert r1 == [
         "--test-flag-1",
         "--test-flag-3",
         "value",
@@ -21,3 +21,5 @@ def test_dict_to_flags() -> None:
         "v2",
         "-vvvvv",
     ]
+    r2 = u.dict_to_flags({"verbose": 0, "empty_list": []})
+    assert r2 == []