dotfiles/python/installer/old_installer.py

"""Install NixOS on a ZFS pool."""

from __future__ import annotations

import curses
import logging
import sys
from collections import defaultdict
from os import getenv
from pathlib import Path
from random import getrandbits
from subprocess import PIPE, Popen, run
from time import sleep
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Sequence

logger = logging.getLogger(__name__)
ESCAPE_KEY = 27


def configure_logger(level: str = "INFO") -> None:
    """Configure the logger.

    Args:
        level (str, optional): The logging level. Defaults to "INFO".
    """
    logging.basicConfig(
        level=level,
        datefmt="%Y-%m-%dT%H:%M:%S%z",
        format="%(asctime)s %(levelname)s %(filename)s:%(lineno)d - %(message)s",
        handlers=[logging.StreamHandler(sys.stdout)],
    )


def bash_wrapper(command: str) -> str:
    """Execute a bash command and capture the output.

    Args:
        command (str): The bash command to be executed.

    Returns:
        Tuple[str, int]: A tuple containing the output of the command (stdout) as a string,
        the error output (stderr) as a string (optional), and the return code as an integer.
    """
    logger.debug(f"running {command=}")
    # This is a acceptable risk
    process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
    output, _ = process.communicate()
    if process.returncode != 0:
        error = f"Failed to run command {command=} return code {process.returncode=}"
        raise RuntimeError(error)

    return output.decode()


def partition_disk(disk: str, swap_size: int, reserve: int = 0) -> None:
    """Partition a disk.

    Args:
        disk (str): The disk to partition.
        swap_size (int): The size of the swap partition in GB.
            minimum value is 1.
        reserve (int, optional): The size of the reserve partition in GB. Defaults to 0.
            minimum value is 0.
    """
    logger.info(f"partitioning {disk=}")
    swap_size = max(swap_size, 1)
    reserve = max(reserve, 0)

    bash_wrapper(f"blkdiscard -f {disk}")

    if reserve > 0:
        msg = f"Creating swap partition on {disk=} with size {swap_size=}GiB and reserve {reserve=}GiB"
        logger.info(msg)

        swap_start = swap_size + reserve
        swap_partition = f"mkpart swap -{swap_start}GiB -{reserve}GiB "
    else:
        logger.info(f"Creating swap partition on {disk=} with size {swap_size=}GiB")
        swap_start = swap_size
        swap_partition = f"mkpart swap -{swap_start}GiB 100% "

    logger.debug(f"{swap_partition=}")

    create_partitions = (
        f"parted --script --align=optimal {disk} -- "
        "mklabel gpt "
        "mkpart EFI 1MiB 4GiB "
        f"mkpart root_pool 4GiB -{swap_start}GiB "
        f"{swap_partition}"
        "set 1 esp on"
    )
    bash_wrapper(create_partitions)

    logger.info(f"{disk=} successfully partitioned")


def create_zfs_pool(pool_disks: Sequence[str], mnt_dir: str) -> None:
    """Create a ZFS pool.

    Args:
        pool_disks (Sequence[str]): A tuple of disks to use for the pool.
        mnt_dir (str): The mount directory.
    """
    if len(pool_disks) <= 0:
        error = "disks must be a tuple of at least length 1"
        raise ValueError(error)

    zpool_create = (
        "zpool create "
        "-o ashift=12 "
        "-o autotrim=on "
        f"-R {mnt_dir} "
        "-O acltype=posixacl "
        "-O canmount=off "
        "-O dnodesize=auto "
        "-O normalization=formD "
        "-O relatime=on "
        "-O xattr=sa "
        "-O mountpoint=legacy "
        "-O compression=zstd "
        "-O atime=off "
        "root_pool "
    )
    if len(pool_disks) == 1:
        zpool_create += pool_disks[0]
    else:
        zpool_create += "mirror "
        zpool_create += " ".join(pool_disks)

    bash_wrapper(zpool_create)
    zpools = bash_wrapper("zpool list -o name")
    if "root_pool" not in zpools.splitlines():
        logger.critical("Failed to create root_pool")
        sys.exit(1)


