thaloco/pyro-mcp
1
Fork 0
Code Pull requests Releases Packages Activity Actions

Add direct GHCR environment pulls

Browse source
This commit is contained in:
Thales Maciel 2026-03-08 16:08:01 -03:00
parent 5d5243df23
commit 75082467f9
5 changed files with 346 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

95
src/pyro_mcp/vm_environments.py
View file

@ -2,6 +2,7 @@
from __future__ import annotations
import hashlib
import json
import os
import shutil
@ -27,6 +28,7 @@ OCI_MANIFEST_ACCEPT = ", ".join(
"application/vnd.docker.distribution.manifest.v2+json",
)
)
OCI_READ_CHUNK_SIZE = 1024 * 1024
@dataclass(frozen=True)
@ -381,6 +383,8 @@ class EnvironmentStore:
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))
resolved_digest: str | None = None
source = "oci://unknown"
try:
manifest, resolved_digest = self._fetch_oci_manifest(spec)
layers = manifest.get("layers")
@ -402,13 +406,12 @@ class EnvironmentStore:
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"
)
if spec.oci_registry is not None and spec.oci_repository is not None:
source = f"oci://{spec.oci_registry}/{spec.oci_repository}"
if resolved_digest is not None:
source = f"{source}@{resolved_digest}"
elif spec.oci_reference is not None:
source = f"{source}:{spec.oci_reference}"
self._write_install_manifest(
temp_dir,
spec=spec,
@ -487,12 +490,14 @@ class EnvironmentStore:
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")
resolved_digest = response_headers.get("docker-content-digest")
if resolved_digest is not None:
self._verify_digest_bytes(payload, resolved_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",
}:
} or isinstance(manifest.get("manifests"), list):
manifests = manifest.get("manifests")
if not isinstance(manifests, list):
raise RuntimeError("OCI index did not contain manifests")
@ -509,22 +514,39 @@ class EnvironmentStore:
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
resolved_digest = response_headers.get("docker-content-digest") or selected
self._verify_digest_bytes(payload, resolved_digest)
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)
digest_algorithm, digest_value = self._split_digest(digest)
if digest_algorithm != "sha256":
raise RuntimeError(f"unsupported OCI blob digest algorithm: {digest_algorithm}")
hasher = hashlib.sha256()
with (
self._open_request(
self._oci_url(
spec.oci_registry,
spec.oci_repository,
f"blobs/{digest}",
),
headers={},
repository=spec.oci_repository,
) as response,
dest.open("wb") as handle,
):
while True:
chunk = response.read(OCI_READ_CHUNK_SIZE)
if not chunk:
break
hasher.update(chunk)
handle.write(chunk)
if hasher.hexdigest() != digest_value:
raise RuntimeError(
f"OCI blob digest mismatch for {digest}; got sha256:{hasher.hexdigest()}"
)
def _request_bytes(
self,
@ -533,10 +555,19 @@ class EnvironmentStore:
headers: dict[str, str],
repository: str,
) -> tuple[bytes, dict[str, str]]:
with self._open_request(url, headers=headers, repository=repository) as response:
return response.read(), {key.lower(): value for key, value in response.headers.items()}
def _open_request(
self,
url: str,
*,
headers: dict[str, str],
repository: str,
) -> Any:
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())
return urllib.request.urlopen(request, timeout=90) # noqa: S310
except urllib.error.HTTPError as exc:
if exc.code != 401:
raise RuntimeError(f"failed to fetch OCI resource {url}: {exc}") from exc
@ -549,8 +580,10 @@ class EnvironmentStore:
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())
try:
return urllib.request.urlopen(authenticated_request, timeout=90) # noqa: S310
except urllib.error.HTTPError as auth_exc:
raise RuntimeError(f"failed to fetch OCI resource {url}: {auth_exc}") from auth_exc
def _fetch_registry_token(self, authenticate: str, repository: str) -> str:
if not authenticate.startswith("Bearer "):
@ -613,3 +646,17 @@ class EnvironmentStore:
def _oci_url(self, registry: str, repository: str, suffix: str) -> str:
return f"https://{registry}/v2/{repository}/{suffix}"
def _split_digest(self, digest: str) -> tuple[str, str]:
algorithm, separator, value = digest.partition(":")
if separator == "" or value == "":
raise RuntimeError(f"invalid OCI digest: {digest}")
return algorithm, value
def _verify_digest_bytes(self, payload: bytes, digest: str) -> None:
algorithm, value = self._split_digest(digest)
if algorithm != "sha256":
raise RuntimeError(f"unsupported OCI digest algorithm: {algorithm}")
actual = hashlib.sha256(payload).hexdigest()
if actual != value:
raise RuntimeError(f"OCI digest mismatch for {digest}; got sha256:{actual}")