nixos-rebuild-ng: add --sudo/--use-remote-sudo flags

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

@@ -47,6 +47,9 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
     parser.add_argument("--upgrade", action="store_true")
     parser.add_argument("--upgrade-all", action="store_true")
     parser.add_argument("--json", action="store_true")
+    parser.add_argument("--sudo", action="store_true")
+    # TODO: add deprecated=True in Python >=3.13
+    parser.add_argument("--use-remote-sudo", dest="sudo", action="store_true")
     parser.add_argument("action", choices=Action.values(), nargs="?")
     parser.add_argument("--verbose", "-v", action="count", default=0)
 
@@ -181,7 +184,7 @@ def execute(argv: list[str]) -> None:
                     no_link=True,
                     **flake_flags,
                 )
-                set_profile(profile, path_to_config)
+                set_profile(profile, path_to_config, sudo=args.sudo)
             else:
                 path_to_config = nixos_build(
                     "system",
@@ -190,10 +193,11 @@ def execute(argv: list[str]) -> None:
                     no_out_link=True,
                     **nix_flags,
                 )
-                set_profile(profile, path_to_config)
+                set_profile(profile, path_to_config, sudo=args.sudo)
             switch_to_configuration(
                 path_to_config,
                 action,
+                sudo=args.sudo,
                 specialisation=args.specialisation,
                 install_bootloader=args.install_bootloader,
             )
@@ -227,6 +231,7 @@ def execute(argv: list[str]) -> None:
                 switch_to_configuration(
                     path_to_config,
                     action,
+                    sudo=args.sudo,
                     specialisation=args.specialisation,
                 )
         case Action.BUILD_VM | Action.BUILD_VM_WITH_BOOTLOADER:

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

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import platform
 import re
 from dataclasses import dataclass
@@ -114,3 +115,17 @@ class Profile:
                 path = Path("/nix/var/nix/profiles/system-profiles") / name
                 path.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
                 return Profile(name, path)
+
+
+@dataclass(frozen=True)
+class SSH:
+    host: str
+    opts: list[str]
+
+    @staticmethod
+    def from_arg(host: str | None) -> SSH | None:
+        if host:
+            opts = os.getenv("SSH_OPTS", "").split()
+            return SSH(host, opts)
+        else:
+            return None

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

@@ -4,7 +4,7 @@ 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 (
@@ -15,6 +15,7 @@ from .models import (
     NRError,
     Profile,
 )
+from .process import run
 from .utils import Args, dict_to_flags, info
 
 FLAKE_FLAGS: Final = ["--extra-experimental-features", "nix-command flakes"]
@@ -280,14 +281,15 @@ 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, 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(["nix-env", "-p", profile.path, "--set", path_to_config], check=True, sudo=sudo)
 
 
 def switch_to_configuration(
     path_to_config: Path,
     action: Action,
+    sudo: bool,
     install_bootloader: bool = False,
     specialisation: str | None = None,
 ) -> None:
@@ -313,6 +315,7 @@ def switch_to_configuration(
             "LOCALE_ARCHIVE": os.getenv("LOCALE_ARCHIVE", ""),
         },
         check=True,
+        sudo=sudo,
     )
 
 

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

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import Any, Sequence
+
+from .models import SSH
+
+
+def run(
+    args: Sequence[str | bytes | os.PathLike[Any]],
+    # make `check` explicit so we always know if the code is aborting on errors
+    check: bool,
+    remote: SSH | None = None,
+    sudo: bool = False,
+    **kwargs: Any,
+) -> subprocess.CompletedProcess[Any]:
+    if sudo:
+        args = ["sudo"] + list(args)
+    if remote:
+        args = ["ssh", *remote.opts, remote.host, "--"] + list(args)
+    return subprocess.run(args, check=check, **kwargs)

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

@@ -70,7 +70,7 @@ def test_parse_args() -> None:
     assert r2.attr == "bar"
 
 
-@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.shutil.which), autospec=True, return_value="/bin/git")
 def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> None:
     nixpkgs_path = tmp_path / "nixpkgs"
@@ -143,7 +143,7 @@ def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> Non
     )
 
 
