Loading pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py +15 −17 Original line number Diff line number Diff line Loading @@ -63,20 +63,6 @@ class BuildAttr: return cls(Path(file or "default.nix"), attr) def _get_hostname(target_host: Remote | None) -> 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() @dataclass(frozen=True) class Flake: path: str Loading @@ -95,9 +81,7 @@ class Flake: 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 _get_hostname(target_host) or "default"}"' ) nixos_attr = f'nixosConfigurations."{attr or cls._get_hostname(target_host) or "default"}"' path = m.group("path") return cls(path, nixos_attr) Loading Loading @@ -126,6 +110,20 @@ class Flake: except FileNotFoundError: return self.path @staticmethod def _get_hostname(target_host: Remote | None) -> 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() @dataclass(frozen=True) class Generation: Loading pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py +9 −7 Original line number Diff line number Diff line Loading @@ -25,7 +25,7 @@ from .models import ( Profile, Remote, ) from .process import SSH_DEFAULT_OPTS, run_wrapper from .process import PRESERVE_ENV, SSH_DEFAULT_OPTS, run_wrapper from .utils import Args, dict_to_flags FLAKE_FLAGS: Final = ["--extra-experimental-features", "nix-command flakes"] Loading Loading @@ -192,9 +192,7 @@ def copy_closure( Also supports copying a closure from a remote to another remote.""" sshopts = os.getenv("NIX_SSHOPTS", "") extra_env = { "NIX_SSHOPTS": " ".join(filter(lambda x: x, [sshopts, *SSH_DEFAULT_OPTS])) } env = {"NIX_SSHOPTS": " ".join(filter(lambda x: x, [sshopts, *SSH_DEFAULT_OPTS]))} def nix_copy_closure(host: Remote, to: bool) -> None: run_wrapper( Loading @@ -205,7 +203,7 @@ def copy_closure( host.host, closure, ], extra_env=extra_env, env=env, ) def nix_copy(to_host: Remote, from_host: Remote) -> None: Loading @@ -221,7 +219,7 @@ def copy_closure( f"ssh://{to_host.host}", closure, ], extra_env=extra_env, env=env, ) match (to_host, from_host): Loading Loading @@ -703,7 +701,11 @@ def switch_to_configuration( run_wrapper( [*cmd, path_to_config / "bin/switch-to-configuration", str(action)], extra_env={"NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0"}, env={ "LOCALE_ARCHIVE": PRESERVE_ENV, "NIXOS_NO_CHECK": PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0", }, remote=target_host, sudo=sudo, ) Loading pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py +163 −43 Original line number Diff line number Diff line Loading @@ -5,10 +5,11 @@ import os import re import shlex import subprocess from collections.abc import Sequence from collections.abc import Mapping, Sequence from dataclasses import dataclass from enum import Enum from ipaddress import AddressValueError, IPv6Address from typing import Final, Self, TextIO, TypedDict, Unpack from typing import Final, Literal, Self, TextIO, TypedDict, Unpack, override from . import tmpdir Loading @@ -23,7 +24,30 @@ SSH_DEFAULT_OPTS: Final = [ "ControlPersist=60", ] type Args = Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]] class _Env(Enum): PRESERVE_ENV = "PRESERVE" @override def __repr__(self) -> str: return self.value PRESERVE_ENV: Final = _Env.PRESERVE_ENV type Arg = str | bytes | os.PathLike[str] | os.PathLike[bytes] type Args = Sequence[Arg] type EnvValue = str | Literal[_Env.PRESERVE_ENV] @dataclass(frozen=True) class _RawShellArg: value: str @override def __str__(self) -> str: return self.value @dataclass(frozen=True) Loading Loading @@ -106,70 +130,84 @@ def run_wrapper( args: Args, *, check: bool = True, extra_env: dict[str, str] | None = None, env: Mapping[str, EnvValue] | 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 process_input = None run_args = args run_args: list[Arg] = list(args) final_args: list[Arg] normalized_env = _normalize_env(env) resolved_env = _resolve_env_local(normalized_env) if remote: if extra_env: extra_env_args = [f"{env}={value}" for env, value in extra_env.items()] args = ["env", *extra_env_args, *args] # Apply env for the *remote command* (not for ssh itself) remote_run_args: list[Arg | _RawShellArg] = [] remote_run_args.extend(run_args) if normalized_env: remote_run_args = _prefix_env_cmd_remote(run_args, normalized_env) if sudo: sudo_args = shlex.split(os.getenv("NIX_SUDOOPTS", "")) if remote.sudo_password: args = ["sudo", "--prompt=", "--stdin", *args] remote_run_args = [ "sudo", "--prompt=", "--stdin", *sudo_args, *remote_run_args, ] process_input = remote.sudo_password + "\n" else: args = ["sudo", *args] run_args = [ remote_run_args = ["sudo", *sudo_args, *remote_run_args] ssh_args: list[Arg] = [ "ssh", *remote.opts, *SSH_DEFAULT_OPTS, remote.ssh_host(), "--", # SSH will join the parameters here and pass it to the shell, so we # need to quote it to avoid issues. # We can't use `shlex.join`, otherwise we will hit MAX_ARG_STRLEN # limits when the command becomes too big. *[shlex.quote(str(a)) for a in args], *[_quote_remote_arg(a) for a in remote_run_args], ] final_args = ssh_args popen_env = None # keep ssh's environment normal else: if extra_env: env = os.environ | extra_env if sudo: # subprocess.run(env=...) would affect sudo, but sudo may drop env # for the target command. # So we inject env via `sudo env ... cmd`. if env is not None and resolved_env: run_args = _prefix_env_cmd(run_args, resolved_env) sudo_args = shlex.split(os.getenv("NIX_SUDOOPTS", "")) # Using --preserve-env is less than ideal since it will cause # the following warn during usage: # > warning: $HOME ('/home/<user>') is not owned by you, # > falling back to the one defined in the 'passwd' file ('/root') # However, right now it is the only way to guarantee the semantics # expected for the commands, e.g. activation with systemd-run # expects access to environment variables like LOCALE_ARCHIVE, # NIXOS_NO_CHECK. # For now, for anyone that the above warn bothers you, please # use `sudo nixos-rebuild` instead of `--sudo` flag. run_args = ["sudo", "--preserve-env", *sudo_args, *run_args] final_args = ["sudo", *sudo_args, *run_args] # No need to pass env to subprocess.run; keep sudo's own env # default. popen_env = None else: # Non-sudo local: we can fully control the environment with # subprocess.run(env=...) final_args = run_args popen_env = None if env is None else resolved_env logger.debug( "calling run with args=%r, kwargs=%r, extra_env=%r", run_args, "calling run with args=%r, kwargs=%r, env=%r", _sanitize_env_run_args(remote_run_args if remote else run_args), kwargs, extra_env, env, ) try: r = subprocess.run( run_args, final_args, check=check, env=env, env=popen_env, input=process_input, # Hope nobody is using NixOS with non-UTF8 encodings, but # "surrogateescape" should still work in those systems. text=True, errors="surrogateescape", **kwargs, Loading @@ -195,13 +233,95 @@ def run_wrapper( raise # SSH does not send the signals to the process when running without usage of # pseudo-TTY (that causes a whole other can of worms), so if the process is # long running (e.g.: a build) this will result in the underlying process # staying alive. # See: https://stackoverflow.com/a/44354466 # Issue: https://github.com/NixOS/nixpkgs/issues/403269 def _resolve_env(env: Mapping[str, EnvValue] | None) -> dict[str, str]: normalized = _normalize_env(env) return _resolve_env_local(normalized) def _normalize_env(env: Mapping[str, EnvValue] | None) -> dict[str, EnvValue]: """ Normalize env mapping, but preserve some environment variables by default. """ return {"PATH": PRESERVE_ENV, **(env or {})} def _resolve_env_local(env: dict[str, EnvValue]) -> dict[str, str]: """ Resolve env mapping where values can be: - PRESERVE_ENV: copy from current os.environ (if present) - str: explicit value """ result: dict[str, str] = {} for k, v in env.items(): if v == PRESERVE_ENV: cur = os.environ.get(k) if cur is not None: result[k] = cur else: result[k] = v return result def _prefix_env_cmd(cmd: Sequence[Arg], resolved_env: dict[str, str]) -> list[Arg]: """ Prefix a command with `env -i K=V ... -- <cmd...>` to set vars for the command. """ if not resolved_env: return list(cmd) assigns = [f"{k}={v}" for k, v in resolved_env.items()] return ["env", "-i", *assigns, *cmd] def _prefix_env_cmd_remote( cmd: Args, env: dict[str, EnvValue], ) -> list[Arg | _RawShellArg]: """ Prefix remote commands with env assignments. Preserve markers are expanded by the remote shell at execution time. """ assigns: list[str | _RawShellArg] = [] for k, v in env.items(): if v is PRESERVE_ENV: assigns.append(_RawShellArg(f"{k}=${{{k}-}}")) else: assigns.append(f"{k}={v}") return ["env", "-i", *assigns, *cmd] def _quote_remote_arg(arg: Arg | _RawShellArg) -> str: if isinstance(arg, _RawShellArg): return str(arg) return shlex.quote(str(arg)) def _sanitize_env_run_args(run_args: list[Arg] | list[Arg | _RawShellArg]) -> list[Arg]: """ Sanitize long or sensitive environment variables from logs. """ sanitized: list[Arg] = [] for value in run_args: if isinstance(value, str) and value.startswith("PATH="): sanitized.append("PATH=<PATH>") elif isinstance(value, str | bytes | os.PathLike): sanitized.append(value) else: sanitized.append(str(value)) return sanitized def _kill_long_running_ssh_process(args: Args, remote: Remote) -> None: """ SSH does not send the signals to the process when running without usage of pseudo-TTY (that causes a whole other can of worms), so if the process is long running (e.g.: a build) this will result in the underlying process staying alive. See: https://stackoverflow.com/a/44354466 Issue: https://github.com/NixOS/nixpkgs/issues/403269 """ logger.info("cleaning-up remote process, please wait...") # We need to escape both the shell and regex here (since pkill interprets Loading pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py +55 −13 Original line number Diff line number Diff line Loading @@ -248,7 +248,6 @@ def test_execute_nix_boot(mock_run: Mock, tmp_path: Path) -> None: | { "env": { "NIXOS_INSTALL_BOOTLOADER": "0", "NIXOS_REBUILD_I_UNDERSTAND_THE_CONSEQUENCES_PLEASE_BREAK_MY_SYSTEM": "1", } } ), Loading Loading @@ -487,7 +486,6 @@ def test_execute_nix_switch_flake(mock_run: Mock, tmp_path: Path) -> None: call( [ "sudo", "--preserve-env", "nix-env", "-p", Path("/nix/var/nix/profiles/system"), Loading @@ -505,21 +503,15 @@ def test_execute_nix_switch_flake(mock_run: Mock, tmp_path: Path) -> None: call( [ "sudo", "--preserve-env", "env", "-i", "NIXOS_INSTALL_BOOTLOADER=1", *nr.nix.SWITCH_TO_CONFIGURATION_CMD_PREFIX, config_path / "bin/switch-to-configuration", "switch", ], check=True, **( DEFAULT_RUN_KWARGS | { "env": { "NIXOS_INSTALL_BOOTLOADER": "1", "NIXOS_REBUILD_I_UNDERSTAND_THE_CONSEQUENCES_PLEASE_BREAK_MY_SYSTEM": "1", } } ), **DEFAULT_RUN_KWARGS, ), ] ) Loading Loading @@ -625,6 +617,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "mktemp", "-d", "-t", Loading @@ -640,6 +635,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "nix-store", "--realise", str(config_path), Loading @@ -656,6 +654,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "readlink", "-f", "/tmp/tmpdir/config", Loading @@ -670,6 +671,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "rm", "-rf", "/tmp/tmpdir", Loading Loading @@ -699,6 +703,9 @@ def test_execute_nix_switch_build_target_host( "user@target-host", "--", "sudo", "env", "-i", "PATH=${PATH-}", "nix-env", "-p", "/nix/var/nix/profiles/system", Loading @@ -714,6 +721,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@target-host", "--", "env", "-i", "PATH=${PATH-}", "test", "-d", "/run/systemd/system", Loading @@ -729,6 +739,10 @@ def test_execute_nix_switch_build_target_host( "--", "sudo", "env", "-i", "PATH=${PATH-}", "LOCALE_ARCHIVE=${LOCALE_ARCHIVE-}", "NIXOS_NO_CHECK=${NIXOS_NO_CHECK-}", "NIXOS_INSTALL_BOOTLOADER=0", *nr.nix.SWITCH_TO_CONFIGURATION_CMD_PREFIX, str(config_path / "bin/switch-to-configuration"), Loading Loading @@ -806,6 +820,9 @@ def test_execute_nix_switch_flake_target_host( "user@localhost", "--", "sudo", "env", "-i", "PATH=${PATH-}", "nix-env", "-p", "/nix/var/nix/profiles/system", Loading @@ -821,6 +838,9 @@ def test_execute_nix_switch_flake_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@localhost", "--", "env", "-i", "PATH=${PATH-}", "test", "-d", "/run/systemd/system", Loading @@ -836,6 +856,10 @@ def test_execute_nix_switch_flake_target_host( "--", "sudo", "env", "-i", "PATH=${PATH-}", "LOCALE_ARCHIVE=${LOCALE_ARCHIVE-}", "NIXOS_NO_CHECK=${NIXOS_NO_CHECK-}", "NIXOS_INSTALL_BOOTLOADER=0", *nr.nix.SWITCH_TO_CONFIGURATION_CMD_PREFIX, str(config_path / "bin/switch-to-configuration"), Loading Loading @@ -912,6 +936,9 @@ def test_execute_nix_switch_flake_build_host( *nr.process.SSH_DEFAULT_OPTS, "user@localhost", "--", "env", "-i", "PATH=${PATH-}", "nix", "--extra-experimental-features", "'nix-command flakes'", Loading Loading @@ -1121,6 +1148,9 @@ def test_execute_build_dry_run_build_and_target_remote( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "nix", "--extra-experimental-features", "'nix-command flakes'", Loading Loading @@ -1304,7 +1334,6 @@ def test_execute_switch_store_path(mock_run: Mock, tmp_path: Path) -> None: | { "env": { "NIXOS_INSTALL_BOOTLOADER": "0", "NIXOS_REBUILD_I_UNDERSTAND_THE_CONSEQUENCES_PLEASE_BREAK_MY_SYSTEM": "1", } } ), Loading Loading @@ -1354,6 +1383,9 @@ def test_execute_switch_store_path_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@remote-host", "--", "env", "-i", "PATH=${PATH-}", "test", "-f", str(config_path / "nixos-version"), Loading @@ -1368,6 +1400,9 @@ def test_execute_switch_store_path_target_host( "user@remote-host", "--", "sudo", "env", "-i", "PATH=${PATH-}", "nix-env", "-p", "/nix/var/nix/profiles/system", Loading @@ -1383,6 +1418,9 @@ def test_execute_switch_store_path_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@remote-host", "--", "env", "-i", "PATH=${PATH-}", "test", "-d", "/run/systemd/system", Loading @@ -1398,6 +1436,10 @@ def test_execute_switch_store_path_target_host( "--", "sudo", "env", "-i", "PATH=${PATH-}", "LOCALE_ARCHIVE=${LOCALE_ARCHIVE-}", "NIXOS_NO_CHECK=${NIXOS_NO_CHECK-}", "NIXOS_INSTALL_BOOTLOADER=0", *nr.nix.SWITCH_TO_CONFIGURATION_CMD_PREFIX, str(config_path / "bin/switch-to-configuration"), Loading pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py +26 −18 Original line number Diff line number Diff line Loading @@ -134,9 +134,7 @@ def test_build_remote( "user@host", Path("/path/to/file"), ], extra_env={ "NIX_SSHOPTS": " ".join(["--ssh opts", *p.SSH_DEFAULT_OPTS]) }, env={"NIX_SSHOPTS": " ".join(["--ssh opts", *p.SSH_DEFAULT_OPTS])}, ), call( ["mktemp", "-d", "-t", "nixos-rebuild.XXXXX"], Loading Loading @@ -208,9 +206,7 @@ def test_build_remote_flake( "user@host", Path("/path/to/file"), ], extra_env={ "NIX_SSHOPTS": " ".join(["--ssh opts", *p.SSH_DEFAULT_OPTS]) }, env={"NIX_SSHOPTS": " ".join(["--ssh opts", *p.SSH_DEFAULT_OPTS])}, ), call( [ Loading Loading @@ -241,7 +237,7 @@ def test_copy_closure(monkeypatch: MonkeyPatch) -> None: n.copy_closure(closure, target_host) mock_run.assert_called_with( ["nix-copy-closure", "--to", "user@target.host", closure], extra_env={"NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS)}, env={"NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS)}, ) monkeypatch.setenv("NIX_SSHOPTS", "--ssh build-opt") Loading @@ -249,15 +245,11 @@ def test_copy_closure(monkeypatch: MonkeyPatch) -> None: n.copy_closure(closure, None, build_host, {"copy_flag": True}) mock_run.assert_called_with( ["nix-copy-closure", "--copy-flag", "--from", "user@build.host", closure], extra_env={ "NIX_SSHOPTS": " ".join(["--ssh build-opt", *p.SSH_DEFAULT_OPTS]) }, env={"NIX_SSHOPTS": " ".join(["--ssh build-opt", *p.SSH_DEFAULT_OPTS])}, ) monkeypatch.setenv("NIX_SSHOPTS", "--ssh build-target-opt") extra_env = { "NIX_SSHOPTS": " ".join(["--ssh build-target-opt", *p.SSH_DEFAULT_OPTS]) } env = {"NIX_SSHOPTS": " ".join(["--ssh build-target-opt", *p.SSH_DEFAULT_OPTS])} with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run: n.copy_closure(closure, target_host, build_host, {"copy_flag": True}) mock_run.assert_called_with( Loading @@ -273,7 +265,7 @@ def test_copy_closure(monkeypatch: MonkeyPatch) -> None: "ssh://user@target.host", closure, ], extra_env=extra_env, env=env, ) Loading Loading @@ -723,7 +715,11 @@ def test_switch_to_configuration_without_systemd_run( ) mock_run.assert_called_with( [profile_path / "bin/switch-to-configuration", "switch"], extra_env={"NIXOS_INSTALL_BOOTLOADER": "0"}, env={ "LOCALE_ARCHIVE": p.PRESERVE_ENV, "NIXOS_NO_CHECK": p.PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "0", }, sudo=False, remote=None, ) Loading Loading @@ -760,7 +756,11 @@ def test_switch_to_configuration_without_systemd_run( config_path / "specialisation/special/bin/switch-to-configuration", "test", ], extra_env={"NIXOS_INSTALL_BOOTLOADER": "1"}, env={ "LOCALE_ARCHIVE": p.PRESERVE_ENV, "NIXOS_NO_CHECK": p.PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "1", }, sudo=True, remote=target_host, ) Loading Loading @@ -791,7 +791,11 @@ def test_switch_to_configuration_with_systemd_run( profile_path / "bin/switch-to-configuration", "switch", ], extra_env={"NIXOS_INSTALL_BOOTLOADER": "0"}, env={ "LOCALE_ARCHIVE": p.PRESERVE_ENV, "NIXOS_NO_CHECK": p.PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "0", }, sudo=False, remote=None, ) Loading @@ -816,7 +820,11 @@ def test_switch_to_configuration_with_systemd_run( config_path / "specialisation/special/bin/switch-to-configuration", "test", ], extra_env={"NIXOS_INSTALL_BOOTLOADER": "1"}, env={ "LOCALE_ARCHIVE": p.PRESERVE_ENV, "NIXOS_NO_CHECK": p.PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "1", }, sudo=True, remote=target_host, ) Loading Loading
pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py +15 −17 Original line number Diff line number Diff line Loading @@ -63,20 +63,6 @@ class BuildAttr: return cls(Path(file or "default.nix"), attr) def _get_hostname(target_host: Remote | None) -> 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() @dataclass(frozen=True) class Flake: path: str Loading @@ -95,9 +81,7 @@ class Flake: 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 _get_hostname(target_host) or "default"}"' ) nixos_attr = f'nixosConfigurations."{attr or cls._get_hostname(target_host) or "default"}"' path = m.group("path") return cls(path, nixos_attr) Loading Loading @@ -126,6 +110,20 @@ class Flake: except FileNotFoundError: return self.path @staticmethod def _get_hostname(target_host: Remote | None) -> 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() @dataclass(frozen=True) class Generation: Loading
pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py +9 −7 Original line number Diff line number Diff line Loading @@ -25,7 +25,7 @@ from .models import ( Profile, Remote, ) from .process import SSH_DEFAULT_OPTS, run_wrapper from .process import PRESERVE_ENV, SSH_DEFAULT_OPTS, run_wrapper from .utils import Args, dict_to_flags FLAKE_FLAGS: Final = ["--extra-experimental-features", "nix-command flakes"] Loading Loading @@ -192,9 +192,7 @@ def copy_closure( Also supports copying a closure from a remote to another remote.""" sshopts = os.getenv("NIX_SSHOPTS", "") extra_env = { "NIX_SSHOPTS": " ".join(filter(lambda x: x, [sshopts, *SSH_DEFAULT_OPTS])) } env = {"NIX_SSHOPTS": " ".join(filter(lambda x: x, [sshopts, *SSH_DEFAULT_OPTS]))} def nix_copy_closure(host: Remote, to: bool) -> None: run_wrapper( Loading @@ -205,7 +203,7 @@ def copy_closure( host.host, closure, ], extra_env=extra_env, env=env, ) def nix_copy(to_host: Remote, from_host: Remote) -> None: Loading @@ -221,7 +219,7 @@ def copy_closure( f"ssh://{to_host.host}", closure, ], extra_env=extra_env, env=env, ) match (to_host, from_host): Loading Loading @@ -703,7 +701,11 @@ def switch_to_configuration( run_wrapper( [*cmd, path_to_config / "bin/switch-to-configuration", str(action)], extra_env={"NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0"}, env={ "LOCALE_ARCHIVE": PRESERVE_ENV, "NIXOS_NO_CHECK": PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0", }, remote=target_host, sudo=sudo, ) Loading
pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py +163 −43 Original line number Diff line number Diff line Loading @@ -5,10 +5,11 @@ import os import re import shlex import subprocess from collections.abc import Sequence from collections.abc import Mapping, Sequence from dataclasses import dataclass from enum import Enum from ipaddress import AddressValueError, IPv6Address from typing import Final, Self, TextIO, TypedDict, Unpack from typing import Final, Literal, Self, TextIO, TypedDict, Unpack, override from . import tmpdir Loading @@ -23,7 +24,30 @@ SSH_DEFAULT_OPTS: Final = [ "ControlPersist=60", ] type Args = Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]] class _Env(Enum): PRESERVE_ENV = "PRESERVE" @override def __repr__(self) -> str: return self.value PRESERVE_ENV: Final = _Env.PRESERVE_ENV type Arg = str | bytes | os.PathLike[str] | os.PathLike[bytes] type Args = Sequence[Arg] type EnvValue = str | Literal[_Env.PRESERVE_ENV] @dataclass(frozen=True) class _RawShellArg: value: str @override def __str__(self) -> str: return self.value @dataclass(frozen=True) Loading Loading @@ -106,70 +130,84 @@ def run_wrapper( args: Args, *, check: bool = True, extra_env: dict[str, str] | None = None, env: Mapping[str, EnvValue] | 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 process_input = None run_args = args run_args: list[Arg] = list(args) final_args: list[Arg] normalized_env = _normalize_env(env) resolved_env = _resolve_env_local(normalized_env) if remote: if extra_env: extra_env_args = [f"{env}={value}" for env, value in extra_env.items()] args = ["env", *extra_env_args, *args] # Apply env for the *remote command* (not for ssh itself) remote_run_args: list[Arg | _RawShellArg] = [] remote_run_args.extend(run_args) if normalized_env: remote_run_args = _prefix_env_cmd_remote(run_args, normalized_env) if sudo: sudo_args = shlex.split(os.getenv("NIX_SUDOOPTS", "")) if remote.sudo_password: args = ["sudo", "--prompt=", "--stdin", *args] remote_run_args = [ "sudo", "--prompt=", "--stdin", *sudo_args, *remote_run_args, ] process_input = remote.sudo_password + "\n" else: args = ["sudo", *args] run_args = [ remote_run_args = ["sudo", *sudo_args, *remote_run_args] ssh_args: list[Arg] = [ "ssh", *remote.opts, *SSH_DEFAULT_OPTS, remote.ssh_host(), "--", # SSH will join the parameters here and pass it to the shell, so we # need to quote it to avoid issues. # We can't use `shlex.join`, otherwise we will hit MAX_ARG_STRLEN # limits when the command becomes too big. *[shlex.quote(str(a)) for a in args], *[_quote_remote_arg(a) for a in remote_run_args], ] final_args = ssh_args popen_env = None # keep ssh's environment normal else: if extra_env: env = os.environ | extra_env if sudo: # subprocess.run(env=...) would affect sudo, but sudo may drop env # for the target command. # So we inject env via `sudo env ... cmd`. if env is not None and resolved_env: run_args = _prefix_env_cmd(run_args, resolved_env) sudo_args = shlex.split(os.getenv("NIX_SUDOOPTS", "")) # Using --preserve-env is less than ideal since it will cause # the following warn during usage: # > warning: $HOME ('/home/<user>') is not owned by you, # > falling back to the one defined in the 'passwd' file ('/root') # However, right now it is the only way to guarantee the semantics # expected for the commands, e.g. activation with systemd-run # expects access to environment variables like LOCALE_ARCHIVE, # NIXOS_NO_CHECK. # For now, for anyone that the above warn bothers you, please # use `sudo nixos-rebuild` instead of `--sudo` flag. run_args = ["sudo", "--preserve-env", *sudo_args, *run_args] final_args = ["sudo", *sudo_args, *run_args] # No need to pass env to subprocess.run; keep sudo's own env # default. popen_env = None else: # Non-sudo local: we can fully control the environment with # subprocess.run(env=...) final_args = run_args popen_env = None if env is None else resolved_env logger.debug( "calling run with args=%r, kwargs=%r, extra_env=%r", run_args, "calling run with args=%r, kwargs=%r, env=%r", _sanitize_env_run_args(remote_run_args if remote else run_args), kwargs, extra_env, env, ) try: r = subprocess.run( run_args, final_args, check=check, env=env, env=popen_env, input=process_input, # Hope nobody is using NixOS with non-UTF8 encodings, but # "surrogateescape" should still work in those systems. text=True, errors="surrogateescape", **kwargs, Loading @@ -195,13 +233,95 @@ def run_wrapper( raise # SSH does not send the signals to the process when running without usage of # pseudo-TTY (that causes a whole other can of worms), so if the process is # long running (e.g.: a build) this will result in the underlying process # staying alive. # See: https://stackoverflow.com/a/44354466 # Issue: https://github.com/NixOS/nixpkgs/issues/403269 def _resolve_env(env: Mapping[str, EnvValue] | None) -> dict[str, str]: normalized = _normalize_env(env) return _resolve_env_local(normalized) def _normalize_env(env: Mapping[str, EnvValue] | None) -> dict[str, EnvValue]: """ Normalize env mapping, but preserve some environment variables by default. """ return {"PATH": PRESERVE_ENV, **(env or {})} def _resolve_env_local(env: dict[str, EnvValue]) -> dict[str, str]: """ Resolve env mapping where values can be: - PRESERVE_ENV: copy from current os.environ (if present) - str: explicit value """ result: dict[str, str] = {} for k, v in env.items(): if v == PRESERVE_ENV: cur = os.environ.get(k) if cur is not None: result[k] = cur else: result[k] = v return result def _prefix_env_cmd(cmd: Sequence[Arg], resolved_env: dict[str, str]) -> list[Arg]: """ Prefix a command with `env -i K=V ... -- <cmd...>` to set vars for the command. """ if not resolved_env: return list(cmd) assigns = [f"{k}={v}" for k, v in resolved_env.items()] return ["env", "-i", *assigns, *cmd] def _prefix_env_cmd_remote( cmd: Args, env: dict[str, EnvValue], ) -> list[Arg | _RawShellArg]: """ Prefix remote commands with env assignments. Preserve markers are expanded by the remote shell at execution time. """ assigns: list[str | _RawShellArg] = [] for k, v in env.items(): if v is PRESERVE_ENV: assigns.append(_RawShellArg(f"{k}=${{{k}-}}")) else: assigns.append(f"{k}={v}") return ["env", "-i", *assigns, *cmd] def _quote_remote_arg(arg: Arg | _RawShellArg) -> str: if isinstance(arg, _RawShellArg): return str(arg) return shlex.quote(str(arg)) def _sanitize_env_run_args(run_args: list[Arg] | list[Arg | _RawShellArg]) -> list[Arg]: """ Sanitize long or sensitive environment variables from logs. """ sanitized: list[Arg] = [] for value in run_args: if isinstance(value, str) and value.startswith("PATH="): sanitized.append("PATH=<PATH>") elif isinstance(value, str | bytes | os.PathLike): sanitized.append(value) else: sanitized.append(str(value)) return sanitized def _kill_long_running_ssh_process(args: Args, remote: Remote) -> None: """ SSH does not send the signals to the process when running without usage of pseudo-TTY (that causes a whole other can of worms), so if the process is long running (e.g.: a build) this will result in the underlying process staying alive. See: https://stackoverflow.com/a/44354466 Issue: https://github.com/NixOS/nixpkgs/issues/403269 """ logger.info("cleaning-up remote process, please wait...") # We need to escape both the shell and regex here (since pkill interprets Loading
pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py +55 −13 Original line number Diff line number Diff line Loading @@ -248,7 +248,6 @@ def test_execute_nix_boot(mock_run: Mock, tmp_path: Path) -> None: | { "env": { "NIXOS_INSTALL_BOOTLOADER": "0", "NIXOS_REBUILD_I_UNDERSTAND_THE_CONSEQUENCES_PLEASE_BREAK_MY_SYSTEM": "1", } } ), Loading Loading @@ -487,7 +486,6 @@ def test_execute_nix_switch_flake(mock_run: Mock, tmp_path: Path) -> None: call( [ "sudo", "--preserve-env", "nix-env", "-p", Path("/nix/var/nix/profiles/system"), Loading @@ -505,21 +503,15 @@ def test_execute_nix_switch_flake(mock_run: Mock, tmp_path: Path) -> None: call( [ "sudo", "--preserve-env", "env", "-i", "NIXOS_INSTALL_BOOTLOADER=1", *nr.nix.SWITCH_TO_CONFIGURATION_CMD_PREFIX, config_path / "bin/switch-to-configuration", "switch", ], check=True, **( DEFAULT_RUN_KWARGS | { "env": { "NIXOS_INSTALL_BOOTLOADER": "1", "NIXOS_REBUILD_I_UNDERSTAND_THE_CONSEQUENCES_PLEASE_BREAK_MY_SYSTEM": "1", } } ), **DEFAULT_RUN_KWARGS, ), ] ) Loading Loading @@ -625,6 +617,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "mktemp", "-d", "-t", Loading @@ -640,6 +635,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "nix-store", "--realise", str(config_path), Loading @@ -656,6 +654,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "readlink", "-f", "/tmp/tmpdir/config", Loading @@ -670,6 +671,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "rm", "-rf", "/tmp/tmpdir", Loading Loading @@ -699,6 +703,9 @@ def test_execute_nix_switch_build_target_host( "user@target-host", "--", "sudo", "env", "-i", "PATH=${PATH-}", "nix-env", "-p", "/nix/var/nix/profiles/system", Loading @@ -714,6 +721,9 @@ def test_execute_nix_switch_build_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@target-host", "--", "env", "-i", "PATH=${PATH-}", "test", "-d", "/run/systemd/system", Loading @@ -729,6 +739,10 @@ def test_execute_nix_switch_build_target_host( "--", "sudo", "env", "-i", "PATH=${PATH-}", "LOCALE_ARCHIVE=${LOCALE_ARCHIVE-}", "NIXOS_NO_CHECK=${NIXOS_NO_CHECK-}", "NIXOS_INSTALL_BOOTLOADER=0", *nr.nix.SWITCH_TO_CONFIGURATION_CMD_PREFIX, str(config_path / "bin/switch-to-configuration"), Loading Loading @@ -806,6 +820,9 @@ def test_execute_nix_switch_flake_target_host( "user@localhost", "--", "sudo", "env", "-i", "PATH=${PATH-}", "nix-env", "-p", "/nix/var/nix/profiles/system", Loading @@ -821,6 +838,9 @@ def test_execute_nix_switch_flake_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@localhost", "--", "env", "-i", "PATH=${PATH-}", "test", "-d", "/run/systemd/system", Loading @@ -836,6 +856,10 @@ def test_execute_nix_switch_flake_target_host( "--", "sudo", "env", "-i", "PATH=${PATH-}", "LOCALE_ARCHIVE=${LOCALE_ARCHIVE-}", "NIXOS_NO_CHECK=${NIXOS_NO_CHECK-}", "NIXOS_INSTALL_BOOTLOADER=0", *nr.nix.SWITCH_TO_CONFIGURATION_CMD_PREFIX, str(config_path / "bin/switch-to-configuration"), Loading Loading @@ -912,6 +936,9 @@ def test_execute_nix_switch_flake_build_host( *nr.process.SSH_DEFAULT_OPTS, "user@localhost", "--", "env", "-i", "PATH=${PATH-}", "nix", "--extra-experimental-features", "'nix-command flakes'", Loading Loading @@ -1121,6 +1148,9 @@ def test_execute_build_dry_run_build_and_target_remote( *nr.process.SSH_DEFAULT_OPTS, "user@build-host", "--", "env", "-i", "PATH=${PATH-}", "nix", "--extra-experimental-features", "'nix-command flakes'", Loading Loading @@ -1304,7 +1334,6 @@ def test_execute_switch_store_path(mock_run: Mock, tmp_path: Path) -> None: | { "env": { "NIXOS_INSTALL_BOOTLOADER": "0", "NIXOS_REBUILD_I_UNDERSTAND_THE_CONSEQUENCES_PLEASE_BREAK_MY_SYSTEM": "1", } } ), Loading Loading @@ -1354,6 +1383,9 @@ def test_execute_switch_store_path_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@remote-host", "--", "env", "-i", "PATH=${PATH-}", "test", "-f", str(config_path / "nixos-version"), Loading @@ -1368,6 +1400,9 @@ def test_execute_switch_store_path_target_host( "user@remote-host", "--", "sudo", "env", "-i", "PATH=${PATH-}", "nix-env", "-p", "/nix/var/nix/profiles/system", Loading @@ -1383,6 +1418,9 @@ def test_execute_switch_store_path_target_host( *nr.process.SSH_DEFAULT_OPTS, "user@remote-host", "--", "env", "-i", "PATH=${PATH-}", "test", "-d", "/run/systemd/system", Loading @@ -1398,6 +1436,10 @@ def test_execute_switch_store_path_target_host( "--", "sudo", "env", "-i", "PATH=${PATH-}", "LOCALE_ARCHIVE=${LOCALE_ARCHIVE-}", "NIXOS_NO_CHECK=${NIXOS_NO_CHECK-}", "NIXOS_INSTALL_BOOTLOADER=0", *nr.nix.SWITCH_TO_CONFIGURATION_CMD_PREFIX, str(config_path / "bin/switch-to-configuration"), Loading
pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py +26 −18 Original line number Diff line number Diff line Loading @@ -134,9 +134,7 @@ def test_build_remote( "user@host", Path("/path/to/file"), ], extra_env={ "NIX_SSHOPTS": " ".join(["--ssh opts", *p.SSH_DEFAULT_OPTS]) }, env={"NIX_SSHOPTS": " ".join(["--ssh opts", *p.SSH_DEFAULT_OPTS])}, ), call( ["mktemp", "-d", "-t", "nixos-rebuild.XXXXX"], Loading Loading @@ -208,9 +206,7 @@ def test_build_remote_flake( "user@host", Path("/path/to/file"), ], extra_env={ "NIX_SSHOPTS": " ".join(["--ssh opts", *p.SSH_DEFAULT_OPTS]) }, env={"NIX_SSHOPTS": " ".join(["--ssh opts", *p.SSH_DEFAULT_OPTS])}, ), call( [ Loading Loading @@ -241,7 +237,7 @@ def test_copy_closure(monkeypatch: MonkeyPatch) -> None: n.copy_closure(closure, target_host) mock_run.assert_called_with( ["nix-copy-closure", "--to", "user@target.host", closure], extra_env={"NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS)}, env={"NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS)}, ) monkeypatch.setenv("NIX_SSHOPTS", "--ssh build-opt") Loading @@ -249,15 +245,11 @@ def test_copy_closure(monkeypatch: MonkeyPatch) -> None: n.copy_closure(closure, None, build_host, {"copy_flag": True}) mock_run.assert_called_with( ["nix-copy-closure", "--copy-flag", "--from", "user@build.host", closure], extra_env={ "NIX_SSHOPTS": " ".join(["--ssh build-opt", *p.SSH_DEFAULT_OPTS]) }, env={"NIX_SSHOPTS": " ".join(["--ssh build-opt", *p.SSH_DEFAULT_OPTS])}, ) monkeypatch.setenv("NIX_SSHOPTS", "--ssh build-target-opt") extra_env = { "NIX_SSHOPTS": " ".join(["--ssh build-target-opt", *p.SSH_DEFAULT_OPTS]) } env = {"NIX_SSHOPTS": " ".join(["--ssh build-target-opt", *p.SSH_DEFAULT_OPTS])} with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run: n.copy_closure(closure, target_host, build_host, {"copy_flag": True}) mock_run.assert_called_with( Loading @@ -273,7 +265,7 @@ def test_copy_closure(monkeypatch: MonkeyPatch) -> None: "ssh://user@target.host", closure, ], extra_env=extra_env, env=env, ) Loading Loading @@ -723,7 +715,11 @@ def test_switch_to_configuration_without_systemd_run( ) mock_run.assert_called_with( [profile_path / "bin/switch-to-configuration", "switch"], extra_env={"NIXOS_INSTALL_BOOTLOADER": "0"}, env={ "LOCALE_ARCHIVE": p.PRESERVE_ENV, "NIXOS_NO_CHECK": p.PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "0", }, sudo=False, remote=None, ) Loading Loading @@ -760,7 +756,11 @@ def test_switch_to_configuration_without_systemd_run( config_path / "specialisation/special/bin/switch-to-configuration", "test", ], extra_env={"NIXOS_INSTALL_BOOTLOADER": "1"}, env={ "LOCALE_ARCHIVE": p.PRESERVE_ENV, "NIXOS_NO_CHECK": p.PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "1", }, sudo=True, remote=target_host, ) Loading Loading @@ -791,7 +791,11 @@ def test_switch_to_configuration_with_systemd_run( profile_path / "bin/switch-to-configuration", "switch", ], extra_env={"NIXOS_INSTALL_BOOTLOADER": "0"}, env={ "LOCALE_ARCHIVE": p.PRESERVE_ENV, "NIXOS_NO_CHECK": p.PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "0", }, sudo=False, remote=None, ) Loading @@ -816,7 +820,11 @@ def test_switch_to_configuration_with_systemd_run( config_path / "specialisation/special/bin/switch-to-configuration", "test", ], extra_env={"NIXOS_INSTALL_BOOTLOADER": "1"}, env={ "LOCALE_ARCHIVE": p.PRESERVE_ENV, "NIXOS_NO_CHECK": p.PRESERVE_ENV, "NIXOS_INSTALL_BOOTLOADER": "1", }, sudo=True, remote=target_host, ) Loading