From 210faf516688a2c66ed7bac07dd072d34b4f0b11 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Thu, 12 Dec 2024 17:33:40 +0000
Subject: [PATCH 01/12] nixos-rebuild-ng: change copy closure logic when
 copying from_host -> to_host

- If user is using Nix 2.18+ (the default), we will switch to use `nix
  copy` since it supports using `--from` and `--to` at the same time
- If user is using older versions of Nix, we will call
  `nix-copy-closure` twice, once with `--from` and another with `--to`

The reason for this change is because the previous logic, that SSH'd to
from_host to copy to to_host using `nix-copy-closure --to` means that
`from_host` needs to have credentials to communicate with `to_host`,
that is not always possible nor expected. It breaks even in simple cases
like `--from` and `--to` is the same host because it is not common for a
host to SSH'd itself. This also makes the behavior more consistent,
since `--from` and `--to` is interpreted from the eval host point of
view EXCEPT when they're used together.

This may of course break in the cases where this old behavior was
assumed, e.g.: deploying with a bastion host. I am not sure how common
this was, and I think this change in behavior should enable more cases
than break, but let's see. For that particular case, running
`nixos-rebuild-ng` inside SSH command (e.g.: `ssh build_host --
nixos-rebuild-ng switch --target-host target_host`) should do the
trick.
---
 pkgs/by-name/ni/nixos-rebuild-ng/package.nix  |  8 +-
 .../src/nixos_rebuild/__init__.py             | 14 ++--
 .../src/nixos_rebuild/constants.py            |  6 ++
 .../nixos-rebuild-ng/src/nixos_rebuild/nix.py | 71 +++++++++++------
 .../ni/nixos-rebuild-ng/src/tests/test_nix.py | 77 +++++++++++++------
 5 files changed, 120 insertions(+), 56 deletions(-)
 create mode 100644 pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/package.nix b/pkgs/by-name/ni/nixos-rebuild-ng/package.nix
index 7dc84de9984e..6cd7eca28691 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/package.nix
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/package.nix
@@ -19,6 +19,11 @@
 }:
 let
   executable = if withNgSuffix then "nixos-rebuild-ng" else "nixos-rebuild";
+  # This version is kind of arbitrary, we use some features that were
+  # implemented in newer versions of Nix, but not necessary 2.18.
+  # However, Lix is a fork of Nix 2.18, so this looks like a good version
+  # to cut specific functionality.
+  withNix218 = lib.versionAtLeast nix.version "2.18";
 in
 python3Packages.buildPythonApplication rec {
   pname = "nixos-rebuild-ng";
@@ -53,8 +58,9 @@ python3Packages.buildPythonApplication rec {
   ];
 
   postPatch = ''
-    substituteInPlace nixos_rebuild/__init__.py \
+    substituteInPlace nixos_rebuild/constants.py \
       --subst-var-by executable ${executable} \
+      --subst-var-by withNix218 ${lib.boolToString withNix218} \
       --subst-var-by withReexec ${lib.boolToString withReexec} \
       --subst-var-by withShellFiles ${lib.boolToString withShellFiles}
 
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index 49ba193a03ad..8ce9c9b2d7f0 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -9,6 +9,7 @@ from subprocess import CalledProcessError, run
 from typing import assert_never
 
 from . import nix
+from .constants import EXECUTABLE, WITH_NIX_2_18, WITH_REEXEC, WITH_SHELL_FILES
 from .models import Action, BuildAttr, Flake, NRError, Profile
 from .process import Remote, cleanup_ssh
 from .utils import Args, LogFormatter
@@ -16,12 +17,6 @@ from .utils import Args, LogFormatter
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 
-# Build-time flags
-# Strings to avoid breaking standalone (e.g.: `python -m nixos_rebuild`) usage
-EXECUTABLE = "@executable@"
-WITH_REEXEC = "@withReexec@"
-WITH_SHELL_FILES = "@withShellFiles@"
-
 
 def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentParser]]:
     common_flags = argparse.ArgumentParser(add_help=False)
@@ -187,7 +182,7 @@ def parse_args(
     }
 
     if args.help or args.action is None:
-        if WITH_SHELL_FILES == "true":
+        if WITH_SHELL_FILES:
             r = run(["man", "8", EXECUTABLE], check=False)
             parser.exit(r.returncode)
         else:
@@ -281,6 +276,9 @@ def reexec(
 def execute(argv: list[str]) -> None:
     args, args_groups = parse_args(argv)
 
+    if not WITH_NIX_2_18:
+        logger.warning("you're using Nix <2.18, some features will not work correctly")
+
     atexit.register(cleanup_ssh)
 
     common_flags = vars(args_groups["common_flags"])
@@ -303,7 +301,7 @@ def execute(argv: list[str]) -> None:
     # Re-exec to a newer version of the script before building to ensure we get
     # the latest fixes
     if (
-        WITH_REEXEC == "true"
+        WITH_REEXEC
         and can_run
         and not args.fast
         and not os.environ.get("_NIXOS_REBUILD_REEXEC")
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py
new file mode 100644
index 000000000000..883c8b41b360
--- /dev/null
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py
@@ -0,0 +1,6 @@
+# Build-time flags
+# Strings to avoid breaking standalone (e.g.: `python -m nixos_rebuild`) usage
+EXECUTABLE = "@executable@"
+WITH_NIX_2_18 = "@withNix218@" == "true"  # type: ignore
+WITH_REEXEC = "@withReexec@" == "true"  # type: ignore
+WITH_SHELL_FILES = "@withShellFiles@" == "true"  # type: ignore
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
index 3100496e2eee..bca9a5b12c71 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
@@ -7,6 +7,7 @@ from string import Template
 from subprocess import PIPE, CalledProcessError
 from typing import Final
 
+from .constants import WITH_NIX_2_18
 from .models import (
     Action,
     BuildAttr,
@@ -139,30 +140,56 @@ def copy_closure(
     """Copy a nix closure to or from host to localhost.
 
     Also supports copying a closure from a remote to another remote."""
-    host = to_host or from_host
-    if not host:
-        return
 
     sshopts = os.getenv("NIX_SSHOPTS", "")
-    run_wrapper(
-        [
-            "nix-copy-closure",
-            *dict_to_flags(copy_flags),
-            "--to" if to_host else "--from",
-            host.host,
-            closure,
-        ],
-        extra_env={
-            # Using raw NIX_SSHOPTS here to avoid messing up with the passed
-            # parameters, and we do not add the SSH_DEFAULT_OPTS in the remote
-            # to remote case, otherwise it will fail because of ControlPath
-            # will not exist in remote
-            "NIX_SSHOPTS": sshopts
-            if from_host and to_host
-            else " ".join(filter(lambda x: x, [*SSH_DEFAULT_OPTS, sshopts]))
-        },
-        remote=from_host if to_host else None,
-    )
+    extra_env = {
+        "NIX_SSHOPTS": " ".join(filter(lambda x: x, [*SSH_DEFAULT_OPTS, sshopts]))
+    }
+
+    def nix_copy_closure(host: Remote, to: bool) -> None:
+        run_wrapper(
+            [
+                "nix-copy-closure",
+                *dict_to_flags(copy_flags),
+                "--to" if to else "--from",
+                host.host,
+                closure,
+            ],
+            extra_env=extra_env,
+        )
+
+    def nix_copy(to_host: Remote, from_host: Remote) -> None:
+        run_wrapper(
+            [
+                "nix",
+                "copy",
+                "--from",
+                f"ssh://{from_host.host}",
+                "--to",
+                f"ssh://{to_host.host}",
+                closure,
+            ],
+            extra_env=extra_env,
+        )
+
+    match (to_host, from_host):
+        case (None, None):
+            return
+        case (Remote(_) as host, None) | (None, Remote(_) as host):
+            nix_copy_closure(host, to=bool(to_host))
+        case (Remote(_), Remote(_)):
+            if WITH_NIX_2_18:
+                # With newer Nix, use `nix copy` instead of `nix-copy-closure`
+                # since it supports `--to` and `--from` at the same time
+                # TODO: once we drop Nix 2.3 from nixpkgs, remove support for
+                # `nix-copy-closure`
+                nix_copy(to_host, from_host)
+            else:
+                # With older Nix, we need to copy from to local and local to
+                # host. This means it is slower and need additional disk space
+                # in local
+                nix_copy_closure(from_host, to=False)
+                nix_copy_closure(to_host, to=True)
 
 
 def edit(flake: Flake | None, **flake_flags: Args) -> None:
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
index a5310a1df8a8..0e1a0ecdcc68 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
@@ -112,7 +112,6 @@ def test_remote_build(mock_run: Any, monkeypatch: Any) -> None:
                 extra_env={
                     "NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS + ["--ssh opts"])
                 },
-                remote=None,
             ),
             call(
                 ["nix-store", "--realise", Path("/path/to/file"), "--build"],
@@ -166,7 +165,6 @@ def test_remote_build_flake(mock_run: Any, monkeypatch: Any) -> None:
                 extra_env={
                     "NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS + ["--ssh opts"])
                 },
-                remote=None,
             ),
             call(
                 [
@@ -185,37 +183,66 @@ def test_remote_build_flake(mock_run: Any, monkeypatch: Any) -> None:
     )
 
 
-@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
-def test_copy_closure(mock_run: Any, monkeypatch: Any) -> None:
+def test_copy_closure(monkeypatch: Any) -> None:
     closure = Path("/path/to/closure")
-    n.copy_closure(closure, None)
-    mock_run.assert_not_called()
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
+        n.copy_closure(closure, None)
+        mock_run.assert_not_called()
 
     target_host = m.Remote("user@target.host", [], None)
     build_host = m.Remote("user@build.host", [], 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)},
-        remote=None,
-    )
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
+        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)},
+        )
 
     monkeypatch.setenv("NIX_SSHOPTS", "--ssh build-opt")
-    n.copy_closure(closure, None, build_host)
-    mock_run.assert_called_with(
-        ["nix-copy-closure", "--from", "user@build.host", closure],
-        extra_env={"NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS + ["--ssh build-opt"])},
-        remote=None,
-    )
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
+        n.copy_closure(closure, None, build_host)
+        mock_run.assert_called_with(
+            ["nix-copy-closure", "--from", "user@build.host", closure],
+            extra_env={
+                "NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS + ["--ssh build-opt"])
+            },
+        )
 
     monkeypatch.setenv("NIX_SSHOPTS", "--ssh build-target-opt")
-    n.copy_closure(closure, target_host, build_host)
-    mock_run.assert_called_with(
-        ["nix-copy-closure", "--to", "user@target.host", closure],
-        remote=build_host,
-        extra_env={"NIX_SSHOPTS": "--ssh build-target-opt"},
-    )
+    monkeypatch.setattr(n, "WITH_NIX_2_18", True)
+    extra_env = {
+        "NIX_SSHOPTS": " ".join(p.SSH_DEFAULT_OPTS + ["--ssh build-target-opt"])
+    }
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
+        n.copy_closure(closure, target_host, build_host)
+        mock_run.assert_called_with(
+            [
+                "nix",
+                "copy",
+                "--from",
+                "ssh://user@build.host",
+                "--to",
+                "ssh://user@target.host",
+                closure,
+            ],
+            extra_env=extra_env,
+        )
+
+    monkeypatch.setattr(n, "WITH_NIX_2_18", False)
+    with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
+        n.copy_closure(closure, target_host, build_host)
+        mock_run.assert_has_calls(
+            [
+                call(
+                    ["nix-copy-closure", "--from", "user@build.host", closure],
+                    extra_env=extra_env,
+                ),
+                call(
+                    ["nix-copy-closure", "--to", "user@target.host", closure],
+                    extra_env=extra_env,
+                ),
+            ]
+        )
 
 
 @patch(get_qualified_name(n.run_wrapper, n), autospec=True)

From 0051a3dc9a83a1436aed830645f75eefeab52be1 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Thu, 12 Dec 2024 19:19:20 +0000
Subject: [PATCH 02/12] nixos-rebuild-ng: use -s as copy_flags

---
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py     | 8 +++++++-
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py        | 6 +++++-
 pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py  | 2 ++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index 8ce9c9b2d7f0..11ebd6dfdfa0 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -56,7 +56,13 @@ def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentPa
 
     copy_flags = argparse.ArgumentParser(add_help=False)
     copy_flags.add_argument(
-        "--use-substitutes", "--substitute-on-destination", "-s", action="store_true"
+        "--use-substitutes",
+        "--substitute-on-destination",
+        "-s",
+        action="store_true",
+        # `-s` is the destination since it has the same meaning in
+        # `nix-copy-closure` and `nix copy`
+        dest="s",
     )
 
     sub_parsers = {
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
index cd89435f2d4f..b1d94114fa51 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
@@ -1,5 +1,5 @@
 import logging
-from typing import TypeAlias, override
+from typing import TypeAlias, assert_never, override
 
 Args: TypeAlias = bool | str | list[str] | int | None
 
@@ -25,6 +25,8 @@ def dict_to_flags(d: dict[str, Args]) -> list[str]:
         match value:
             case None | False | 0 | []:
                 continue
+            case True if len(key) == 1:
+                flags.append(f"-{key}")
             case True:
                 flags.append(flag)
             case int():
@@ -36,4 +38,6 @@ def dict_to_flags(d: dict[str, Args]) -> list[str]:
                 flags.append(flag)
                 for v in value:
                     flags.append(v)
+            case _:
+                assert_never(value)
     return flags
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
index 0e5eb7437f65..36b8024a4e56 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
@@ -9,6 +9,7 @@ def test_dict_to_flags() -> None:
             "test_flag_3": "value",
             "test_flag_4": ["v1", "v2"],
             "test_flag_5": None,
+            "t": True,
             "verbose": 5,
         }
     )
@@ -19,6 +20,7 @@ def test_dict_to_flags() -> None:
         "--test-flag-4",
         "v1",
         "v2",
+        "-t",
         "-vvvvv",
     ]
     r2 = u.dict_to_flags({"verbose": 0, "empty_list": []})

From 84c6a4e1cc80bc1cdba942e093fc7665c156e6c3 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Thu, 12 Dec 2024 19:21:18 +0000
Subject: [PATCH 03/12] nixos-rebuild-ng: add --debug flag

---
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py      | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index 11ebd6dfdfa0..59e563a7ae22 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -81,6 +81,9 @@ def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentPa
         allow_abbrev=False,
     )
     main_parser.add_argument("--help", "-h", action="store_true", help="Show manpage")
+    main_parser.add_argument(
+        "--debug", action="store_true", help="Enable debug logging"
+    )
     main_parser.add_argument(
         "--file", "-f", help="Enable and build the NixOS system from the specified file"
     )