-@patch(get_qualified_name(nr.nix.run, nr.nix), autospec=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()
@@ -163,6 +163,7 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
             "--flake",
             "/path/to/config#hostname",
             "--install-bootloader",
+            "--sudo",
             "--verbose",
         ]
     )
@@ -187,6 +188,7 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
             ),
             call(
                 [
+                    "sudo",
                     "nix-env",
                     "-p",
                     Path("/nix/var/nix/profiles/system"),
@@ -196,7 +198,7 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
                 check=True,
             ),
             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,
             ),
@@ -204,7 +206,7 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
     )
 
 
-@patch(get_qualified_name(nr.nix.run, nr.nix), autospec=True)
+@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()
@@ -236,7 +238,7 @@ def test_execute_switch_rollback(mock_run: Any, tmp_path: Path) -> None:
     )
 
 
-@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(

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

@@ -3,7 +3,7 @@ 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
 
@@ -98,3 +98,12 @@ def test_profile_from_name(mock_mkdir: Any) -> None:
         Path("/nix/var/nix/profiles/system-profiles/something"),
     )
     mock_mkdir.assert_called_once()
+
+
+def test_ssh_from_name(monkeypatch: Any) -> None:
+    assert m.SSH.from_arg("user@localhost") == m.SSH("user@localhost", [])
+
+    monkeypatch.setenv("SSH_OPTS", "-f foo -b bar")
+    assert m.SSH.from_arg("user@localhost") == m.SSH(
+        "user@localhost", ["-f", "foo", "-b", "bar"]
+    )

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

@@ -6,8 +6,8 @@ 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
 
@@ -297,10 +297,12 @@ def test_rollback_temporary_profile(tmp_path: Path) -> None:
 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, 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],
+        check=True,
+        sudo=False,
     )
 
 
@@ -315,6 +317,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
         n.switch_to_configuration(
             profile_path,
             m.Action.SWITCH,
+            sudo=False,
             specialisation=None,
             install_bootloader=False,
         )
@@ -322,12 +325,14 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
         [profile_path / "bin/switch-to-configuration", "switch"],
         env={"NIXOS_INSTALL_BOOTLOADER": "0", "LOCALE_ARCHIVE": ""},
         check=True,
+        sudo=False,
     )
 
     with pytest.raises(m.NRError) as e:
         n.switch_to_configuration(
             config_path,
             m.Action.BOOT,
+            sudo=False,
             specialisation="special",
         )
     assert (
@@ -342,6 +347,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
         n.switch_to_configuration(
             Path("/path/to/config"),
             m.Action.TEST,
+            sudo=True,
             install_bootloader=True,
             specialisation="special",
         )
@@ -352,6 +358,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
         ],
         env={"NIXOS_INSTALL_BOOTLOADER": "1", "LOCALE_ARCHIVE": "/path/to/locale"},
         check=True,
+        sudo=True,
     )
 
 

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

@@ -0,0 +1,47 @@
+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(["test", "--with", "flags"], check=True)
+    mock_run.assert_called_with(["test", "--with", "flags"], check=True)
+
+    p.run(["test", "--with", "flags"], check=False, sudo=True)
+    mock_run.assert_called_with(["sudo", "test", "--with", "flags"], check=False)
+
+    p.run(
+        ["test", "--with", "flags"],
+        check=True,
+        remote=m.SSH("user@localhost", ["--ssh", "opt"]),
+    )
+    mock_run.assert_called_with(
+        ["ssh", "--ssh", "opt", "user@localhost", "--", "test", "--with", "flags"],
+        check=True,
+    )
+
+    p.run(
+        ["test", "--with", "flags"],
+        check=True,
+        sudo=True,
+        remote=m.SSH("user@localhost", ["--ssh", "opt"]),
+    )
+    mock_run.assert_called_with(
+        [
+            "ssh",
+            "--ssh",
+            "opt",
+            "user@localhost",
+            "--",
+            "sudo",
+            "test",
+            "--with",
+            "flags",
+        ],
+        check=True,
+    )

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

@@ -1,6 +1,6 @@
 import argparse
 
-from nixos_rebuild import utils as u
+import nixos_rebuild.utils as u
 
 
 def test_dict_to_flags() -> None: