ops/t4exec.py

#!/usr/bin/env python3
"""Tunnel envelope executor — https://aitun.cc/DED976KN (protocol v1).

Commands:
  status                 GET /health + /status
  run '<code>' [tmo]     execute python one-liner/snippet remotely
  runfile <py> [tmo]     execute a local .py file remotely
  get <remote> [tmo]     chunked download (8 MB per round-trip)
  put <local> <remote>   chunked upload (45 KB per round-trip)
"""
import base64
import json
import os
import sys
import urllib.request
import zlib

# Tunnel endpoint: override with `export T4_URL=https://aitun.cc/XXXX` when a
# new tunnel code is issued (aitun URLs rotate every Colab session).
URL = os.environ.get("T4_URL", "https://aitun.cc/ZMPBD8M2")
CRC_MODE = None  # auto-detected against known-good envelope


def b64u(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).decode().rstrip("=")


def b64d(s: str) -> bytes:
    return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))


def _crc(p_b64: str, raw: bytes) -> str:
    # Server-side check verified live: crc32 of the DECODED payload (raw code
    # bytes), NOT of the b64 text. 'print(123, 456)' -> crc c03b37d5 accepted.
    return format(zlib.crc32(raw) & 0xFFFFFFFF, "08x")


def envelope(code: str) -> bytes:
    # NOTE: payload must be a SINGLE b64url line. The 76-char wrapped form is
    # in the spec but the current relay corrupts multi-line payloads
    # (crc mismatch server-side) — single-line verified against crc 9e30c354.
    raw = code.encode()
    p = b64u(raw)
    head = f"samaicmdbegin\nv=1\nenc=b64url\ncrc={_crc(p, raw)}\n"
    return (head + p + "\nsamaicmdend\n").encode()


def _detect_crc():
    """Self-test: print(40+2) must yield the known-good crc 9e30c354."""
    global CRC_MODE
    raw = b"print(40+2)"
    p = b64u(raw)
    if format(zlib.crc32(p.encode()) & 0xFFFFFFFF, "08x") == "9e30c354":
        CRC_MODE = "b64"
    elif format(zlib.crc32(raw) & 0xFFFFFFFF, "08x") == "9e30c354":
        CRC_MODE = "raw"
    else:  # pragma: no cover
        raise SystemExit("FATAL: cannot reproduce known-good crc 9e30c354")
    assert envelope("print(40+2)").decode().splitlines()[3] == "crc=9e30c354"


def execute(code: str, timeout: int = 600) -> dict:
    url = f"{URL}/execute?respenc=b64url&timeout={timeout}"
    req = urllib.request.Request(
        url, data=envelope(code), headers={"Content-Type": "text/plain"}
    )
    with urllib.request.urlopen(req, timeout=timeout + 60) as r:
        body = json.loads(r.read().decode())
    out = {}
    # server flips between b64 keys and plaintext keys — accept both
    for k in ("stdout_b64", "stderr_b64"):
        v = body.get(k)
        if v:
            out[k[: -len("_b64")]] = b64d(v).decode(errors="replace")
    for k in ("stdout", "stderr"):
        v = body.get(k)
        if v and k not in out:
            out[k] = v if isinstance(v, str) else str(v)
    for k in ("status", "error", "execution_time", "return_code",
              "success", "error_type"):
        if k in body:
            out[k] = body[k]
    for k in ("error_b64", "error_type_b64", "traceback_b64"):
        v = body.get(k)
        if v and k[: -len("_b64")] not in out:
            try:
                out[k[: -len("_b64")]] = b64d(v).decode(errors="replace")
            except Exception:
                pass
    return out


def get(remote: str, timeout: int = 600) -> bytes:
    """Chunked download: remote prints B64:<chunk> per 8 MB slice."""
    chunks = []
    off = 0
    reader = (
        "import base64\n"
        f"f=open({remote!r},'rb'); f.seek({{off}}); c=f.read(8000000); f.close()\n"
        "import sys; sys.stdout.write('B64:'+base64.b64encode(c).decode()+"
        "'\\nEOF' if len(c)<8000000 else '\\n')"
    )
    while True:
        r = execute(reader.format(off=off), timeout=timeout)
        so = r.get("stdout", "")
        if "B64:" not in so:
            raise RuntimeError(f"get failed at off={off}: {r}")
        for line in so.splitlines():
            if line.startswith("B64:"):
                chunks.append(b64d(line[4:]))
        if "EOF" in so:
            break
        off += 8000000
    return b"".join(chunks)


def put(local: str, remote: str, timeout: int = 600) -> None:
    data = open(local, "rb").read()
    init = (
        "import os\n"
        f"os.makedirs(os.path.dirname({remote!r}) or '.', exist_ok=True)\n"
        f"open({remote!r},'wb').close()\n"
        "print('ok')"
    )
    r0 = execute(init, timeout=60)
    if r0.get("stdout", "").strip() != "ok":
        raise RuntimeError(f"put init failed: {r0}")
    app = (
        "import base64\n"
        f"f=open({remote!r},'ab'); "
        "f.write(base64.urlsafe_b64decode(sys_stdin + '='*(-len(sys_stdin)%4))); "
        "f.close(); print('ok')"
    ).replace("sys_stdin", "__PAYLOAD__")
    n = (len(data) + 44999) // 45000
    for i in range(n):
        chunk_b64 = b64u(data[i * 45000 : (i + 1) * 45000])
        code = app.replace("__PAYLOAD__", repr(chunk_b64))
        r = execute(code, timeout=timeout)
        if r.get("stdout", "").strip() != "ok":
            raise RuntimeError(f"put chunk {i+1}/{n} failed: {r}")
    print(f"put ok: {local} -> {remote} ({len(data)} bytes, {n} chunks)")


def main():
    _detect_crc()
    cmd = sys.argv[1] if len(sys.argv) > 1 else "status"
    if cmd == "status":
        for ep in ("/health", "/status"):
            try:
                with urllib.request.urlopen(URL + ep, timeout=15) as r:
                    print(ep, "->", r.read().decode()[:2500])
            except Exception as e:
                print(ep, "ERROR:", e)
    elif cmd == "run":
        code = sys.argv[2]
        tmo = int(sys.argv[3]) if len(sys.argv) > 3 else 600
        r = execute(code, tmo)
        print(r.get("stdout", ""))
        if r.get("stderr"):
            print("[stderr]", r["stderr"], file=sys.stderr)
        if r.get("error"):
            print("[error]", r["error"], file=sys.stderr)
    elif cmd == "runfile":
        tmo = int(sys.argv[3]) if len(sys.argv) > 3 else 600
        r = execute(open(sys.argv[2]).read(), tmo)
        print(r.get("stdout", ""))
        if r.get("stderr"):
            print("[stderr]", r["stderr"], file=sys.stderr)
    elif cmd == "get":
        tmo = int(sys.argv[3]) if len(sys.argv) > 3 else 600
        data = get(sys.argv[2], tmo)
        sys.stdout.buffer.write(data)
    elif cmd == "put":
        put(sys.argv[2], sys.argv[3])
    else:
        raise SystemExit(__doc__)


if __name__ == "__main__":
    main()