def create_zfs_datasets() -> None:
    """Create ZFS datasets."""
    bash_wrapper("zfs create -o canmount=noauto -o reservation=10G root_pool/root")
    bash_wrapper("zfs create root_pool/home")
    bash_wrapper("zfs create root_pool/var -o reservation=1G")
    bash_wrapper("zfs create -o compression=zstd-9 -o reservation=10G root_pool/nix")
    datasets = bash_wrapper("zfs list -o name")

    expected_datasets = {
        "root_pool/root",
        "root_pool/home",
        "root_pool/var",
        "root_pool/nix",
    }
    missing_datasets = expected_datasets.difference(datasets.splitlines())
    if missing_datasets:
        logger.critical(f"Failed to create pools {missing_datasets}")
        sys.exit(1)


def get_cpu_manufacturer() -> str:
    """Get the CPU manufacturer."""
    output = bash_wrapper("cat /proc/cpuinfo")

    id_vendor = {"AuthenticAMD": "amd", "GenuineIntel": "intel"}

    for line in output.splitlines():
        if "vendor_id" in line:
            return id_vendor[line.split(": ")[1].strip()]
    error = "Failed to get CPU manufacturer"
    raise RuntimeError(error)


def get_boot_drive_id(disk: str) -> str:
    """Get the boot drive ID."""
    output = bash_wrapper(f"lsblk -o UUID {disk}-part1")
    return output.splitlines()[1]


def create_nix_hardware_file(mnt_dir: str, disks: Sequence[str], *, encrypt: bool) -> None:
    """Create a NixOS hardware file."""
    cpu_manufacturer = get_cpu_manufacturer()

    devices = ""
    if encrypt:
        disk = disks[0]

        devices = (
            f'     luks.devices."luks-root-pool-{disk.split("/")[-1]}-part2"'
            "= {\n"
            f'        device = "{disk}-part2";\n'
            "        bypassWorkqueues = true;\n"
            "        allowDiscards = true;\n"
            "      };\n"
        )

    host_id = format(getrandbits(32), "08x")

    nix_hardware = (
        "{ config, lib, modulesPath, ... }:\n"
        "{\n"
        '  imports = [ (modulesPath + "/installer/scan/not-detected.nix") ];\n\n'
        "  boot = {\n"
        "    initrd = {\n"
        "      availableKernelModules = [ \n"
        '        "ahci"\n'
        '        "ehci_pci"\n'
        '        "nvme"\n'
        '        "sd_mod"\n'
        '        "usb_storage"\n'
        '        "usbhid"\n'
        '        "xhci_pci"\n'
        "      ];\n"
        "      kernelModules = [ ];\n"
        f" {devices}"
        "    };\n"
        f'    kernelModules = [ "kvm-{cpu_manufacturer}" ];\n'
        "    extraModulePackages = [ ];\n"
        "  };\n\n"
        "  fileSystems = {\n"
        '    "/" = lib.mkDefault {\n      device = "root_pool/root";\n      fsType = "zfs";\n    };\n\n'
        '    "/home" = {\n      device = "root_pool/home";\n      fsType = "zfs";\n    };\n\n'
        '    "/var" = {\n      device = "root_pool/var";\n      fsType = "zfs";\n    };\n\n'
        '    "/nix" = {\n      device = "root_pool/nix";\n      fsType = "zfs";\n    };\n\n'
        '    "/boot" = {\n'
        f'      device = "/dev/disk/by-uuid/{get_boot_drive_id(disks[0])}";\n'
        '      fsType = "vfat";\n'
        "      options = [\n"
        '        "fmask=0077"\n'
        '        "dmask=0077"\n'
        "      ];\n"
        "    };\n"
        "  };\n\n"
        "  swapDevices = [ ];\n\n"
        "  networking.useDHCP = lib.mkDefault true;\n\n"
        '  nixpkgs.hostPlatform = lib.mkDefault "x86_64-linux";\n'
        f"  hardware.cpu.{cpu_manufacturer}.updateMicrocode = lib.mkDefault "
        "config.hardware.enableRedistributableFirmware;\n"
        f'  networking.hostId = "{host_id}";\n'
        "}\n"
    )

    Path(f"{mnt_dir}/etc/nixos/hardware-configuration.nix").write_text(nix_hardware)


