starting safe_reboot

common/global/default.nix

@@ -16,6 +16,7 @@
     ./nh.nix
     ./nix.nix
     ./programs.nix
+    ./safe_reboot.nix
     ./ssh.nix
     ./snapshot_manager.nix
   ];
@@ -49,6 +50,11 @@
       PYTHONPATH = "${inputs.self}/";
     };
 
+    safe_reboot = {
+      enable = lib.mkDefault true;
+      datasetPrefix = "root_pool/";
+    };
+
     zfs = {
       trim.enable = lib.mkDefault true;
       autoScrub.enable = lib.mkDefault true;

common/global/safe_reboot.nix

@@ -0,0 +1,56 @@
+{
+  config,
+  inputs,
+  lib,
+  pkgs,
+  ...
+}:
+let
+  cfg = config.services.safe_reboot;
+  python_command =
+    lib.escapeShellArgs (
+      [
+        "${pkgs.my_python}/bin/python"
+        "-m"
+        "python.tools.safe_reboot"
+      ]
+      ++ lib.optionals (cfg.drivePath != null) [ cfg.drivePath ]
+      ++ [
+        "--dataset-prefix"
+        cfg.datasetPrefix
+        "--check-only"
+      ]
+    );
+in
+{
+  options.services.safe_reboot = {
+    enable = lib.mkEnableOption "Safe reboot dataset/drive validation";
+    datasetPrefix = lib.mkOption {
+      type = lib.types.str;
+      default = "root_pool/";
+      description = "Dataset prefix that must have exec enabled before rebooting.";
+    };
+    drivePath = lib.mkOption {
+      type = lib.types.nullOr lib.types.str;
+      default = null;
+      description = "Drive path that must exist before rebooting. Set to null to skip.";
+    };
+  };
+
+  config = lib.mkIf cfg.enable {
+    systemd.services.safe-reboot-check = {
+      description = "Safe reboot validation";
+      before = [ "systemd-reboot.service" ];
+      wantedBy = [ "reboot.target" ];
+      partOf = [ "reboot.target" ];
+      path = [ pkgs.zfs ];
+      environment = {
+        PYTHONPATH = "${inputs.self}/";
+      };
+      serviceConfig = {
+        Type = "oneshot";
+        ExecStart = python_command;
+      };
+    };
+  };
+}

python/tools/safe_reboot.py

@@ -0,0 +1,111 @@
+"""Safe reboot helper."""
+
+from __future__ import annotations
+
+import logging
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated
+
+import typer
+
+from python.common import bash_wrapper, configure_logger
+from python.zfs import Dataset, get_datasets
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+logger = logging.getLogger(__name__)
+
+
+def get_root_pool_datasets(dataset_prefix: str) -> list[Dataset]:
+    """Return datasets that start with the provided prefix."""
+    return [dataset for dataset in get_datasets() if dataset.name.startswith(dataset_prefix)]
+
+
+def get_non_executable_datasets(datasets: Sequence[Dataset]) -> list[str]:
+    """Return dataset names that have exec disabled."""
+    return [dataset.name for dataset in datasets if dataset.exec.lower() != "on"]
+
+
+def drive_present(drive: str) -> bool:
+    """Check whether the provided drive exists."""
+    drive_path = drive.strip()
+    if not drive_path:
+        error = "Drive path cannot be empty"
+        raise ValueError(error)
+
+    return Path(drive_path).exists()
+
+
+def reboot_system() -> None:
+    """Call systemctl reboot."""
+    output, return_code = bash_wrapper("systemctl reboot")
+    if return_code != 0:
+        raise RuntimeError(output.strip() or "Failed to issue reboot command")
+
+
+def validate_state(drive: str | None, dataset_prefix: str) -> list[str]:
+    """Validate dataset and drive state."""
+    datasets = get_root_pool_datasets(dataset_prefix)
+
+    errors: list[str] = []
+    if not datasets:
+        errors.append(f"No datasets found with prefix {dataset_prefix}")
+    else:
+        non_exec_datasets = get_non_executable_datasets(datasets)
+        if non_exec_datasets:
+            errors.append(f"Datasets missing exec=on: {', '.join(non_exec_datasets)}")
+
+    if drive:
+        try:
+            if not drive_present(drive):
+                errors.append(f"Drive {drive} is not present")
+        except ValueError as err:
+            errors.append(str(err))
+
+    return errors
+
+
+def reboot(
+    drive: Annotated[str | None, typer.Argument(help="Drive that must exist before rebooting.")] = None,
+    dataset_prefix: Annotated[
+        str,
+        typer.Option(
+            "--dataset-prefix",
+            "-p",
+            help="Datasets with this prefix are validated.",
+        ),
+    ] = "root_pool/",
+    dry_run: Annotated[
+        bool,
+        typer.Option(
+            "--check-only",
+            help="Only validate state without issuing the reboot command.",
+        ),
+    ] = False,
+) -> None:
+    """Validate datasets and drive before rebooting."""
+    configure_logger()
+    logger.info("Starting safe reboot checks")
+
+    if errors := validate_state(drive, dataset_prefix):
+        for error in errors:
+            logger.error(error)
+        sys.exit(1)
+
+    if dry_run:
+        logger.info("All checks passed")
+        return
+
+    logger.info("All checks passed, issuing reboot")
+    reboot_system()
+
+
+def cli() -> None:
+    """CLI entry point."""
+    typer.run(reboot)
+
+
+if __name__ == "__main__":
+    cli()

tests/test_safe_reboot.py

@@ -0,0 +1,96 @@
+"""Tests for safe_reboot."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from python.tools.safe_reboot import reboot
+from python.zfs.dataset import Dataset
+
+if TYPE_CHECKING:
+    from pytest_mock import MockerFixture
+
+SAFE_REBOOT = "python.tools.safe_reboot"
+
+
+def create_dataset(mocker: MockerFixture, name: str, exec_state: str) -> Dataset:
+    """Create a mock dataset."""
+    dataset = mocker.MagicMock(spec=Dataset)
+    dataset.name = name
+    dataset.exec = exec_state
+    return dataset
+
+
+def test_reboot_reboots_when_checks_pass(mocker: MockerFixture) -> None:
+    """The command should reboot when all checks pass."""
+    dataset = create_dataset(mocker, "root_pool/root", "on")
+    mocker.patch(f"{SAFE_REBOOT}.get_datasets", return_value=(dataset,))
+    mocker.patch(f"{SAFE_REBOOT}.drive_present", return_value=True)
+    mock_bash = mocker.patch(f"{SAFE_REBOOT}.bash_wrapper", return_value=("", 0))
+
+    reboot("/dev/disk/root-drive")
+
+    mock_bash.assert_called_once_with("systemctl reboot")
+
+
+def test_reboot_reboots_without_drive_requirement(mocker: MockerFixture) -> None:
+    """The command should reboot even when no drive is provided."""
+    dataset = create_dataset(mocker, "root_pool/root", "on")
+    mocker.patch(f"{SAFE_REBOOT}.get_datasets", return_value=(dataset,))
+    mock_bash = mocker.patch(f"{SAFE_REBOOT}.bash_wrapper", return_value=("", 0))
+
+    reboot(None)
+
+    mock_bash.assert_called_once_with("systemctl reboot")
+
+
+def test_reboot_errors_on_non_exec_dataset(mocker: MockerFixture) -> None:
+    """The command should exit when a dataset lacks exec."""
+    dataset = create_dataset(mocker, "root_pool/root", "off")
+    mocker.patch(f"{SAFE_REBOOT}.get_datasets", return_value=(dataset,))
+    mocker.patch(f"{SAFE_REBOOT}.drive_present", return_value=True)
+    mocker.patch(f"{SAFE_REBOOT}.bash_wrapper", return_value=("", 0))
+
+    with pytest.raises(SystemExit) as excinfo:
+        reboot("/dev/disk/root-drive")
+
+    assert excinfo.value.code == 1
+
+
+def test_reboot_errors_when_driver_missing(mocker: MockerFixture) -> None:
+    """The command should exit when the requested driver is absent."""
+    dataset = create_dataset(mocker, "root_pool/root", "on")
+    mocker.patch(f"{SAFE_REBOOT}.get_datasets", return_value=(dataset,))
+    mocker.patch(f"{SAFE_REBOOT}.drive_present", return_value=False)
+    mocker.patch(f"{SAFE_REBOOT}.bash_wrapper", return_value=("", 0))
+
+    with pytest.raises(SystemExit) as excinfo:
+        reboot("/dev/disk/root-drive")
+
+    assert excinfo.value.code == 1
+
+
+def test_reboot_errors_when_no_datasets_found(mocker: MockerFixture) -> None:
+    """The command should exit when no datasets match the prefix."""
+    mocker.patch(f"{SAFE_REBOOT}.get_datasets", return_value=())
+    mocker.patch(f"{SAFE_REBOOT}.drive_present", return_value=True)
+    mocker.patch(f"{SAFE_REBOOT}.bash_wrapper", return_value=("", 0))
+
+    with pytest.raises(SystemExit) as excinfo:
+        reboot("/dev/disk/root-drive")
+
+    assert excinfo.value.code == 1
+
+
+def test_reboot_check_only_skips_reboot(mocker: MockerFixture) -> None:
+    """The command should only validate when --check-only is provided."""
+    dataset = create_dataset(mocker, "root_pool/root", "on")
+    mocker.patch(f"{SAFE_REBOOT}.get_datasets", return_value=(dataset,))
+    mocker.patch(f"{SAFE_REBOOT}.drive_present", return_value=True)
+    mock_bash = mocker.patch(f"{SAFE_REBOOT}.bash_wrapper", return_value=("", 0))
+
+    reboot("/dev/disk/root-drive", check_only=True)
+
+    mock_bash.assert_not_called()