local-ci-cd-system/src/ci_orchestrator/winimage/uup.py

"""Acquire a Windows Server 2025 install ISO from UUP dump (Linux host).

UUP dump is deterministic on its *input*: a build is pinned by build number and
the component files come MS-signed from the official Windows Update CDN. The
final ISO is assembled locally (its hash is not canonical), so we cache the
result per build number under ``<CI_ROOT>/iso/ws2025-<build>.iso``.

Flow (mirrors what uupdump.net's generated package does):

1. ``resolve_build`` — query the uupdump API for Windows Server 2025 builds and
   pick the newest (or match a pinned build string).
2. ``build_iso`` — download the per-set conversion package, run its
   ``uup_download_linux.sh`` (aria2 download + wimlib conversion), then move the
   produced ISO into the cache.

The downloader requires ``aria2c``, ``cabextract``, ``wimlib-imagex``,
``chntpw`` and ``genisoimage``/``mkisofs`` on PATH (installed by
``setup-host-linux.sh``).
"""

from __future__ import annotations

import json
import shutil
import subprocess
import urllib.parse
import urllib.request
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import click

API_BASE = "https://api.uupdump.net"
GET_URL = "https://uupdump.net/get.php"
_USER_AGENT = "ci-orchestrator/1.0 (+deploy-win)"

#: Tools the UUP → ISO conversion needs (uup_download_linux.sh checks these).
REQUIRED_TOOLS = ("aria2c", "cabextract", "wimlib-imagex", "chntpw")

#: Title prefix that identifies a Windows Server 2025 set (excludes "version 23H2").
_WS2025_PREFIX = "Windows Server 2025"


@dataclass(frozen=True)
class UupBuild:
    """A UUP dump build set."""

    uuid: str
    build: str  # e.g. "26100.32860"
    title: str
    arch: str


def _http_json(url: str, *, timeout: float = 60.0) -> dict[str, Any]:
    """GET ``url`` and parse the JSON body."""
    req = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT, "Accept": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        raw = resp.read().decode("utf-8")
    data: Any = json.loads(raw)
    if not isinstance(data, dict):
        raise click.ClickException(f"uupdump API returned unexpected JSON from {url}")
    return data


def _build_sort_key(build: str) -> tuple[int, ...]:
    """Numeric sort key for a build string like ``26100.32860``."""
    parts: list[int] = []
    for chunk in build.split("."):
        try:
            parts.append(int(chunk))
        except ValueError:
            parts.append(0)
    return tuple(parts)


def list_ws2025_builds(*, arch: str = "amd64", timeout: float = 60.0) -> list[UupBuild]:
    """Return all Windows Server 2025 builds for ``arch``, newest first."""
    url = f"{API_BASE}/listid.php?search={urllib.parse.quote(_WS2025_PREFIX)}"
    data = _http_json(url, timeout=timeout)
    builds_raw = data.get("response", {}).get("builds", {})
    if not isinstance(builds_raw, dict):
        raise click.ClickException("uupdump API: missing 'builds' in listid response")

    out: list[UupBuild] = []
    for entry in builds_raw.values():
        if not isinstance(entry, dict):
            continue
        title = str(entry.get("title", ""))
        if not title.startswith(_WS2025_PREFIX):
            continue
        if str(entry.get("arch", "")) != arch:
            continue
        out.append(
            UupBuild(
                uuid=str(entry["uuid"]),
                build=str(entry["build"]),
                title=title,
                arch=arch,
            )
        )
    out.sort(key=lambda b: _build_sort_key(b.build), reverse=True)
    return out


def resolve_build(
    *, build: str | None = None, arch: str = "amd64", timeout: float = 60.0
) -> UupBuild:
    """Resolve the newest WS2025 build, or the one matching ``build`` if pinned."""
    builds = list_ws2025_builds(arch=arch, timeout=timeout)
    if not builds:
        raise click.ClickException(
            f"No Windows Server 2025 ({arch}) builds found on uupdump."
        )
    if build is None:
        return builds[0]
    for b in builds:
        if b.build == build:
            return b
    available = ", ".join(b.build for b in builds)
    raise click.ClickException(
        f"WS2025 build {build!r} not found on uupdump. Available: {available}"
    )


