clan-cli: Move flash.py to clan_lib/flash

pkgs/clan-cli/clan_lib/flash/automount.py

@@ -0,0 +1,63 @@
+import logging
+import shutil
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+
+from clan_lib.cmd import Log, RunOpts, run
+from clan_lib.errors import ClanError
+from clan_lib.machines.machines import Machine
+
+log = logging.getLogger(__name__)
+
+
+@contextmanager
+def pause_automounting(
+    devices: list[Path], machine: Machine, request_graphical: bool = False
+) -> Generator[None]:
+    """
+    Pause automounting on the device for the duration of this context
+    manager
+    """
+
+    if shutil.which("udevadm") is None:
+        msg = "udev is required to disable automounting"
+        log.warning(msg)
+        yield None
+        return
+
+    inhibit_path = Path(__file__).parent / "inhibit.sh"
+    if not inhibit_path.exists():
+        msg = f"{inhibit_path} not found"
+        raise ClanError(msg)
+
+    str_devs = [str(dev) for dev in devices]
+    cmd = [str(inhibit_path), "enable", *str_devs]
+
+    result = run(
+        cmd,
+        RunOpts(
+            log=Log.BOTH,
+            check=False,
+            needs_user_terminal=True,
+            prefix=machine.name,
+            requires_root_perm=True,
+            graphical_perm=request_graphical,
+        ),
+    )
+    if result.returncode != 0:
+        machine.error("Failed to inhibit automounting")
+    yield None
+    cmd = [str(inhibit_path), "disable", *str_devs]
+    result = run(
+        cmd,
+        RunOpts(
+            log=Log.BOTH,
+            check=False,
+            prefix=machine.name,
+            requires_root_perm=True,
+            graphical_perm=request_graphical,
+        ),
+    )
+    if result.returncode != 0:
+        machine.error("Failed to re-enable automounting")

pkgs/clan-cli/clan_lib/flash/flash.py

@@ -0,0 +1,176 @@
+import json
+import logging
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+from clan_cli.facts.generate import generate_facts
+from clan_cli.vars.generate import generate_vars
+from clan_cli.vars.upload import populate_secret_vars
+
+from clan_lib.api import API
+from clan_lib.cmd import Log, RunOpts, cmd_with_root, run
+from clan_lib.errors import ClanError
+from clan_lib.machines.machines import Machine
+from clan_lib.nix import nix_shell
+
+from .automount import pause_automounting
+from .list import list_keymaps, list_languages
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class SystemConfig:
+    language: str | None = field(default=None)
+    keymap: str | None = field(default=None)
+    ssh_keys_path: list[str] | None = field(default=None)
+
+
+@dataclass
+class Disk:
+    name: str
+    device: str
+
+
+# TODO: unify this with machine install
+@API.register
+def run_machine_flash(
+    machine: Machine,
+    *,
+    mode: str,
+    disks: list[Disk],
+    system_config: SystemConfig,
+    dry_run: bool,
+    write_efi_boot_entries: bool,
+    debug: bool,
+    extra_args: list[str] | None = None,
+    graphical: bool = False,
+) -> None:
+    """Flash a machine with the given configuration.
+    Args:
+        machine: The Machine instance to flash.
+        mode: The mode to use for flashing (e.g., "install", "reinstall
+        disks: List of Disk instances representing the disks to flash.
+        system_config: SystemConfig instance containing language, keymap, and SSH keys.
+        dry_run: If True, perform a dry run without making changes.
+        write_efi_boot_entries: If True, write EFI boot entries.
+        debug: If True, enable debug mode.
+        extra_args: Additional arguments to pass to the disko-install command.
+        graphical: If True, run the command in graphical mode.
+    Raises:
+        ClanError: If the language or keymap is invalid, or if there are issues with
+            reading SSH keys, or if disko-install fails.
+    """
+    devices = [Path(disk.device) for disk in disks]
+    with pause_automounting(devices, machine, request_graphical=graphical):
+        if extra_args is None:
+            extra_args = []
+        system_config_nix: dict[str, Any] = {}
+
+        generate_facts([machine])
+        generate_vars([machine])
+
+        if system_config.language:
+            if system_config.language not in list_languages():
+                msg = (
+                    f"Language '{system_config.language}' is not a valid language. "
+                    f"Run 'clan flash list languages' to see a list of possible languages."
+                )
+                raise ClanError(msg)
+            system_config_nix["i18n"] = {"defaultLocale": system_config.language}
+
+        if system_config.keymap:
+            if system_config.keymap not in list_keymaps():
+                msg = (
+                    f"Keymap '{system_config.keymap}' is not a valid keymap. "
+                    f"Run 'clan flash list keymaps' to see a list of possible keymaps."
+                )
+                raise ClanError(msg)
+            system_config_nix["console"] = {"keyMap": system_config.keymap}
+            system_config_nix["services"] = {
+                "xserver": {"xkb": {"layout": system_config.keymap}}
+            }
+
+        if system_config.ssh_keys_path:
+            root_keys = []
+            for key_path in (Path(x) for x in system_config.ssh_keys_path):
+                try:
+                    root_keys.append(key_path.read_text())
+                except OSError as e:
+                    msg = f"Cannot read SSH public key file: {key_path}: {e}"
+                    raise ClanError(msg) from e
+            system_config_nix["users"] = {
+                "users": {"root": {"openssh": {"authorizedKeys": {"keys": root_keys}}}}
+            }
+
+        from clan_cli.vars.generate import Generator
+
+        for generator in Generator.generators_from_flake(machine.name, machine.flake):
+            for file in generator.files:
+                if file.needed_for == "partitioning":
+                    msg = f"Partitioning time secrets are not supported with `clan flash write`: clan.core.vars.generators.{generator.name}.files.{file.name}"
+                    raise ClanError(msg)
+
+        with TemporaryDirectory(prefix="disko-install-") as _tmpdir:
+            tmpdir = Path(_tmpdir)
+            upload_dir = machine.secrets_upload_directory
+
+            if upload_dir.startswith("/"):
+                local_dir = tmpdir / upload_dir[1:]
+            else:
+                local_dir = tmpdir / upload_dir
+
+            local_dir.mkdir(parents=True)
+            machine.secret_facts_store.upload(local_dir)
+            populate_secret_vars(machine, local_dir)
+            disko_install = []
+
+            if os.geteuid() != 0:
+                wrapper = " ".join(
+                    [
+                        "disko_install=$(command -v disko-install);",
+                        "exec",
+                        *cmd_with_root(['"$disko_install" "$@"'], graphical=graphical),
+                    ]
+                )
+                disko_install.extend(["bash", "-c", wrapper])
+
+            disko_install.append("disko-install")
+            if write_efi_boot_entries:
+                disko_install.append("--write-efi-boot-entries")
+            if dry_run:
+                disko_install.append("--dry-run")
+            if debug:
+                disko_install.append("--debug")
+            for disk in disks:
+                disko_install.extend(["--disk", disk.name, disk.device])
+
+                log.info("Will flash disk %s: %s", disk.name, disk.device)
+
+            disko_install.extend(["--extra-files", str(local_dir), upload_dir])
+            disko_install.extend(["--flake", str(machine.flake) + "#" + machine.name])
+            disko_install.extend(["--mode", str(mode)])
+            disko_install.extend(
+                [
+                    "--system-config",
+                    json.dumps(system_config_nix),
+                ]
+            )
+            disko_install.extend(["--option", "dry-run", "true"])
+            disko_install.extend(extra_args)
+
+            cmd = nix_shell(
+                ["disko"],
+                disko_install,
+            )
+            run(
+                cmd,
+                RunOpts(
+                    log=Log.BOTH,
+                    error_msg=f"Failed to flash {machine}",
+                    needs_user_terminal=True,
+                ),
+            )

pkgs/clan-cli/clan_lib/flash/inhibit.sh

