dotfiles/python/installer/__main__.py

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

from __future__ import annotations

import curses
import logging
import sys
from argparse import ArgumentParser
from os import environ, getenv
from pathlib import Path
from random import getrandbits
from subprocess import run
from time import sleep
from typing import TYPE_CHECKING

from python.installer.tui import draw_menu
from python.logging_config import configure_logger
from python.process import require_commands, run_output

if TYPE_CHECKING:
    from collections.abc import Sequence

logger = logging.getLogger(__name__)

REQUIRED_COMMANDS = (
    "blkdiscard",
    "cryptsetup",
    "find",
    "lsblk",
    "mkfs.vfat",
    "mount",
    "nixos-generate-config",
    "nixos-install",
    "parted",
    "readlink",
    "zfs",
    "zpool",
)


def configure_terminal() -> None:
    """Force a terminal type every live system has a terminfo entry for.

    Terminals such as kitty advertise TERM values (xterm-kitty) that live
    systems have no terminfo entry for, so setupterm fails before the TUI
    can start.
    """
    environ["TERM"] = "xterm-256color"

    if getenv("TERMINFO") or getenv("TERMINFO_DIRS"):
        return
    terminfo_fallback_directories = (
        Path("/run/current-system/sw/share/terminfo"),
        Path("/usr/share/terminfo"),
        Path("/etc/terminfo"),
        Path("/lib/terminfo"),
    )

    existing_directories = [str(directory) for directory in terminfo_fallback_directories if directory.is_dir()]
    if existing_directories:
        environ["TERMINFO_DIRS"] = ":".join(existing_directories)


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)

    run_output(("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 = (
        "parted",
        "--script",
        "--align=optimal",
        disk,
        "--",
        "mklabel",
        "gpt",
        "mkpart",
        "EFI",
        "1MiB",
        "4GiB",
        "mkpart",
        "root_pool",
        "4GiB",
        f"-{swap_start}GiB",
        *swap_partition.split(),
        "set",
        "1",
        "esp",
        "on",
    )
    run_output(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",
        "-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.append(pool_disks[0])
    else:
        zpool_create.append("mirror")
        zpool_create.extend(pool_disks)

    run_output(zpool_create)
    zpools = run_output(("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."""
    run_output(("zfs", "create", "-o", "canmount=noauto", "-o", "reservation=10G", "root_pool/root"))
    run_output(("zfs", "create", "root_pool/home"))
    run_output(("zfs", "create", "-o", "reservation=1G", "root_pool/var"))
    run_output(("zfs", "create", "-o", "compression=zstd-9", "-o", "reservation=10G", "root_pool/nix"))
    datasets = run_output(("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 = Path("/proc/cpuinfo").read_text()

    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 = run_output(("lsblk", "-o", "UUID", f"{disk}-part1"))
    return output.splitlines()[1]


def create_nix_hardware_file(mnt_dir: str, disks: Sequence[str], encrypt: str | None) -> 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: str | None) -> None:
    """Install NixOS."""
    run_output(("mount", "-o", "X-mount.mkdir", "-t", "zfs", "root_pool/root", mnt_dir))
    run_output(("mount", "-o", "X-mount.mkdir", "-t", "zfs", "root_pool/home", f"{mnt_dir}/home"))
    run_output(("mount", "-o", "X-mount.mkdir", "-t", "zfs", "root_pool/var", f"{mnt_dir}/var"))
    run_output(("mount", "-o", "X-mount.mkdir", "-t", "zfs", "root_pool/nix", f"{mnt_dir}/nix"))

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

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

    run_output(("nixos-generate-config", "--root", mnt_dir))

    create_nix_hardware_file(mnt_dir, disks, encrypt)

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


def installer(
    disks: Sequence[str],
    swap_size: int,
    reserve: int,
    encrypt_key: str | None,
) -> None:
    """Main."""
    logger.info("Starting installation")
    require_commands(REQUIRED_COMMANDS)
    disks = tuple(sorted(disks))

    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,
            )

    # nixos-install rejects mount points under world-writable paths like /tmp
    mnt_dir = "/mnt"

    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_key)

    logger.info("Installation complete")


def main(argv: Sequence[str] | None = None) -> None:
    """Main."""
    parser = ArgumentParser(description="Install this NixOS configuration onto a ZFS root pool.")
    parser.add_argument(
        "--check",
        action="store_true",
        help="validate that the live environment has the external installer commands and exit",
    )
    parser.add_argument("--log-level", default=getenv("LOG_LEVEL", "DEBUG"), help="Python log level")
    args = parser.parse_args(argv)

    configure_logger(args.log_level)

    if args.check:
        require_commands(REQUIRED_COMMANDS)
        logger.info("installer runtime dependencies are available")
        return

    configure_terminal()
    state = curses.wrapper(draw_menu)

    encrypt_key = getenv("ENCRYPT_KEY")

    if not encrypt_key:
        encrypt_key = state.encryption_password

    if not state.selected_device_ids:
        logger.error("No disks selected; exiting without installing")
        sys.exit(1)

    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()