def install_nixos(mnt_dir: str, disks: Sequence[str], *, encrypt: bool) -> None:
    """Install NixOS."""
    bash_wrapper(f"mount -o X-mount.mkdir -t zfs root_pool/root {mnt_dir}")
    bash_wrapper(f"mount -o X-mount.mkdir -t zfs root_pool/home {mnt_dir}/home")
    bash_wrapper(f"mount -o X-mount.mkdir -t zfs root_pool/var {mnt_dir}/var")
    bash_wrapper(f"mount -o X-mount.mkdir -t zfs root_pool/nix {mnt_dir}/nix")

    for disk in disks:
        bash_wrapper(f"mkfs.vfat -n EFI {disk}-part1")

    # set up mirroring afterwards if more than one disk
    boot_partition = (
        f"mount -t vfat -o fmask=0077,dmask=0077,iocharset=iso8859-1,X-mount.mkdir {disks[0]}-part1 {mnt_dir}/boot"
    )
    bash_wrapper(boot_partition)

    bash_wrapper(f"nixos-generate-config --root {mnt_dir}")

    create_nix_hardware_file(mnt_dir, disks, encrypt=encrypt)

    run(("nixos-install", "--root", mnt_dir), check=True)


def installer(
    disks: set[str],
    swap_size: int,
    reserve: int,
    encrypt_key: str | None,
) -> None:
    """Main."""
    logger.info("Starting installation")

    for disk in disks:
        partition_disk(disk, swap_size, reserve)

        if encrypt_key:
            sleep(1)
            key_input = encrypt_key.encode()
            run(
                ("cryptsetup", "luksFormat", "--type", "luks2", f"{disk}-part2", "-"),
                input=key_input,
                check=True,
            )
            run(
                (
                    "cryptsetup",
                    "luksOpen",
                    f"{disk}-part2",
                    f"luks-root-pool-{disk.split('/')[-1]}-part2",
                    "-",
                ),
                input=key_input,
                check=True,
            )

    mnt_dir = "/tmp/nix_install"  # noqa: S108

    Path(mnt_dir).mkdir(parents=True, exist_ok=True)

    if encrypt_key:
        pool_disks = [f"/dev/mapper/luks-root-pool-{disk.split('/')[-1]}-part2" for disk in disks]
    else:
        pool_disks = [f"{disk}-part2" for disk in disks]

    create_zfs_pool(pool_disks, mnt_dir)

    create_zfs_datasets()

    install_nixos(mnt_dir, disks, encrypt=bool(encrypt_key))

    logger.info("Installation complete")


class Cursor:
    """Track cursor position and constrain movement to screen bounds."""

    def __init__(self) -> None:
        """Initialize cursor position and screen dimensions."""
        self.x_position = 0
        self.y_position = 0
        self.height = 0
        self.width = 0

    def set_height(self, height: int) -> None:
        """Set the maximum screen height."""
        self.height = height

    def set_width(self, width: int) -> None:
        """Set the maximum screen width."""
        self.width = width

    def x_bounce_check(self, cursor: int) -> int:
        """Clamp an x position to the screen width."""
        cursor = max(0, cursor)
        return min(self.width - 1, cursor)

    def y_bounce_check(self, cursor: int) -> int:
        """Clamp a y position to the screen height."""
        cursor = max(0, cursor)
        return min(self.height - 1, cursor)

    def set_x(self, x: int) -> None:
        """Set the cursor x position."""
        self.x_position = self.x_bounce_check(x)

    def set_y(self, y: int) -> None:
        """Set the cursor y position."""
        self.y_position = self.y_bounce_check(y)

    def get_x(self) -> int:
        """Get the cursor x position."""
        return self.x_position

    def get_y(self) -> int:
        """Get the cursor y position."""
        return self.y_position

    def move_up(self) -> None:
        """Move the cursor up one row."""
        self.set_y(self.y_position - 1)

    def move_down(self) -> None:
        """Move the cursor down one row."""
        self.set_y(self.y_position + 1)

    def move_left(self) -> None:
        """Move the cursor left one column."""
        self.set_x(self.x_position - 1)

    def move_right(self) -> None:
        """Move the cursor right one column."""
        self.set_x(self.x_position + 1)

    def navigation(self, key: int) -> None:
        """Move the cursor for a curses navigation key."""
        action = {
            curses.KEY_DOWN: self.move_down,
            curses.KEY_UP: self.move_up,
            curses.KEY_RIGHT: self.move_right,
            curses.KEY_LEFT: self.move_left,
        }

        action.get(key, lambda: None)()


class State:
    """State class to store the state of the program."""

    def __init__(self) -> None:
        """Initialize installer menu state."""
        self.key = 0
        self.cursor = Cursor()

        self.swap_size = 0
        self.show_swap_input = False

        self.reserve_size = 0
        self.show_reserve_input = False

        self.selected_device_ids = set()

    def get_selected_devices(self) -> tuple[str]:
        """Get selected devices."""
        return tuple(self.selected_device_ids)


def get_device(raw_device: str) -> dict[str, str]:
    """Parse an lsblk key-value device row."""
    raw_device_components = raw_device.split(" ")
    return {thing.split("=")[0].lower(): thing.split("=")[1].strip('"') for thing in raw_device_components}


def get_devices() -> list[dict[str, str]]:
    """Get a list of devices."""
    # --bytes
    raw_devices = bash_wrapper("lsblk --paths --pairs").splitlines()
    return [get_device(raw_device) for raw_device in raw_devices]


def get_device_id_mapping() -> dict[str, set[str]]:
    """Get a list of device ids.

    Returns:
        list[str]: the list of device ids
    """
    device_ids = bash_wrapper("find /dev/disk/by-id -type l").splitlines()

    device_id_mapping: dict[str, set[str]] = defaultdict(set)

    for device_id in device_ids:
        device = bash_wrapper(f"readlink -f {device_id}").strip()
        device_id_mapping[device].add(device_id)

    return device_id_mapping


def calculate_device_menu_padding(devices: list[dict[str, str]], column: str, padding: int = 0) -> int:
    """Calculate the width needed for a device menu column."""
    return max(len(device[column]) for device in devices) + padding


def draw_device_ids(
    state: State,
    row_number: int,
    menu_start_x: int,
    std_screen: curses.window,
    menu_width: list[int],
    device_ids: set[str],
) -> tuple[State, int]:
    """Draw selectable device IDs for a device row."""
    for device_id in sorted(device_ids):
        row_number = row_number + 1
        if row_number == state.cursor.get_y() and state.cursor.get_x() in menu_width:
            std_screen.attron(curses.A_BOLD)
            if state.key == ord(" "):
                if device_id not in state.selected_device_ids:
                    state.selected_device_ids.add(device_id)
                else:
                    state.selected_device_ids.remove(device_id)

        if device_id in state.selected_device_ids:
            std_screen.attron(curses.color_pair(7))

        std_screen.addstr(row_number, menu_start_x, f"  {device_id}")

        std_screen.attroff(curses.color_pair(7))
        std_screen.attroff(curses.A_BOLD)

    return state, row_number


