Prepare inputs, then verify with networking disabled

import base64
import os
import time

import httpx

PREPARE_INPUTS = r"""
mkdir -p /fixture
printf 'expected-v1\n' > /fixture/expected.txt
"""

VERIFY_OFFLINE = r"""
test "$(cat /fixture/expected.txt)" = expected-v1
for interface in /sys/class/net/*; do
    test "${interface##*/}" = lo
done
printf 'fixture-ok network-disabled\n'
"""


def main():
    with httpx.Client(
        base_url=os.getenv(
            "CONTREE_BASE_URL", "https://api.tokenfactory.nebius.com/sandboxes"
        ).rstrip("/") + "/",
        headers={
            "Authorization": f"Bearer {os.environ['NEBIUS_API_KEY']}",
            "Project": os.environ["NEBIUS_PROJECT_ID"],
        },
        timeout=15,
    ) as client:
        prepared, _ = run_operation(
            client,
            os.environ["IMAGE_UUID"],
            PREPARE_INPUTS,
            networking=True,
            disposable=False,
        )
        checkpoint = prepared.get("result_image_uuid")
        if not checkpoint:
            raise RuntimeError("Preparation returned no saved image")
        print(f"Checkpoint: {checkpoint}", flush=True)

        verified, output = run_operation(
            client,
            checkpoint,
            VERIFY_OFFLINE,
            networking=False,
            disposable=True,
        )
        if (
            output != "fixture-ok network-disabled\n"
            or verified.get("result_image_uuid")
        ):
            raise RuntimeError("Verification output or disposable result was unexpected")
        print(output, end="")


def run_operation(client, image, command, *, networking, disposable):
    # Submit once. A transport error can leave the submission outcome uncertain.
    response = client.post(
        "v1/instances",
        json={
            "image": image,
            "command": "set -e\n" + command,
            "shell": True,
            "networking": {"enabled": networking},
            "disposable": disposable,
            "timeout": 30,
            "truncate_output_at": 1048576,
        },
    )
    response.raise_for_status()
    if response.status_code != 201:
        raise RuntimeError(f"Unexpected create status: {response.status_code}")
    operation_id = response.json()["uuid"]
    print(f"Operation: {operation_id}", flush=True)
    deadline = time.monotonic() + 90
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise TimeoutError(f"Inspect operation {operation_id} before retrying")
        response = client.get(f"v1/operations/{operation_id}", timeout=min(15, remaining))
        response.raise_for_status()
        operation = response.json()
        status = operation["status"]
        if status == "SUCCESS":
            break
        if status not in {"PENDING", "ASSIGNED", "EXECUTING"}:
            raise RuntimeError(f"Operation {operation_id} ended as {status}")
        retry_after = response.headers.get("Retry-After", "1")
        delay = max(1, int(retry_after)) if retry_after.isdecimal() else 1
        time.sleep(max(0, min(delay, deadline - time.monotonic())))

    result = operation["metadata"]["result"]
    state = result["state"]
    stdout = decode_output(result["stdout"])
    stderr = decode_output(result["stderr"])
    if (
        state["exit_code"] != 0
        or state.get("timed_out")
        or state.get("signal") not in (None, -1)
    ):
        raise RuntimeError(f"Process failed in {operation_id}: {state}; stderr={stderr}")
    return operation, stdout


def decode_output(stream):
    if stream.get("truncated"):
        raise RuntimeError("Output was truncated; increase the output limit")
    if stream["encoding"] == "base64":
        return base64.b64decode(stream["value"], validate=True).decode("utf-8")
    if stream["encoding"] in {"ascii", "utf-8", "utf8"}:
        return stream["value"]
    raise ValueError(f"Unsupported output encoding: {stream['encoding']}")


if __name__ == "__main__":
    main()