mirror of
https://github.com/RichieCahill/dotfiles.git
synced 2026-04-21 14:49:10 -04:00
112 lines
3.0 KiB
Python
112 lines
3.0 KiB
Python
"""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()
|