From 2b2d9135c52d0bf59890b378d1732150000f5c5c Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Mon, 1 Dec 2025 20:22:50 -0500 Subject: [PATCH] starting safe_reboot --- common/global/default.nix | 6 ++ common/global/safe_reboot.nix | 56 +++++++++++++++++ python/tools/safe_reboot.py | 111 ++++++++++++++++++++++++++++++++++ tests/test_safe_reboot.py | 96 +++++++++++++++++++++++++++++ 4 files changed, 269 insertions(+) create mode 100644 common/global/safe_reboot.nix create mode 100644 python/tools/safe_reboot.py create mode 100644 tests/test_safe_reboot.py diff --git a/common/global/default.nix b/common/global/default.nix index 3134a4e..5a77207 100644 --- a/common/global/default.nix +++ b/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; diff --git a/common/global/safe_reboot.nix b/common/global/safe_reboot.nix new file mode 100644 index 0000000..300f001 --- /dev/null +++ b/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; + }; + }; + }; +} diff --git a/python/tools/safe_reboot.py b/python/tools/safe_reboot.py new file mode 100644 index 0000000..8fdcc18 --- /dev/null +++ b/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() diff --git a/tests/test_safe_reboot.py b/tests/test_safe_reboot.py new file mode 100644 index 0000000..46371c7 --- /dev/null +++ b/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()