def draw_device_menu(
    std_screen: curses.window,
    devices: list[dict[str, str]],
    device_id_mapping: dict[str, set[str]],
    state: State,
    menu_start_y: int = 0,
    menu_start_x: int = 0,
) -> tuple[State, int]:
    """Draw the device menu and handle user input.

    Args:
        std_screen (curses.window): the curses window to draw on
        devices (list[dict[str, str]]): the list of devices to draw
        device_id_mapping (dict[str, set[str]]): the list of device ids to draw
        state (State): the state object to update
        menu_start_y (int, optional): the y position to start drawing the menu. Defaults to 0.
        menu_start_x (int, optional): the x position to start drawing the menu. Defaults to 0.

    Returns:
        State: the updated state object
    """
    padding = 2

    name_padding = calculate_device_menu_padding(devices, "name", padding)
    size_padding = calculate_device_menu_padding(devices, "size", padding)
    type_padding = calculate_device_menu_padding(devices, "type", padding)
    mountpoints_padding = calculate_device_menu_padding(devices, "mountpoints", padding)

    device_header = (
        f"{'Name':{name_padding}}{'Size':{size_padding}}{'Type':{type_padding}}{'Mountpoints':{mountpoints_padding}}"
    )

    menu_width = range(menu_start_x, len(device_header) + menu_start_x)

    std_screen.addstr(menu_start_y, menu_start_x, device_header, curses.color_pair(5))
    devises_list_start = menu_start_y + 1

    row_number = devises_list_start

    for device in devices:
        row_number = row_number + 1
        device_name = device["name"]
        device_row = (
            f"{device_name:{name_padding}}"
            f"{device['size']:{size_padding}}"
            f"{device['type']:{type_padding}}"
            f"{device['mountpoints']:{mountpoints_padding}}"
        )
        std_screen.addstr(row_number, menu_start_x, device_row)

        state, row_number = draw_device_ids(
            state=state,
            row_number=row_number,
            menu_start_x=menu_start_x,
            std_screen=std_screen,
            menu_width=menu_width,
            device_ids=device_id_mapping[device_name],
        )

    return state, row_number


def debug_menu(std_screen: curses.window, key: int) -> None:
    """Draw debug information for the current curses screen."""
    height, width = std_screen.getmaxyx()
    width_height = f"Width: {width}, Height: {height}"
    std_screen.addstr(height - 4, 0, width_height, curses.color_pair(5))

    key_pressed = f"Last key pressed: {key}"[: width - 1]
    if key == 0:
        key_pressed = "No key press detected..."[: width - 1]
    std_screen.addstr(height - 3, 0, key_pressed)

    for i in range(8):
        std_screen.addstr(height - 2, i * 3, f"{i}██", curses.color_pair(i))


def status_bar(
    std_screen: curses.window,
    cursor: Cursor,
    width: int,
    height: int,
) -> None:
    """Draw the footer status bar."""
    std_screen.attron(curses.A_REVERSE)
    std_screen.attron(curses.color_pair(3))

    status_bar = f"Press 'q' to exit | STATUS BAR | Pos: {cursor.get_x()}, {cursor.get_y()}"
    std_screen.addstr(height - 1, 0, status_bar)
    std_screen.addstr(height - 1, len(status_bar), " " * (width - len(status_bar) - 1))

    std_screen.attroff(curses.color_pair(3))
    std_screen.attroff(curses.A_REVERSE)


def set_color() -> None:
    """Initialize curses color pairs."""
    curses.start_color()
    curses.use_default_colors()
    for i in range(curses.COLORS):
        curses.init_pair(i + 1, i, -1)


def get_text_input(std_screen: curses.window, prompt: str, y: int, x: int) -> str:
    """Read text input from a curses screen."""
    curses.echo()
    std_screen.addstr(y, x, prompt)
    input_str = ""
    while True:
        key = std_screen.getch()
        if key == ord("\n"):
            break
        if key == ESCAPE_KEY:
            input_str = ""
            break
        if key in (curses.KEY_BACKSPACE, ord("\b"), 127):
            input_str = input_str[:-1]
            std_screen.addstr(y, x + len(prompt), input_str + " ")
        else:
            input_str += chr(key)
        std_screen.refresh()
    curses.noecho()
    return input_str