def check_tools() -> list[str]:
    """Return the list of missing conversion tools (empty when all present)."""
    missing = [t for t in REQUIRED_TOOLS if shutil.which(t) is None]
    if shutil.which("genisoimage") is None and shutil.which("mkisofs") is None:
        missing.append("genisoimage|mkisofs")
    return missing


def cache_path(dest_dir: Path, build: UupBuild) -> Path:
    """Path of the cached ISO for ``build`` under ``dest_dir``."""
    return dest_dir / f"ws2025-{build.build}.iso"


def _download_package(build: UupBuild, lang: str, edition: str, dest: Path, *, timeout: float) -> None:
    """Download the uupdump conversion package zip for ``build``."""
    url = (
        f"{GET_URL}?id={build.uuid}"
        f"&pack={urllib.parse.quote(lang)}"
        f"&edition={urllib.parse.quote(edition)}"
        f"&autodl=2"
    )
    req = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT})
    with urllib.request.urlopen(req, timeout=timeout) as resp, dest.open("wb") as fh:
        shutil.copyfileobj(resp, fh)
    if dest.stat().st_size < 1024:
        raise click.ClickException(f"uupdump package download too small: {dest}")


def build_iso(
    build: UupBuild,
    *,
    dest_dir: Path,
    lang: str = "en-us",
    edition: str = "SERVERSTANDARD",
    work_dir: Path | None = None,
    keep_work: bool = False,
) -> Path:
    """Build (or reuse a cached) WS2025 ISO for ``build`` and return its path.

    Downloads the uupdump conversion package, runs its ``uup_download_linux.sh``
    (aria2 download of the MS-signed UUP set + wimlib conversion to ISO), then
    moves the produced ISO to ``<dest_dir>/ws2025-<build>.iso``.
    """
    dest_dir.mkdir(parents=True, exist_ok=True)
    cached = cache_path(dest_dir, build)
    if cached.is_file():
        click.echo(f"  [uup] cached ISO present, reusing: {cached}")
        return cached

    missing = check_tools()
    if missing:
        raise click.ClickException(
            "Missing UUP conversion tools: "
            + ", ".join(missing)
            + ". Install them (setup-host-linux.sh installs aria2 wimtools "
            "cabextract chntpw genisoimage)."
        )

    work = work_dir or (dest_dir / f".uup-{build.build}")
    if work.exists():
        shutil.rmtree(work)
    work.mkdir(parents=True)

    pkg = work / "uup-package.zip"
    click.echo(f"  [uup] downloading conversion package for build {build.build} ({build.uuid})")
    _download_package(build, lang, edition, pkg, timeout=120.0)
    with zipfile.ZipFile(pkg) as zf:
        zf.extractall(work)

    script = work / "uup_download_linux.sh"
    if not script.is_file():
        raise click.ClickException(f"uup_download_linux.sh missing from package: {work}")
    script.chmod(0o755)

    click.echo("  [uup] running converter (downloads ~5 GB + builds ISO; this is slow) ...")
    proc = subprocess.run(
        ["bash", str(script)],
        cwd=str(work),
        check=False,
    )
    if proc.returncode != 0:
        raise click.ClickException(
            f"UUP conversion failed (exit {proc.returncode}). See {work}/aria2_download.log."
        )

    produced = _find_produced_iso(work)
    if produced is None:
        raise click.ClickException(f"UUP conversion produced no ISO under {work}")

    shutil.move(str(produced), str(cached))
    click.echo(f"  [uup] ISO ready: {cached}")
    if not keep_work:
        shutil.rmtree(work, ignore_errors=True)
    return cached


def _find_produced_iso(work: Path) -> Path | None:
    """Find the ISO the converter produced (largest *.iso, excluding the package)."""
    candidates = [
        p
        for p in work.rglob("*.iso")
        if p.name != "uup-package.zip" and p.is_file()
    ]
    candidates += [p for p in work.rglob("*.ISO") if p.is_file()]
    if not candidates:
        return None
    return max(candidates, key=lambda p: p.stat().st_size)