sovereign-orchestrator/app/iso_builder.py

"""
Sovereign Orchestrator - ISO Builder

Generates answers.toml files and invokes `proxmox-auto-install-assistant`
to produce custom auto-installer ISOs.  Build state is tracked in-memory
and exposed through the API.
"""

from __future__ import annotations

import asyncio
import logging
import os
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import toml

from app.config import settings
from app.models import BuildState, BuildStatus, ISOConfig

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# In-memory build registry
# ---------------------------------------------------------------------------

_builds: dict[str, BuildStatus] = {}


def get_all_builds() -> list[BuildStatus]:
    """Return all builds, newest first."""
    return sorted(
        _builds.values(), key=lambda b: b.created_at, reverse=True
    )


def get_build(build_id: str) -> Optional[BuildStatus]:
    """Look up a single build by ID."""
    return _builds.get(build_id)


def get_latest_build() -> Optional[BuildStatus]:
    """Return the most recent build, or None."""
    builds = get_all_builds()
    return builds[0] if builds else None


def get_active_builds() -> list[BuildStatus]:
    """Return builds that are still in progress."""
    return [
        b for b in _builds.values()
        if b.state in (BuildState.PENDING, BuildState.BUILDING, BuildState.UPLOADING)
    ]


# ---------------------------------------------------------------------------
# answers.toml generation
# ---------------------------------------------------------------------------

def generate_answers_toml(config: ISOConfig) -> str:
    """Render an ``answers.toml`` string from the given configuration.

    The output format matches the Proxmox auto-installer specification
    with [global], [network], and [disk-setup] sections.
    """
    doc: dict = {
        "global": {
            "keyboard": config.keyboard,
            "country": config.country,
            "timezone": config.timezone,
            "mailto": config.mailto,
            "root-password": config.root_password,
            "fqdn": config.fqdn,
            "reboot-on-error": config.reboot_on_error,
        },
        "network": {
            "source": config.network_source.value,
        },
        "disk-setup": {
            "filesystem": config.filesystem.value,
            "disk-list": config.disk_list,
        },
    }

    if config.root_ssh_keys:
        doc["global"]["root-ssh-keys"] = config.root_ssh_keys

    return toml.dumps(doc)


# ---------------------------------------------------------------------------
# Base ISO discovery
# ---------------------------------------------------------------------------

_BASE_ISO_SEARCH_PATHS = [
    Path(settings.iso_storage_path),
    Path("/var/lib/vz/template/iso"),
    Path("/usr/share/proxmox-installer"),
    Path.home() / "isos",
]


def find_base_iso() -> Optional[Path]:
    """Locate the first Proxmox VE ISO in well-known directories.

    Searches for files matching ``proxmox-ve_*.iso``.
    """
    for search_dir in _BASE_ISO_SEARCH_PATHS:
        if not search_dir.is_dir():
            continue
        candidates = sorted(
            search_dir.glob("proxmox-ve_*.iso"), reverse=True
        )
        if candidates:
            logger.info("Found base ISO: %s", candidates[0])
            return candidates[0]
    logger.warning("No base Proxmox VE ISO found in search paths")
    return None


# ---------------------------------------------------------------------------
# ISO build pipeline
# ---------------------------------------------------------------------------

def _ensure_build_dir() -> Path:
    """Create and return the build output directory."""
    build_dir = Path(settings.build_output_path)
    build_dir.mkdir(parents=True, exist_ok=True)
    return build_dir


async def build_iso(config: ISOConfig) -> BuildStatus:
    """Kick off an asynchronous ISO build and return the BuildStatus.

    The actual subprocess work runs in the background so the API can
    return immediately.
    """
    build = BuildStatus(iso_config=config)
    _builds[build.id] = build
    build.log(f"Build {build.id} created")

    # Launch the background task
    asyncio.create_task(_run_build(build))

    return build


async def _run_build(build: BuildStatus) -> None:
    """Execute the full ISO build pipeline."""
    try:
        build.state = BuildState.BUILDING
        build.log("Build started")

        build_dir = _ensure_build_dir()
        work_dir = build_dir / build.id
        work_dir.mkdir(parents=True, exist_ok=True)

        # 1. Write answers.toml
        answers_path = work_dir / "answers.toml"
        answers_content = generate_answers_toml(build.iso_config)
        answers_path.write_text(answers_content)
        build.log(f"answers.toml written to {answers_path}")
        logger.debug("answers.toml content:\n%s", answers_content)

        # 2. Locate base ISO
        base_iso = find_base_iso()
        if base_iso is None:
            raise FileNotFoundError(
                "No base Proxmox VE ISO found. Place a proxmox-ve_*.iso in "
                f"{settings.iso_storage_path}"
            )
        build.log(f"Using base ISO: {base_iso}")

        # 3. Determine output filename
        fqdn_slug = build.iso_config.fqdn.replace(".", "-")
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
        output_filename = f"pve-autoinstall-{fqdn_slug}-{timestamp}.iso"
        output_path = work_dir / output_filename

        # 4. Copy the base ISO to the output path first (prepare-iso works
        #    on a copy so we don't mutate the original)
        build.log(f"Copying base ISO to {output_path}")
        await asyncio.to_thread(shutil.copy2, str(base_iso), str(output_path))
        build.log("Base ISO copied")

        # 5. Run proxmox-auto-install-assistant prepare-iso
        cmd = [
            "proxmox-auto-install-assistant",
            "prepare-iso",
            str(output_path),
            "--fetch-from", "iso",
            "--answer-file", str(answers_path),
        ]
        build.log(f"Running: {' '.join(cmd)}")

        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            cwd=str(work_dir),
        )

        # Stream output into build logs
        assert proc.stdout is not None
        while True:
            line = await proc.stdout.readline()
            if not line:
                break
            decoded = line.decode("utf-8", errors="replace").rstrip()
            build.log(decoded)
            logger.debug("[build %s] %s", build.id, decoded)

        return_code = await proc.wait()

        if return_code != 0:
            raise RuntimeError(
                f"proxmox-auto-install-assistant exited with code {return_code}"
            )

        build.iso_filename = output_filename
        build.state = BuildState.COMPLETED
        build.log(f"Build completed: {output_filename}")
        logger.info("Build %s completed: %s", build.id, output_path)

    except Exception as exc:
        build.state = BuildState.FAILED
        build.error = str(exc)
        build.log(f"Build failed: {exc}")
        logger.exception("Build %s failed", build.id)


# ---------------------------------------------------------------------------
# Utility: get the output ISO path for a build
# ---------------------------------------------------------------------------

def get_build_iso_path(build_id: str) -> Optional[Path]:
    """Return the absolute path to the output ISO for a completed build."""
    build = _builds.get(build_id)
    if not build or not build.iso_filename:
        return None
    return Path(settings.build_output_path) / build.id / build.iso_filename