dotfiles/python/installer/build.py

"""Build the one-file installer binary."""

from __future__ import annotations

import logging
import os
import shutil
import stat
import subprocess
from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Sequence

logger = logging.getLogger(__name__)

BINARY_NAME = "nixos-installer"
DEFAULT_INTERPRETER = "/lib64/ld-linux-x86-64.so.2"
INSTALLER_SOURCE_FILES = (
    Path("python/__init__.py"),
    Path("python/logging_config.py"),
    Path("python/process.py"),
    Path("python/installer/__init__.py"),
    Path("python/installer/__main__.py"),
    Path("python/installer/tui.py"),
)


class InstallerBuildError(RuntimeError):
    """Raised when the installer binary cannot be built."""


class MissingSourceFileError(InstallerBuildError):
    """Raised when a required source file is missing."""

    def __init__(self, path: Path) -> None:
        """Store the missing path."""
        super().__init__(f"Required installer source file is missing: {path}")
        self.path = path


class MissingToolError(InstallerBuildError):
    """Raised when a required build tool is missing."""

    def __init__(self, tool: str) -> None:
        """Store the missing tool name."""
        super().__init__(f"Required build tool is missing from PATH: {tool}")
        self.tool = tool


def repo_root() -> Path:
    """Return the repository root for direct script usage."""
    return Path(__file__).resolve().parents[2]


def require_tool(tool: str) -> str:
    """Return the path to a tool or raise."""
    tool_path = shutil.which(tool)
    if tool_path is None:
        raise MissingToolError(tool)
    return tool_path


def copy_installer_source(source_root: Path, destination: Path) -> None:
    """Copy only the installer files into a minimal staging tree."""
    if destination.exists():
        shutil.rmtree(destination)
    destination.mkdir(parents=True)

    for relative_path in INSTALLER_SOURCE_FILES:
        source = source_root / relative_path
        if not source.is_file():
            raise MissingSourceFileError(source)

        target = destination / relative_path
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(source, target)


def run_command(command: Sequence[str], *, env: dict[str, str] | None = None) -> None:
    """Run a build command."""
    logger.info("running command=%s", command)
    subprocess.run(command, check=True, env=env)


def pyinstaller_environment(staged_source: Path, build_root: Path) -> dict[str, str]:
    """Return environment variables for PyInstaller."""
    env = os.environ.copy()
    home = build_root / "home"
    home.mkdir(parents=True, exist_ok=True)

    env["HOME"] = str(home)
    if existing_pythonpath := env.get("PYTHONPATH"):
        env["PYTHONPATH"] = f"{staged_source}{os.pathsep}{existing_pythonpath}"
    else:
        env["PYTHONPATH"] = str(staged_source)

    return env


def build_installer(
    *,
    source_root: Path,
    build_root: Path,
    output: Path,
    interpreter: str,
    patch_elf: bool,
) -> Path:
    """Build the one-file installer binary."""
    pyinstaller = require_tool("pyinstaller")
    if patch_elf:
        patchelf = require_tool("patchelf")

    source_root = source_root.resolve()
    build_root = build_root.resolve()
    output = output.resolve()

    staged_source = build_root / "source"
    dist_dir = build_root / "dist"
    work_dir = build_root / "work"
    spec_dir = build_root / "spec"

    copy_installer_source(source_root, staged_source)
    output.parent.mkdir(parents=True, exist_ok=True)

    run_command(
        (
            pyinstaller,
            "--clean",
            "--onefile",
            "--name",
            BINARY_NAME,
            "--paths",
            str(staged_source),
            "--distpath",
            str(dist_dir),
            "--workpath",
            str(work_dir),
            "--specpath",
            str(spec_dir),
            str(staged_source / "python/installer/__main__.py"),
        ),
        env=pyinstaller_environment(staged_source, build_root),
    )

    built_binary = dist_dir / BINARY_NAME
    shutil.copy2(built_binary, output)
    output.chmod(output.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)

    if patch_elf:
        run_command((patchelf, "--set-interpreter", interpreter, "--remove-rpath", str(output)))

    return output


def parse_args(argv: Sequence[str] | None = None) -> Namespace:
    """Parse command-line arguments."""
    parser = ArgumentParser(description="Build the one-file NixOS installer binary.")
    parser.add_argument("--source-root", type=Path, default=repo_root(), help="repo or staged source root")
    parser.add_argument("--build-root", type=Path, default=Path("build/nixos-installer"), help="temporary build root")
    parser.add_argument("--output", type=Path, default=Path("dist/nixos-installer"), help="output binary path")
    parser.add_argument("--interpreter", default=DEFAULT_INTERPRETER, help="ELF interpreter path for the USB binary")
    parser.add_argument("--skip-patchelf", action="store_true", help="do not patch the final ELF binary")
    parser.add_argument("--log-level", default="INFO", help="Python log level")
    return parser.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> None:
    """Build the installer binary from the command line."""
    args = parse_args(argv)
    logging.basicConfig(level=args.log_level, format="%(levelname)s %(message)s")
    build_installer(
        source_root=args.source_root,
        build_root=args.build_root,
        output=args.output,
        interpreter=args.interpreter,
        patch_elf=not args.skip_patchelf,
    )


if __name__ == "__main__":
    main()