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

This commit is contained in:
Thiago Kenji Okada 2024-11-17 18:26:45 +00:00
parent 3b41ec0691
commit 6c6d08dc4f
9 changed files with 125 additions and 15 deletions

View File

@ -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:

View File

@ -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

View File

@ -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,
)

View File

@ -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)

View File

@ -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(

View File

@ -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"]
)

View File

@ -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,
)

View File

@ -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,
)

View File

@ -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: