Generate and download a binary artifact

import asyncio
import gzip
import hashlib
from pathlib import Path

from contree_sdk import Contree

PAYLOAD = b"untrusted input\x00\x01\xff\n"
GENERATE_ARTIFACT = r"""
import gzip
import hashlib
from pathlib import Path

output_directory = Path("/out")
output_directory.mkdir(parents=True, exist_ok=True)
input_bytes = Path("/input.bin").read_bytes()
compressed_bytes = gzip.compress(input_bytes, mtime=0)
(output_directory / "result.bin").write_bytes(compressed_bytes)
digest = hashlib.sha256(compressed_bytes).hexdigest()
(output_directory / "result.sha256").write_text(digest + "\n")
"""


async def main() -> None:
    client = Contree()
    base = await client.images.oci("docker.io/library/python:3.12-alpine")
    run = await base.run(
        shell="python /generate_artifact.py",
        files={
            "/input.bin": PAYLOAD,
            "/generate_artifact.py": GENERATE_ARTIFACT.encode(),
        },
        disposable=False,
        timeout=30,
    )
    if run.exit_code != 0 or run.uuid is None:
        raise RuntimeError('Artifact generation failed or returned no saved image')
    print(f"Artifact saved in image {run.uuid}; downloading")
    destination = Path("downloaded-result.bin")
    await run.download("/out/result.bin", destination)
    remote_digest = (await run.read("/out/result.sha256")).decode().strip()
    actual = destination.read_bytes()
    if actual[:8] != b"\x1f\x8b\x08\x00\x00\x00\x00\x00":
        raise RuntimeError('Downloaded file has an unexpected gzip header')
    if gzip.decompress(actual) != PAYLOAD:
        raise RuntimeError('Downloaded artifact does not contain the original input')
    digest = hashlib.sha256(actual).hexdigest()
    if digest != remote_digest:
        raise RuntimeError('Downloaded bytes do not match the remote SHA-256 digest')
    print(f"Downloaded {destination}: {len(actual)} bytes; contents and digest match")
    print(f"image={run.uuid} bytes={len(actual)} sha256={digest}")


asyncio.run(main())