Refactor public API around environments

src/pyro_mcp/vm_environments.py

@@ -0,0 +1,615 @@
+"""Official environment catalog and local cache management."""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import tarfile
+import tempfile
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from pyro_mcp.runtime import DEFAULT_PLATFORM, RuntimePaths
+
+DEFAULT_ENVIRONMENT_VERSION = "1.0.0"
+DEFAULT_CATALOG_VERSION = "1.0.0"
+OCI_MANIFEST_ACCEPT = ", ".join(
+    (
+        "application/vnd.oci.image.index.v1+json",
+        "application/vnd.oci.image.manifest.v1+json",
+        "application/vnd.docker.distribution.manifest.list.v2+json",
+        "application/vnd.docker.distribution.manifest.v2+json",
+    )
+)
+
+
+@dataclass(frozen=True)
+class VmEnvironment:
+    """Catalog entry describing a curated Linux environment."""
+
+    name: str
+    version: str
+    description: str
+    default_packages: tuple[str, ...]
+    distribution: str
+    distribution_version: str
+    source_profile: str
+    platform: str = DEFAULT_PLATFORM
+    source_url: str | None = None
+    oci_registry: str | None = None
+    oci_repository: str | None = None
+    oci_reference: str | None = None
+    source_digest: str | None = None
+    compatibility: str = ">=1.0.0,<2.0.0"
+
+
+@dataclass(frozen=True)
+class InstalledEnvironment:
+    """Resolved environment artifact locations."""
+
+    name: str
+    version: str
+    install_dir: Path
+    kernel_image: Path
+    rootfs_image: Path
+    source: str
+    source_digest: str | None
+    installed: bool
+
+
+CATALOG: dict[str, VmEnvironment] = {
+    "debian:12": VmEnvironment(
+        name="debian:12",
+        version=DEFAULT_ENVIRONMENT_VERSION,
+        description="Debian 12 environment with Git preinstalled for common agent workflows.",
+        default_packages=("bash", "coreutils", "git"),
+        distribution="debian",
+        distribution_version="12",
+        source_profile="debian-git",
+        oci_registry="ghcr.io",
+        oci_repository="thaloco/pyro-environments/debian-12",
+        oci_reference=DEFAULT_ENVIRONMENT_VERSION,
+    ),
+    "debian:12-base": VmEnvironment(
+        name="debian:12-base",
+        version=DEFAULT_ENVIRONMENT_VERSION,
+        description="Minimal Debian 12 environment for shell and core Unix tooling.",
+        default_packages=("bash", "coreutils"),
+        distribution="debian",
+        distribution_version="12",
+        source_profile="debian-base",
+        oci_registry="ghcr.io",
+        oci_repository="thaloco/pyro-environments/debian-12-base",
+        oci_reference=DEFAULT_ENVIRONMENT_VERSION,
+    ),
+    "debian:12-build": VmEnvironment(
+        name="debian:12-build",
+        version=DEFAULT_ENVIRONMENT_VERSION,
+        description="Debian 12 environment with Git and common build tools preinstalled.",
+        default_packages=("bash", "coreutils", "git", "gcc", "make", "cmake", "python3"),
+        distribution="debian",
+        distribution_version="12",
+        source_profile="debian-build",
+        oci_registry="ghcr.io",
+        oci_repository="thaloco/pyro-environments/debian-12-build",
+        oci_reference=DEFAULT_ENVIRONMENT_VERSION,
+    ),
+}
+
+
+def _default_cache_dir() -> Path:
+    return Path(
+        os.environ.get(
+            "PYRO_ENVIRONMENT_CACHE_DIR",
+            str(Path.home() / ".cache" / "pyro-mcp" / "environments"),
+        )
+    )
+
+
+def _manifest_profile_digest(runtime_paths: RuntimePaths, profile_name: str) -> str | None:
+    profiles = runtime_paths.manifest.get("profiles")
+    if not isinstance(profiles, dict):
+        return None
+    profile = profiles.get(profile_name)
+    if not isinstance(profile, dict):
+        return None
+    rootfs = profile.get("rootfs")
+    if not isinstance(rootfs, dict):
+        return None
+    raw_digest = rootfs.get("sha256")
+    return raw_digest if isinstance(raw_digest, str) else None
+
+
+def get_environment(name: str, *, runtime_paths: RuntimePaths | None = None) -> VmEnvironment:
+    """Resolve a curated environment by name."""
+    try:
+        spec = CATALOG[name]
+    except KeyError as exc:
+        known = ", ".join(sorted(CATALOG))
+        raise ValueError(f"unknown environment {name!r}; expected one of: {known}") from exc
+    if runtime_paths is None:
+        return spec
+    return VmEnvironment(
+        name=spec.name,
+        version=spec.version,
+        description=spec.description,
+        default_packages=spec.default_packages,
+        distribution=spec.distribution,
+        distribution_version=spec.distribution_version,
+        source_profile=spec.source_profile,
+        platform=spec.platform,
+        source_url=spec.source_url,
+        oci_registry=spec.oci_registry,
+        oci_repository=spec.oci_repository,
+        oci_reference=spec.oci_reference,
+        source_digest=_manifest_profile_digest(runtime_paths, spec.source_profile),
+        compatibility=spec.compatibility,
+    )
+
+
+def list_environments(*, runtime_paths: RuntimePaths | None = None) -> list[dict[str, object]]:
+    """Return catalog metadata in a JSON-safe format."""
+    return [
+        _serialize_environment(get_environment(name, runtime_paths=runtime_paths))
+        for name in sorted(CATALOG)
+    ]
+
+
+def _serialize_environment(environment: VmEnvironment) -> dict[str, object]:
+    return {
+        "name": environment.name,
+        "version": environment.version,
+        "description": environment.description,
+        "default_packages": list(environment.default_packages),
+        "distribution": environment.distribution,
+        "distribution_version": environment.distribution_version,
+        "platform": environment.platform,
+        "oci_registry": environment.oci_registry,
+        "oci_repository": environment.oci_repository,
+        "oci_reference": environment.oci_reference,
+        "source_digest": environment.source_digest,
+        "compatibility": environment.compatibility,
+    }
+
+
+class EnvironmentStore:
+    """Install and inspect curated environments in a local cache."""
+
+    def __init__(
+        self,
+        *,
+        runtime_paths: RuntimePaths,
+        cache_dir: Path | None = None,
+    ) -> None:
+        self._runtime_paths = runtime_paths
+        self._cache_dir = cache_dir or _default_cache_dir()
+        raw_platform = self._runtime_paths.manifest.get("platform", DEFAULT_PLATFORM)
+        platform = raw_platform if isinstance(raw_platform, str) else DEFAULT_PLATFORM
+        self._platform_dir = self._cache_dir / platform
+
+    @property
+    def cache_dir(self) -> Path:
+        return self._cache_dir
+
+    @property
+    def catalog_version(self) -> str:
+        return DEFAULT_CATALOG_VERSION
+
+    def list_environments(self) -> list[dict[str, object]]:
+        environments: list[dict[str, object]] = []
+        for name in sorted(CATALOG):
+            environments.append(self.inspect_environment(name))
+        return environments
+
+    def pull_environment(self, name: str) -> dict[str, object]:
+        installed = self.ensure_installed(name)
+        return {
+            **self.inspect_environment(name),
+            "install_dir": str(installed.install_dir),
+            "kernel_image": str(installed.kernel_image),
+            "rootfs_image": str(installed.rootfs_image),
+            "source": installed.source,
+        }
+
+    def inspect_environment(self, name: str) -> dict[str, object]:
+        spec = get_environment(name, runtime_paths=self._runtime_paths)
+        install_dir = self._install_dir(spec)
+        metadata_path = install_dir / "environment.json"
+        installed = metadata_path.exists() and (install_dir / "vmlinux").exists()
+        payload = _serialize_environment(spec)
+        payload.update(
+            {
+                "catalog_version": self.catalog_version,
+                "installed": installed,
+                "cache_dir": str(self._cache_dir),
+                "install_dir": str(install_dir),
+            }
+        )
+        if installed:
+            payload["install_manifest"] = str(metadata_path)
+        return payload
+
+    def ensure_installed(self, name: str) -> InstalledEnvironment:
+        spec = get_environment(name, runtime_paths=self._runtime_paths)
+        self._platform_dir.mkdir(parents=True, exist_ok=True)
+        install_dir = self._install_dir(spec)
+        metadata_path = install_dir / "environment.json"
+        if metadata_path.exists():
+            kernel_image = install_dir / "vmlinux"
+            rootfs_image = install_dir / "rootfs.ext4"
+            if kernel_image.exists() and rootfs_image.exists():
+                metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+                source = str(metadata.get("source", "cache"))
+                raw_digest = metadata.get("source_digest")
+                digest = raw_digest if isinstance(raw_digest, str) else None
+                return InstalledEnvironment(
+                    name=spec.name,
+                    version=spec.version,
+                    install_dir=install_dir,
+                    kernel_image=kernel_image,
+                    rootfs_image=rootfs_image,
+                    source=source,
+                    source_digest=digest,
+                    installed=True,
+                )
+
+        source_dir = self._runtime_paths.artifacts_dir / spec.source_profile
+        if source_dir.exists():
+            return self._install_from_local_source(spec, source_dir)
+        if (
+            spec.oci_registry is not None
+            and spec.oci_repository is not None
+            and spec.oci_reference is not None
+        ):
+            return self._install_from_oci(spec)
+        if spec.source_url is not None:
+            return self._install_from_archive(spec, spec.source_url)
+        raise RuntimeError(
+            f"environment {spec.name!r} is not installed and no downloadable source is configured"
+        )
+
+    def prune_environments(self) -> dict[str, object]:
+        deleted: list[str] = []
+        if not self._platform_dir.exists():
+            return {"deleted_environment_dirs": [], "count": 0}
+        for child in self._platform_dir.iterdir():
+            if child.name.startswith(".partial-"):
+                shutil.rmtree(child, ignore_errors=True)
+                deleted.append(child.name)
+                continue
+            if not child.is_dir():
+                continue
+            marker = child / "environment.json"
+            if not marker.exists():
+                shutil.rmtree(child, ignore_errors=True)
+                deleted.append(child.name)
+                continue
+            metadata = json.loads(marker.read_text(encoding="utf-8"))
+            raw_name = metadata.get("name")
+            raw_version = metadata.get("version")
+            if not isinstance(raw_name, str) or not isinstance(raw_version, str):
+                shutil.rmtree(child, ignore_errors=True)
+                deleted.append(child.name)
+                continue
+            try:
+                spec = get_environment(raw_name, runtime_paths=self._runtime_paths)
+            except ValueError:
+                shutil.rmtree(child, ignore_errors=True)
+                deleted.append(child.name)
+                continue
+            if spec.version != raw_version:
+                shutil.rmtree(child, ignore_errors=True)
+                deleted.append(child.name)
+        return {"deleted_environment_dirs": sorted(deleted), "count": len(deleted)}
+
+    def _install_dir(self, spec: VmEnvironment) -> Path:
+        normalized = spec.name.replace(":", "_")
+        return self._platform_dir / f"{normalized}-{spec.version}"
+
+    def _install_from_local_source(
+        self, spec: VmEnvironment, source_dir: Path
+    ) -> InstalledEnvironment:
+        install_dir = self._install_dir(spec)
+        temp_dir = Path(tempfile.mkdtemp(prefix=".partial-", dir=self._platform_dir))
+        try:
+            self._link_or_copy(source_dir / "vmlinux", temp_dir / "vmlinux")
+            self._link_or_copy(source_dir / "rootfs.ext4", temp_dir / "rootfs.ext4")
+            self._write_install_manifest(
+                temp_dir,
+                spec=spec,
+                source="bundled-runtime-source",
+                source_digest=spec.source_digest,
+            )
+            shutil.rmtree(install_dir, ignore_errors=True)
+            temp_dir.replace(install_dir)
+        except Exception:
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            raise
+        return InstalledEnvironment(
+            name=spec.name,
+            version=spec.version,
+            install_dir=install_dir,
+            kernel_image=install_dir / "vmlinux",
+            rootfs_image=install_dir / "rootfs.ext4",
+            source="bundled-runtime-source",
+            source_digest=spec.source_digest,
+            installed=True,
+        )
+
+    def _install_from_archive(self, spec: VmEnvironment, archive_url: str) -> InstalledEnvironment:
+        install_dir = self._install_dir(spec)
+        temp_dir = Path(tempfile.mkdtemp(prefix=".partial-", dir=self._platform_dir))
+        archive_path = temp_dir / "environment.tgz"
+        try:
+            urllib.request.urlretrieve(archive_url, archive_path)  # noqa: S310
+            self._extract_archive(archive_path, temp_dir)
+            kernel_image = self._locate_artifact(temp_dir, "vmlinux")
+            rootfs_image = self._locate_artifact(temp_dir, "rootfs.ext4")
+            if kernel_image.parent != temp_dir:
+                shutil.move(str(kernel_image), temp_dir / "vmlinux")
+            if rootfs_image.parent != temp_dir:
+                shutil.move(str(rootfs_image), temp_dir / "rootfs.ext4")
+            self._write_install_manifest(
+                temp_dir,
+                spec=spec,
+                source=archive_url,
+                source_digest=spec.source_digest,
+            )
+            archive_path.unlink(missing_ok=True)
+            shutil.rmtree(install_dir, ignore_errors=True)
+            temp_dir.replace(install_dir)
+        except Exception:
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            raise
+        return InstalledEnvironment(
+            name=spec.name,
+            version=spec.version,
+            install_dir=install_dir,
+            kernel_image=install_dir / "vmlinux",
+            rootfs_image=install_dir / "rootfs.ext4",
+            source=archive_url,
+            source_digest=spec.source_digest,
+            installed=True,
+        )
+
+    def _install_from_oci(self, spec: VmEnvironment) -> InstalledEnvironment:
+        install_dir = self._install_dir(spec)
+        temp_dir = Path(tempfile.mkdtemp(prefix=".partial-", dir=self._platform_dir))
+        try:
+            manifest, resolved_digest = self._fetch_oci_manifest(spec)
+            layers = manifest.get("layers")
+            if not isinstance(layers, list) or not layers:
+                raise RuntimeError("OCI manifest did not contain any layers")
+            for index, layer in enumerate(layers):
+                if not isinstance(layer, dict):
+                    raise RuntimeError("OCI manifest layer entry is malformed")
+                raw_digest = layer.get("digest")
+                if not isinstance(raw_digest, str):
+                    raise RuntimeError("OCI manifest layer is missing a digest")
+                blob_path = temp_dir / f"layer-{index}.tar"
+                self._download_oci_blob(spec, raw_digest, blob_path)
+                self._extract_tar_archive(blob_path, temp_dir)
+                blob_path.unlink(missing_ok=True)
+            kernel_image = self._locate_artifact(temp_dir, "vmlinux")
+            rootfs_image = self._locate_artifact(temp_dir, "rootfs.ext4")
+            if kernel_image.parent != temp_dir:
+                shutil.move(str(kernel_image), temp_dir / "vmlinux")
+            if rootfs_image.parent != temp_dir:
+                shutil.move(str(rootfs_image), temp_dir / "rootfs.ext4")
+            source = (
+                f"oci://{spec.oci_registry}/{spec.oci_repository}:{spec.oci_reference}"
+                if spec.oci_registry is not None
+                and spec.oci_repository is not None
+                and spec.oci_reference is not None
+                else "oci://unknown"
+            )
+            self._write_install_manifest(
+                temp_dir,
+                spec=spec,
+                source=source,
+                source_digest=resolved_digest or spec.source_digest,
+            )
+            shutil.rmtree(install_dir, ignore_errors=True)
+            temp_dir.replace(install_dir)
+        except Exception:
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            raise
+        return InstalledEnvironment(
+            name=spec.name,
+            version=spec.version,
+            install_dir=install_dir,
+            kernel_image=install_dir / "vmlinux",
+            rootfs_image=install_dir / "rootfs.ext4",
+            source=source,
+            source_digest=resolved_digest or spec.source_digest,
+            installed=True,
+        )
+
+    def _write_install_manifest(
+        self,
+        install_dir: Path,
+        *,
+        spec: VmEnvironment,
+        source: str,
+        source_digest: str | None,
+    ) -> None:
+        payload = {
+            "catalog_version": self.catalog_version,
+            "name": spec.name,
+            "version": spec.version,
+            "source": source,
+            "source_digest": source_digest,
+            "installed_at": int(time.time()),
+        }
+        (install_dir / "environment.json").write_text(
+            json.dumps(payload, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+    def _extract_archive(self, archive_path: Path, dest_dir: Path) -> None:
+        self._extract_tar_archive(archive_path, dest_dir)
+
+    def _locate_artifact(self, root: Path, name: str) -> Path:
+        for candidate in root.rglob(name):
+            if candidate.is_file():
+                return candidate
+        raise RuntimeError(f"environment archive did not contain {name}")
+
+    def _link_or_copy(self, source: Path, dest: Path) -> None:
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        relative_target = os.path.relpath(source, start=dest.parent)
+        try:
+            dest.symlink_to(relative_target)
+        except OSError:
+            shutil.copy2(source, dest)
+
+    def _fetch_oci_manifest(
+        self, spec: VmEnvironment
+    ) -> tuple[dict[str, Any], str | None]:
+        if spec.oci_registry is None or spec.oci_repository is None or spec.oci_reference is None:
+            raise RuntimeError("OCI source metadata is incomplete")
+        headers = {"Accept": OCI_MANIFEST_ACCEPT}
+        payload, response_headers = self._request_bytes(
+            self._oci_url(
+                spec.oci_registry,
+                spec.oci_repository,
+                f"manifests/{spec.oci_reference}",
+            ),
+            headers=headers,
+            repository=spec.oci_repository,
+        )
+        manifest = json.loads(payload.decode("utf-8"))
+        if not isinstance(manifest, dict):
+            raise RuntimeError("OCI manifest response was not a JSON object")
+        resolved_digest = response_headers.get("Docker-Content-Digest")
+        media_type = manifest.get("mediaType")
+        if media_type in {
+            "application/vnd.oci.image.index.v1+json",
+            "application/vnd.docker.distribution.manifest.list.v2+json",
+        }:
+            manifests = manifest.get("manifests")
+            if not isinstance(manifests, list):
+                raise RuntimeError("OCI index did not contain manifests")
+            selected = self._select_oci_manifest_descriptor(manifests)
+            payload, response_headers = self._request_bytes(
+                self._oci_url(
+                    spec.oci_registry,
+                    spec.oci_repository,
+                    f"manifests/{selected}",
+                ),
+                headers=headers,
+                repository=spec.oci_repository,
+            )
+            manifest = json.loads(payload.decode("utf-8"))
+            if not isinstance(manifest, dict):
+                raise RuntimeError("OCI child manifest response was not a JSON object")
+            resolved_digest = response_headers.get("Docker-Content-Digest") or selected
+        return manifest, resolved_digest
+
+    def _download_oci_blob(self, spec: VmEnvironment, digest: str, dest: Path) -> None:
+        if spec.oci_registry is None or spec.oci_repository is None:
+            raise RuntimeError("OCI source metadata is incomplete")
+        payload, _ = self._request_bytes(
+            self._oci_url(
+                spec.oci_registry,
+                spec.oci_repository,
+                f"blobs/{digest}",
+            ),
+            headers={},
+            repository=spec.oci_repository,
+        )
+        dest.write_bytes(payload)
+
+    def _request_bytes(
+        self,
+        url: str,
+        *,
+        headers: dict[str, str],
+        repository: str,
+    ) -> tuple[bytes, dict[str, str]]:
+        request = urllib.request.Request(url, headers=headers, method="GET")
+        try:
+            with urllib.request.urlopen(request, timeout=90) as response:  # noqa: S310
+                return response.read(), dict(response.headers.items())
+        except urllib.error.HTTPError as exc:
+            if exc.code != 401:
+                raise RuntimeError(f"failed to fetch OCI resource {url}: {exc}") from exc
+            authenticate = exc.headers.get("WWW-Authenticate")
+            if authenticate is None:
+                raise RuntimeError("OCI registry denied access without an auth challenge") from exc
+            token = self._fetch_registry_token(authenticate, repository)
+            authenticated_request = urllib.request.Request(
+                url,
+                headers={**headers, "Authorization": f"Bearer {token}"},
+                method="GET",
+            )
+            with urllib.request.urlopen(authenticated_request, timeout=90) as response:  # noqa: S310
+                return response.read(), dict(response.headers.items())
+
+    def _fetch_registry_token(self, authenticate: str, repository: str) -> str:
+        if not authenticate.startswith("Bearer "):
+            raise RuntimeError("unsupported OCI authentication scheme")
+        params = self._parse_authenticate_parameters(authenticate[len("Bearer ") :])
+        realm = params.get("realm")
+        if realm is None:
+            raise RuntimeError("OCI auth challenge did not include a token realm")
+        query = {
+            "service": params.get("service", ""),
+            "scope": params.get("scope", f"repository:{repository}:pull"),
+        }
+        token_url = f"{realm}?{urllib.parse.urlencode(query)}"
+        with urllib.request.urlopen(token_url, timeout=90) as response:  # noqa: S310
+            payload = json.loads(response.read().decode("utf-8"))
+        if not isinstance(payload, dict):
+            raise RuntimeError("OCI auth token response was not a JSON object")
+        raw_token = payload.get("token") or payload.get("access_token")
+        if not isinstance(raw_token, str) or raw_token == "":
+            raise RuntimeError("OCI auth token response did not include a bearer token")
+        return raw_token
+
+    def _parse_authenticate_parameters(self, raw: str) -> dict[str, str]:
+        params: dict[str, str] = {}
+        for segment in raw.split(","):
+            if "=" not in segment:
+                continue
+            key, value = segment.split("=", 1)
+            params[key.strip()] = value.strip().strip('"')
+        return params
+
+    def _select_oci_manifest_descriptor(self, manifests: list[Any]) -> str:
+        for manifest in manifests:
+            if not isinstance(manifest, dict):
+                continue
+            platform = manifest.get("platform")
+            if not isinstance(platform, dict):
+                continue
+            os_name = platform.get("os")
+            architecture = platform.get("architecture")
+            raw_digest = manifest.get("digest")
+            if (
+                isinstance(os_name, str)
+                and isinstance(architecture, str)
+                and isinstance(raw_digest, str)
+                and os_name == "linux"
+                and architecture in {"amd64", "x86_64"}
+            ):
+                return raw_digest
+        raise RuntimeError("OCI index did not contain a linux/amd64 manifest")
+
+    def _extract_tar_archive(self, archive_path: Path, dest_dir: Path) -> None:
+        dest_root = dest_dir.resolve()
+        with tarfile.open(archive_path, "r:*") as archive:
+            for member in archive.getmembers():
+                member_path = (dest_dir / member.name).resolve()
+                if not member_path.is_relative_to(dest_root):
+                    raise RuntimeError(f"unsafe archive member path: {member.name}")
+            archive.extractall(dest_dir, filter="data")
+
+    def _oci_url(self, registry: str, repository: str, suffix: str) -> str:
+        return f"https://{registry}/v2/{repository}/{suffix}"