@@ -0,0 +1,78 @@
+#!/usr/bin/env bash
+
+# A function to log error messages
+log_fatal() {
+  echo "FATAL: $*" >&2
+}
+
+# A function to log warning messages
+log_warning() {
+  echo "WARNING: $*" >&2
+}
+
+
+# Function to enable inhibition (pause automounting)
+enable_inhibition() {
+  local devices=("$@")
+
+  if ! command -v udevadm &> /dev/null; then
+    log_fatal "udev is required to disable automounting"
+    exit 1
+  fi
+
+  if [ "$EUID" -ne 0 ]; then
+    log_fatal "Root privileges are required to disable automounting"
+    exit 1
+  fi
+
+  local rules_dir="/run/udev/rules.d"
+  mkdir -p "$rules_dir"
+
+  for device in "${devices[@]}"; do
+    local devpath="$device"
+    local rule_file="$rules_dir/90-udisks-inhibit-${devpath//\//_}.rules"
+    echo 'SUBSYSTEM=="block", ENV{DEVNAME}=="'"$devpath"'*", ENV{UDISKS_IGNORE}="1"' > "$rule_file"
+    sync
+  done
+
+  udevadm control --reload
+  udevadm trigger --settle --subsystem-match=block
+}
+
+# Function to disable inhibition (cleanup)
+disable_inhibition() {
+  local devices=("$@")
+  local rules_dir="/run/udev/rules.d"
+  
+  for device in "${devices[@]}"; do
+    local devpath="$device"
+    local rule_file="$rules_dir/90-udisks-inhibit-${devpath//\//_}.rules"
+    rm -f "$rule_file" || log_warning "Could not remove file: $rule_file"
+  done
+
+  udevadm control --reload || log_warning "Could not reload udev rules"
+  udevadm trigger --settle --subsystem-match=block || log_warning "Could not trigger udev settle"
+}
+
+if [ "$#" -lt 2 ]; then
+  echo "Usage: $0 [enable|disable] /dev/sdX ..."
+  exit 1
+fi
+
+action=$1
+shift
+devices=("$@")
+
+case "$action" in
+  enable)
+    enable_inhibition "${devices[@]}"
+    ;;
+  disable)
+    disable_inhibition "${devices[@]}"
+    ;;
+  *)
+    echo "Invalid action: $action"
+    echo "Usage: $0 [enable|disable] /dev/sdX ..."
+    exit 1
+    ;;
+esac

pkgs/clan-cli/clan_lib/flash/list.py

@@ -0,0 +1,73 @@
+import logging
+import os
+from pathlib import Path
+from typing import TypedDict
+
+from clan_lib.api import API
+from clan_lib.cmd import Log, RunOpts, run
+from clan_lib.errors import ClanError
+from clan_lib.nix import nix_build
+
+log = logging.getLogger(__name__)
+
+
+class FlashOptions(TypedDict):
+    languages: list[str]
+    keymaps: list[str]
+
+
+@API.register
+def get_machine_flash_options() -> FlashOptions:
+    """Retrieve available languages and keymaps for flash configuration.
+    Returns:
+        FlashOptions: A dictionary containing lists of available languages and keymaps.
+    Raises:
+        ClanError: If the locale file or keymaps directory does not exist.
+    """
+    return {"languages": list_languages(), "keymaps": list_keymaps()}
+
+
+def list_languages() -> list[str]:
+    cmd = nix_build(["nixpkgs#glibcLocales"])
+    result = run(cmd, RunOpts(log=Log.STDERR, error_msg="Failed to find glibc locales"))
+    locale_file = Path(result.stdout.strip()) / "share" / "i18n" / "SUPPORTED"
+
+    if not locale_file.exists():
+        msg = f"Locale file '{locale_file}' does not exist."
+        raise ClanError(msg)
+
+    with locale_file.open() as f:
+        lines = f.readlines()
+
+    languages = []
+    for line in lines:
+        if line.startswith("#"):
+            continue
+        if "SUPPORTED-LOCALES" in line:
+            continue
+        # Split by '/' and take the first part
+        language = line.split("/")[0].strip()
+        languages.append(language)
+
+    return languages
+
+
+def list_keymaps() -> list[str]:
+    cmd = nix_build(["nixpkgs#kbd"])
+    result = run(cmd, RunOpts(log=Log.STDERR, error_msg="Failed to find kbdinfo"))
+    keymaps_dir = Path(result.stdout.strip()) / "share" / "keymaps"
+
+    if not keymaps_dir.exists():
+        msg = f"Keymaps directory '{keymaps_dir}' does not exist."
+        raise ClanError(msg)
+
+    keymap_files = []
+
+    for _root, _, files in os.walk(keymaps_dir):
+        for file in files:
+            if file.endswith(".map.gz"):
+                # Remove '.map.gz' ending
+                name_without_ext = file[:-7]
+                keymap_files.append(name_without_ext)
+
+    return keymap_files