def swap_size_input(
    std_screen: curses.window,
    state: State,
    swap_offset: int,
) -> State:
    """Handle swap size input."""
    swap_size_text = "Swap size (GB): "
    std_screen.addstr(swap_offset, 0, f"{swap_size_text}{state.swap_size}")
    if state.key == ord("\n") and state.cursor.get_y() == swap_offset:
        state.show_swap_input = True

    if state.show_swap_input:
        swap_size_str = get_text_input(std_screen, swap_size_text, swap_offset, 0)
        try:
            state.swap_size = int(swap_size_str)
            state.show_swap_input = False
        except ValueError:
            std_screen.addstr(swap_offset, 0, "Invalid input. Press any key to continue.")
            std_screen.getch()
            state.show_swap_input = False

    return state


def reserve_size_input(
    std_screen: curses.window,
    state: State,
    reserve_offset: int,
) -> State:
    """Handle reserve size input."""
    reserve_size_text = "reserve size (GB): "
    std_screen.addstr(reserve_offset, 0, f"{reserve_size_text}{state.reserve_size}")
    if state.key == ord("\n") and state.cursor.get_y() == reserve_offset:
        state.show_reserve_input = True

    if state.show_reserve_input:
        reserve_size_str = get_text_input(std_screen, reserve_size_text, reserve_offset, 0)
        try:
            state.reserve_size = int(reserve_size_str)
            state.show_reserve_input = False
        except ValueError:
            std_screen.addstr(reserve_offset, 0, "Invalid input. Press any key to continue.")
            std_screen.getch()
            state.show_reserve_input = False

    return state


def draw_menu(std_screen: curses.window) -> State:
    """Draw the menu and handle user input.

    Args:
        std_screen (curses.window): the curses window to draw on

    Returns:
        State: the state object
    """
    # Clear and refresh the screen for a blank canvas
    std_screen.clear()
    std_screen.refresh()

    set_color()

    state = State()

    devices = get_devices()

    device_id_mapping = get_device_id_mapping()

    # Loop where k is the last character pressed
    while state.key != ord("q"):
        std_screen.clear()
        height, width = std_screen.getmaxyx()

        state.cursor.set_height(height)
        state.cursor.set_width(width)

        state.cursor.navigation(state.key)

        state, device_menu_size = draw_device_menu(
            std_screen=std_screen,
            state=state,
            devices=devices,
            device_id_mapping=device_id_mapping,
        )

        swap_offset = device_menu_size + 2

        swap_size_input(
            std_screen=std_screen,
            state=state,
            swap_offset=swap_offset,
        )
        reserve_size_input(
            std_screen=std_screen,
            state=state,
            reserve_offset=swap_offset + 1,
        )

        status_bar(std_screen, state.cursor, width, height)

        debug_menu(std_screen, state.key)

        std_screen.move(state.cursor.get_y(), state.cursor.get_x())

        std_screen.refresh()

        state.key = std_screen.getch()

    return state


def main() -> None:
    """Run the installer menu and start installation."""
    configure_logger("DEBUG")

    state = curses.wrapper(draw_menu)

    encrypt_key = getenv("ENCRYPT_KEY")

    logger.info("installing_nixos")
    logger.info(f"disks: {state.selected_device_ids}")
    logger.info(f"swap_size: {state.swap_size}")
    logger.info(f"reserve: {state.reserve_size}")
    logger.info(f"encrypted: {bool(encrypt_key)}")

    sleep(3)

    installer(
        disks=state.get_selected_devices(),
        swap_size=state.swap_size,
        reserve=state.reserve_size,
        encrypt_key=encrypt_key,
    )


if __name__ == "__main__":
    main()