@@ -198,8 +201,8 @@ def parse_args(
     def parser_warn(msg: str) -> None:
         print(f"{parser.prog}: warning: {msg}", file=sys.stderr)
 
-    # This flag affects both nix and this script
-    if args.verbose:
+    # verbose affects both nix commands and this script, debug only this script
+    if args.verbose or args.debug:
         logger.setLevel(logging.DEBUG)
 
     # https://github.com/NixOS/nixpkgs/blob/master/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh#L56

From 49afd97aab9c4504bc9c81f5f8be0dfccd9d5ea4 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Thu, 12 Dec 2024 22:06:36 +0000
Subject: [PATCH 04/12] nixos-rebuild-ng: run list_generations in parallel

---
 .../nixos-rebuild-ng/src/nixos_rebuild/nix.py | 33 +++++++++++--------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
index bca9a5b12c71..54ba8ab6d6b4 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
@@ -1,5 +1,6 @@
 import logging
 import os
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from importlib.resources import files
 from pathlib import Path
@@ -337,9 +338,8 @@ def list_generations(profile: Profile) -> list[GenerationJson]:
     Will be formatted in a way that is expected by the output of
     `nixos-rebuild list-generations --json`.
     """
-    generations = get_generations(profile)
-    result = []
-    for generation in reversed(generations):
+
+    def get_generation_info(generation: Generation) -> GenerationJson:
         generation_path = (
             profile.path.parent / f"{profile.path.name}-{generation.id}-link"
         )
@@ -367,19 +367,24 @@ def list_generations(profile: Profile) -> list[GenerationJson]:
             logger.debug("could not get configuration revision: %s", ex)
             configuration_revision = "Unknown"
 
-        result.append(
-            GenerationJson(
-                generation=generation.id,
-                date=generation.timestamp,
-                nixosVersion=nixos_version,
-                kernelVersion=kernel_version,
-                configurationRevision=configuration_revision,
-                specialisations=specialisations,
-                current=generation.current,
-            )
+        return GenerationJson(
+            generation=generation.id,
+            date=generation.timestamp,
+            nixosVersion=nixos_version,
+            kernelVersion=kernel_version,
+            configurationRevision=configuration_revision,
+            specialisations=specialisations,
+            current=generation.current,
         )
 
-    return result
+    # This can be surprisingly slow, especially with lots of generations,
+    # but it is basically IO work so we can run in parallel
+    with ThreadPoolExecutor() as executor:
+        return sorted(
+            executor.map(get_generation_info, get_generations(profile)),
+            key=lambda x: x["generation"],
+            reverse=True,
+        )
 
 
 def repl(attr: str, build_attr: BuildAttr, **nix_flags: Args) -> None:

From df1de458e0a18f1f3bde8b95c1910d4ccf216348 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Fri, 13 Dec 2024 20:43:19 +0000
Subject: [PATCH 05/12] nixos-rebuild-ng: remove --verbose flag hack

---
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py        | 4 ++--
 pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py  | 5 ++++-
 pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py      | 4 ++--
 pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py     | 5 ++++-
 4 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index 59e563a7ae22..b2f9cf22dfe5 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -20,7 +20,7 @@ logger.setLevel(logging.INFO)
 
 def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentParser]]:
     common_flags = argparse.ArgumentParser(add_help=False)
-    common_flags.add_argument("--verbose", "-v", action="count", default=0)
+    common_flags.add_argument("--verbose", "-v", action="count", dest="v", default=0)
     common_flags.add_argument("--max-jobs", "-j")
     common_flags.add_argument("--cores")
     common_flags.add_argument("--log-format")
@@ -202,7 +202,7 @@ def parse_args(
         print(f"{parser.prog}: warning: {msg}", file=sys.stderr)
 
     # verbose affects both nix commands and this script, debug only this script
-    if args.verbose or args.debug:
+    if args.v or args.debug:
         logger.setLevel(logging.DEBUG)
 
     # https://github.com/NixOS/nixpkgs/blob/master/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh#L56
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
index b1d94114fa51..87d3f21596a3 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
@@ -29,8 +29,11 @@ def dict_to_flags(d: dict[str, Args]) -> list[str]:
                 flags.append(f"-{key}")
             case True:
                 flags.append(flag)
+            case int() if len(key) == 1:
+                flags.append(f"-{key * value}")
             case int():
-                flags.append(f"-{key[0] * value}")
+                for i in range(value):
+                    flags.append(flag)
             case str():
                 flags.append(flag)
                 flags.append(value)
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
index 2c06b128da86..dff9c24fb1d2 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
@@ -67,12 +67,12 @@ def test_parse_args() -> None:
         ]
     )
     assert nr.logger.level == logging.DEBUG
-    assert r2.verbose == 3
+    assert r2.v == 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
+    assert g2["common_flags"].v == 3
 
 
 @patch.dict(nr.process.os.environ, {}, clear=True)
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
index 36b8024a4e56..5f73c33acf98 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
@@ -10,7 +10,8 @@ def test_dict_to_flags() -> None:
             "test_flag_4": ["v1", "v2"],
             "test_flag_5": None,
             "t": True,
-            "verbose": 5,
+            "v": 5,
+            "verbose": 2,
         }
     )
     assert r1 == [
@@ -22,6 +23,8 @@ def test_dict_to_flags() -> None:
         "v2",
         "-t",
         "-vvvvv",
+        "--verbose",
+        "--verbose",
     ]
     r2 = u.dict_to_flags({"verbose": 0, "empty_list": []})
     assert r2 == []

From abda7147ff534761e1059d811edf09e97c12e871 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Thu, 12 Dec 2024 23:36:02 +0000
Subject: [PATCH 06/12] nixos-rebuild-ng: remove tabulate as dependency

---
 pkgs/by-name/ni/nixos-rebuild-ng/package.nix  |  7 ---
 .../src/nixos_rebuild/__init__.py             | 16 +-----
 .../src/nixos_rebuild/utils.py                | 51 ++++++++++++++++++-
 .../nixos-rebuild-ng/src/tests/test_utils.py  | 21 ++++++++
 4 files changed, 72 insertions(+), 23 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/package.nix b/pkgs/by-name/ni/nixos-rebuild-ng/package.nix
index 6cd7eca28691..df33313193fe 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/package.nix
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/package.nix
@@ -35,10 +35,6 @@ python3Packages.buildPythonApplication rec {
     setuptools
   ];
 
-  dependencies = with python3Packages; [
-    tabulate
-  ];
-
   nativeBuildInputs = lib.optionals withShellFiles [
     installShellFiles
     python3Packages.shtab
@@ -94,9 +90,6 @@ python3Packages.buildPythonApplication rec {
           mypy
           pytest
           ruff
-          types-tabulate
-          # dependencies
-          tabulate
         ]
       );
     in
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index b2f9cf22dfe5..dec1a5a252d9 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -12,7 +12,7 @@ from . import nix
 from .constants import EXECUTABLE, WITH_NIX_2_18, WITH_REEXEC, WITH_SHELL_FILES
 from .models import Action, BuildAttr, Flake, NRError, Profile
 from .process import Remote, cleanup_ssh
-from .utils import Args, LogFormatter
+from .utils import Args, LogFormatter, tabulate
 
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
@@ -448,8 +448,6 @@ def execute(argv: list[str]) -> None:
             if args.json:
                 print(json.dumps(generations, indent=2))
             else:
-                from tabulate import tabulate
-
                 headers = {
                     "generation": "Generation",
                     "date": "Build-date",
@@ -459,17 +457,7 @@ def execute(argv: list[str]) -> None:
                     "specialisations": "Specialisation",
                     "current": "Current",
                 }
-                # Not exactly the same format as legacy nixos-rebuild but close
-                # enough
-                table = tabulate(
-                    generations,
-                    headers=headers,
-                    tablefmt="plain",
-                    numalign="left",
-                    stralign="left",
-                    disable_numparse=True,
-                )
-                print(table)
+                print(tabulate(generations, headers=headers))
         case Action.REPL:
             if flake:
                 nix.repl_flake("toplevel", flake, **flake_build_flags)
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
index 87d3f21596a3..11a0d9188294 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/utils.py
@@ -1,5 +1,6 @@
 import logging
-from typing import TypeAlias, assert_never, override
+from collections.abc import Mapping, Sequence
+from typing import Any, TypeAlias, assert_never, override
 
 Args: TypeAlias = bool | str | list[str] | int | None
 
@@ -18,7 +19,7 @@ class LogFormatter(logging.Formatter):
         return formatter.format(record)
 
 
-def dict_to_flags(d: dict[str, Args]) -> list[str]:
+def dict_to_flags(d: Mapping[str, Args]) -> list[str]:
     flags = []
     for key, value in d.items():
         flag = f"--{'-'.join(key.split('_'))}"
@@ -44,3 +45,49 @@ def dict_to_flags(d: dict[str, Args]) -> list[str]:
             case _:
                 assert_never(value)
     return flags
+
+
+def remap_dicts(
+    dicts: Sequence[Mapping[str, Any]],
+    mappings: Mapping[str, str],
+) -> list[dict[str, Any]]:
+    return [{mappings.get(k, k): v for k, v in d.items()} for d in dicts]
+
+
+def tabulate(
+    data: Sequence[Mapping[str, Any]],
+    headers: Mapping[str, str] | None = None,
+) -> str:
+    """Convert a sequence of mappings in a tabular-style format for terminal.
+
+    It expects that all mappings (dicts) have the same keys as the first one,
+    otherwise it will misbehave.
+    """
+    if not data:
+        return ""
+
+    if headers:
+        data = remap_dicts(data, headers)
+
+    data_headers = list(data[0].keys())
+
+    column_widths = [
+        max(
+            len(str(header)),
+            *(len(str(row.get(header, ""))) for row in data),
+        )
+        for header in data_headers
+    ]
+
+    def format_row(row: Mapping[str, Any]) -> str:
+        s = (2 * " ").join(
+            f"{str(row[header]).ljust(width)}"
+            for header, width in zip(data_headers, column_widths)
+        )
+        return s.strip()
+
+    result = [format_row(dict(zip(data_headers, data_headers)))]
+    for row in data:
+        result.append(format_row(row))
+
+    return "\n".join(result)
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
index 5f73c33acf98..b1cd52e0b121 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_utils.py
@@ -1,3 +1,5 @@
+import textwrap
+
 import nixos_rebuild.utils as u
 
 
@@ -28,3 +30,22 @@ def test_dict_to_flags() -> None:
     ]
     r2 = u.dict_to_flags({"verbose": 0, "empty_list": []})
     assert r2 == []
+
+
+def test_remap_dicts() -> None:
+    assert u.remap_dicts(
+        [{"foo": 1, "bar": True}, {"qux": "keep"}],
+        {"foo": "Foo", "bar": "Bar"},
+    ) == [{"Foo": 1, "Bar": True}, {"qux": "keep"}]
+
+
+def test_tabulate() -> None:
+    assert u.tabulate([]) == ""
+    assert u.tabulate([{}]) == "\n"
+    assert u.tabulate(
+        [{"foo": 12345, "bar": ["abc", "cde"]}, {"foo": 345, "bar": 456}],
+        {"foo": "Foo", "bar": "Bar"},
+    ) == textwrap.dedent("""\
+        Foo    Bar
+        12345  ['abc', 'cde']
+        345    456""")

From b1203ae015f84805f73b5824d557faae96c98826 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Fri, 13 Dec 2024 22:13:08 +0000
Subject: [PATCH 07/12] nixos-rebuild-ng: add "Breaking changes" section

---
 pkgs/by-name/ni/nixos-rebuild-ng/README.md | 41 ++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/README.md b/pkgs/by-name/ni/nixos-rebuild-ng/README.md
index 0bf680f0cced..f6aa75a12b35 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/README.md
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/README.md
@@ -101,6 +101,47 @@ ruff check --fix .
 ruff format .
 ```
 
+## Breaking changes
+
+While `nixos-rebuild-ng` tries to be as much of a clone of the original as
+possible, there are still some breaking changes that were done in other to
+improve the user experience. If they break your workflow in some way that is
+not possible to fix, please open an issue and we can discuss a solution.
+
+- For `--build-host` and `--target-host`, `nixos-rebuild-ng` does not allocate
+  a pseudo-TTY via SSH (e.g.: `ssh -t`) anymore. The reason for this is because
+  pseudo-TTY breaks some expectations from SSH, like it mangles stdout and
+  stderr, and can
+  [break terminal output](https://github.com/NixOS/nixpkgs/issues/336967) in
+  some situations.
+  The issue is that `sudo` needs a TTY to ask for password, otherwise it will
+  fail. The solution for this is a new flag, `--ask-sudo-password`, that when
+  used with `--target-host` (`--build-host` doesn't need `sudo`), will ask for
+  the `sudo` password for the target host using Python's
+  [getpass](https://docs.python.org/3/library/getpass.html) and forward it to
+  every `sudo` request. Keep in mind that there is no check, so if you type
+  your password wrong, it will fail during activation (this can be improved
+  though)
+- When `--build-host` and `--target-host` is used together, we will use `nix
+  copy` (or 2 `nix-copy-closure` if you're using Nix <2.18) instead of SSH'ing
+  to build host and using `nix-copy-closure --to target-host`. The reason for
+  this is documented in PR
+  [#364698](https://github.com/NixOS/nixpkgs/pull/364698). If you do need the
+  previous behavior, you can simulate it using `ssh build-host --
+  nixos-rebuild-ng switch --target-host target-host`. If that is not the case,
+  please open an issue
+- We do some additional validation of flags, like exiting with an error when
+  `--build-host` or `--target-host` is used with `repl`, since the user could
+  assume that the `repl` would be run remotely while it always run the local
+  machine. `nixos-rebuild` silently ignored those flags, so this
+  [may cause some issues](https://github.com/NixOS/nixpkgs/pull/363922) for
+  wrappers
+- For now we are not supporting `build-vm` or `build-vm-with-bootloader` with
+  `--build-host` anymore. Support for this should be easy to add if you have
+  a use case, the only reason is that this is not supported is because in the
+  original the result would not be symlinked in the current directory, making
+  it kind useless (unless you looked at the result in `/nix/store`)
+
 ## Caveats
 
 - Bugs in the profile manipulation can cause corruption of your profile that

From 11bbfac1c9f27310afc56114d145cfc53b57335e Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Sat, 14 Dec 2024 12:16:42 +0000
Subject: [PATCH 08/12] nixos-rebuild-ng: support --target-host for
 build-vm/build-vm-with-bootloader

---
 pkgs/by-name/ni/nixos-rebuild-ng/README.md    |  5 ---
 .../src/nixos_rebuild/__init__.py             | 34 +++++++++----------
 2 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/README.md b/pkgs/by-name/ni/nixos-rebuild-ng/README.md
index f6aa75a12b35..52df11d42400 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/README.md
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/README.md
@@ -136,11 +136,6 @@ not possible to fix, please open an issue and we can discuss a solution.
   machine. `nixos-rebuild` silently ignored those flags, so this
   [may cause some issues](https://github.com/NixOS/nixpkgs/pull/363922) for
   wrappers
-- For now we are not supporting `build-vm` or `build-vm-with-bootloader` with
-  `--build-host` anymore. Support for this should be easy to add if you have
-  a use case, the only reason is that this is not supported is because in the
-  original the result would not be symlinked in the current directory, making
-  it kind useless (unless you looked at the result in `/nix/store`)
 
 ## Caveats
 
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index dec1a5a252d9..d24b1a8afa71 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -240,6 +240,8 @@ def parse_args(
         Action.BUILD.value,
         Action.DRY_BUILD.value,
         Action.DRY_ACTIVATE.value,
+        Action.BUILD_VM.value,
+        Action.BUILD_VM_WITH_BOOTLOADER.value,
     ):
         parser.error(
             f"--target-host/--build-host is not supported with '{args.action}'"
@@ -337,13 +339,21 @@ def execute(argv: list[str]) -> None:
             | Action.BUILD
             | Action.DRY_BUILD
             | Action.DRY_ACTIVATE
+            | Action.BUILD_VM
+            | Action.BUILD_VM_WITH_BOOTLOADER
         ):
             logger.info("building the system configuration...")
 
-            attr = "config.system.build.toplevel"
             dry_run = action == Action.DRY_BUILD
             no_link = action in (Action.SWITCH, Action.BOOT)
             rollback = bool(args.rollback)
+            match action:
+                case Action.BUILD_VM:
+                    attr = "config.system.build.vm"
+                case Action.BUILD_VM_WITH_BOOTLOADER:
+                    attr = "config.system.build.vmWithBootLoader"
+                case _:
+                    attr = "config.system.build.toplevel"
 
             match (action, rollback, build_host, flake):
                 case (Action.SWITCH | Action.BOOT, True, _, _):
@@ -413,6 +423,7 @@ def execute(argv: list[str]) -> None:
                         target_host=target_host,
                         sudo=args.sudo,
                     )
+
             if action in (Action.SWITCH, Action.BOOT, Action.TEST, Action.DRY_ACTIVATE):
                 nix.switch_to_configuration(
                     path_to_config,
@@ -422,23 +433,12 @@ def execute(argv: list[str]) -> None:
                     specialisation=args.specialisation,
                     install_bootloader=args.install_bootloader,
                 )
-        case Action.BUILD_VM | Action.BUILD_VM_WITH_BOOTLOADER:
-            logger.info("building the system configuration...")
-            attr = "vm" if action == Action.BUILD_VM else "vmWithBootLoader"
-            if flake:
-                path_to_config = nix.build_flake(
-                    f"config.system.build.{attr}",
-                    flake,
-                    **flake_build_flags,
+            elif action in (Action.BUILD_VM, Action.BUILD_VM_WITH_BOOTLOADER):
+                # If you get `not-found`, please open an issue
+                vm_path = next(path_to_config.glob("bin/run-*-vm"), "not-found")
+                print(
+                    f"Done. The virtual machine can be started by running '{vm_path}'"
                 )
-            else:
-                path_to_config = nix.build(
-                    f"config.system.build.{attr}",
-                    build_attr,
-                    **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:
             nix.edit(flake, **flake_build_flags)
         case Action.DRY_RUN:

From 38c8c39b6d322634441addf3ba62dd93a7963173 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Sat, 14 Dec 2024 14:03:21 +0000
Subject: [PATCH 09/12] nixos-rebuild-ng: improve assertion for unreachable
 code

---
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py     | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index d24b1a8afa71..50c4d9ba787f 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -404,10 +404,14 @@ def execute(argv: list[str]) -> None:
                         dry_run=dry_run,
                         **build_flags,
                     )
-                case m:
+                case never:
                     # should never happen, but mypy is not smart enough to
                     # handle this with assert_never
-                    raise NRError(f"invalid match for build: {m}")
+                    # https://github.com/python/mypy/issues/16650
+                    # https://github.com/python/mypy/issues/16722
+                    raise AssertionError(
+                        f"expected code to be unreachable, but got: {never}"
+                    )
 
             if not rollback:
                 nix.copy_closure(

From 37422786020d0b66cb0e308dd34e08544116dd97 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Sun, 15 Dec 2024 00:24:36 +0000
Subject: [PATCH 10/12] nixos-rebuild-ng: capture output of `git rev-parse`

Fixt #365222.
---
 pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py | 3 ++-
 pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py   | 2 +-
 pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py    | 4 ++--
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
index 54ba8ab6d6b4..c7033808394b 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
@@ -246,7 +246,8 @@ def get_nixpkgs_rev(nixpkgs_path: Path | None) -> str | None:
         r = run_wrapper(
             ["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
             check=False,
-            stdout=PIPE,
+            # https://github.com/NixOS/nixpkgs/issues/365222
+            capture_output=True,
         )
     except FileNotFoundError:
         # Git is not included in the closure so we need to check
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
index dff9c24fb1d2..39a50ecc405a 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
@@ -109,7 +109,7 @@ def test_execute_nix_boot(mock_run: Any, tmp_path: Path) -> None:
             call(
                 ["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
                 check=False,
-                stdout=PIPE,
+                capture_output=True,
                 **DEFAULT_RUN_KWARGS,
             ),
             call(
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
index 0e1a0ecdcc68..3cee46045863 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
@@ -289,14 +289,14 @@ def test_get_nixpkgs_rev() -> None:
         mock_run.assert_called_with(
             ["git", "-C", path, "rev-parse", "--short", "HEAD"],
             check=False,
-            stdout=PIPE,
+            capture_output=True,
         )
 
     expected_calls = [
         call(
             ["git", "-C", path, "rev-parse", "--short", "HEAD"],
             check=False,
-            stdout=PIPE,
+            capture_output=True,
         ),
         call(
             ["git", "-C", path, "diff", "--quiet"],

From fb69f66d0c961c5791967b571a5f3eba34d25480 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Sun, 15 Dec 2024 05:47:33 +0000
Subject: [PATCH 11/12] nixos-rebuild-ng: WITH_NIX_2_18 should default to true

---
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py     | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py
index 883c8b41b360..b02d918ab3b4 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/constants.py
@@ -1,6 +1,9 @@
 # Build-time flags
-# Strings to avoid breaking standalone (e.g.: `python -m nixos_rebuild`) usage
+# Use strings to avoid breaking standalone (e.g.: `python -m nixos_rebuild`)
+# usage
 EXECUTABLE = "@executable@"
-WITH_NIX_2_18 = "@withNix218@" == "true"  # type: ignore
+# Use either `== "true"` if the default (e.g.: `python -m nixos_rebuld`) is
+# `False` or `!= "false"` if the default is `True`
+WITH_NIX_2_18 = "@withNix218@" != "false"  # type: ignore
 WITH_REEXEC = "@withReexec@" == "true"  # type: ignore
 WITH_SHELL_FILES = "@withShellFiles@" == "true"  # type: ignore

From 0212452e27161b583dfee83b5bf2f52141a379bb Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Sun, 15 Dec 2024 12:38:55 +0000
Subject: [PATCH 12/12] nixos-rebuild-ng: split get_generations in 2 functions

---
 .../nixos-rebuild-ng/src/nixos_rebuild/nix.py | 117 +++++++++---------
 .../ni/nixos-rebuild-ng/src/tests/test_nix.py |  58 ++++++---
 2 files changed, 98 insertions(+), 77 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
index c7033808394b..c1f2e5d5c711 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
@@ -266,68 +266,72 @@ def get_nixpkgs_rev(nixpkgs_path: Path | None) -> str | None:
         return None
 
 
-def _parse_generation_from_nix_store(path: Path, profile: Profile) -> Generation:
-    entry_id = path.name.split("-")[1]
-    current = path.name == profile.path.readlink().name
-    timestamp = datetime.fromtimestamp(path.stat().st_ctime).strftime(
-        "%Y-%m-%d %H:%M:%S"
-    )
-
-    return Generation(
-        id=int(entry_id),
-        timestamp=timestamp,
-        current=current,
-    )
-
-
-def _parse_generation_from_nix_env(line: str) -> Generation:
-    parts = line.split()
-
-    entry_id = parts[0]
-    timestamp = f"{parts[1]} {parts[2]}"
-    current = "(current)" in parts
-
-    return Generation(
-        id=int(entry_id),
-        timestamp=timestamp,
-        current=current,
-    )
-
-
-def get_generations(
-    profile: Profile,
-    target_host: Remote | None = None,
-    using_nix_env: bool = False,
-    sudo: bool = False,
-) -> list[Generation]:
+def get_generations(profile: Profile) -> list[Generation]:
     """Get all NixOS generations from profile.
 
     Includes generation ID (e.g.: 1, 2), timestamp (e.g.: when it was created)
     and if this is the current active profile or not.
-
-    If `lock_profile = True` this command will need root to run successfully.
     """
     if not profile.path.exists():
         raise NRError(f"no profile '{profile.name}' found")
 
-    result = []
-    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_wrapper(
-            ["nix-env", "-p", profile.path, "--list-generations"],
-            stdout=PIPE,
-            remote=target_host,
-            sudo=sudo,
+    def parse_path(path: Path, profile: Profile) -> Generation:
+        entry_id = path.name.split("-")[1]
+        current = path.name == profile.path.readlink().name
+        timestamp = datetime.fromtimestamp(path.stat().st_ctime).strftime(
+            "%Y-%m-%d %H:%M:%S"
         )
-        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)
+
+        return Generation(
+            id=int(entry_id),
+            timestamp=timestamp,
+            current=current,
+        )
+
+    return sorted(
+        [parse_path(p, profile) for p in profile.path.parent.glob("system-*-link")],
+        key=lambda d: d.id,
+    )
+
+
+def get_generations_from_nix_env(
+    profile: Profile,
+    target_host: Remote | None = None,
+    sudo: bool = False,
+) -> list[Generation]:
+    """Get all NixOS generations from profile with nix-env. Needs root.
+
+    Includes generation ID (e.g.: 1, 2), timestamp (e.g.: when it was created)
+    and if this is the current active profile or not.
+    """
+    if not profile.path.exists():
+        raise NRError(f"no profile '{profile.name}' found")
+
+    # Using `nix-env --list-generations` needs root to lock the profile
+    r = run_wrapper(
+        ["nix-env", "-p", profile.path, "--list-generations"],
+        stdout=PIPE,
+        remote=target_host,
+        sudo=sudo,
+    )
+
+    def parse_line(line: str) -> Generation:
+        parts = line.split()
+
+        entry_id = parts[0]
+        timestamp = f"{parts[1]} {parts[2]}"
+        current = "(current)" in parts
+
+        return Generation(
+            id=int(entry_id),
+            timestamp=timestamp,
+            current=current,
+        )
+
+    return sorted(
+        [parse_line(line) for line in r.stdout.splitlines()],
+        key=lambda d: d.id,
+    )
 
 
 def list_generations(profile: Profile) -> list[GenerationJson]:
@@ -437,11 +441,8 @@ def rollback_temporary_profile(
     sudo: bool,
 ) -> Path | None:
     "Rollback a temporary Nix profile, like one created by `nixos-rebuild test`."
-    generations = get_generations(
-        profile,
-        target_host=target_host,
-        using_nix_env=True,
-        sudo=sudo,
+    generations = get_generations_from_nix_env(
+        profile, target_host=target_host, sudo=sudo
     )
     previous_gen_id = None
     for generation in generations:
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
index 3cee46045863..34f47029c103 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
@@ -327,7 +327,7 @@ def test_get_nixpkgs_rev() -> None:
         mock_run.assert_has_calls(expected_calls)
 
 
-def test_get_generations_from_nix_store(tmp_path: Path) -> None:
+def test_get_generations(tmp_path: Path) -> None:
     nixos_path = tmp_path / "nixos-system"
     nixos_path.mkdir()
 
@@ -337,20 +337,17 @@ def test_get_generations_from_nix_store(tmp_path: Path) -> None:
     (tmp_path / "system-3-link").symlink_to(nixos_path)
     (tmp_path / "system-2-link").symlink_to(nixos_path)
 
-    assert n.get_generations(
-        m.Profile("system", tmp_path / "system"),
-        using_nix_env=False,
-    ) == [
+    assert n.get_generations(m.Profile("system", tmp_path / "system")) == [
         m.Generation(id=1, current=False, timestamp=ANY),
         m.Generation(id=2, current=True, timestamp=ANY),
         m.Generation(id=3, current=False, timestamp=ANY),
     ]
 
 
-@patch(
-    get_qualified_name(n.run_wrapper, n),
-    autospec=True,
-    return_value=CompletedProcess(
+def test_get_generations_from_nix_env(tmp_path: Path) -> None:
+    path = tmp_path / "test"
+    path.touch()
+    return_value = CompletedProcess(
         [],
         0,
         stdout=textwrap.dedent("""\
@@ -358,17 +355,40 @@ def test_get_generations_from_nix_store(tmp_path: Path) -> None:
         2083   2024-11-07 22:59:41
         2084   2024-11-07 23:54:17   (current)
         """),
-    ),
-)
-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), 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"),
-    ]
+    with patch(
+        get_qualified_name(n.run_wrapper, n), autospec=True, return_value=return_value
+    ) as mock_run:
+        assert n.get_generations_from_nix_env(m.Profile("system", path)) == [
+            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"),
+        ]
+        mock_run.assert_called_with(
+            ["nix-env", "-p", path, "--list-generations"],
+            stdout=PIPE,
+            remote=None,
+            sudo=False,
+        )
+
+    remote = m.Remote("user@host", [], "password")
+    with patch(
+        get_qualified_name(n.run_wrapper, n), autospec=True, return_value=return_value
+    ) as mock_run:
+        assert n.get_generations_from_nix_env(
+            m.Profile("system", path), remote, 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"),
+        ]
+        mock_run.assert_called_with(
+            ["nix-env", "-p", path, "--list-generations"],
+            stdout=PIPE,
+            remote=remote,
+            sudo=True,
+        )
 
 
 @patch(