Run Codex in a CoreWeave sandbox

import os
from pathlib import Path
import shlex
import sys

from cwsandbox import AuthStrategy, Endpoint, Sandbox, Service

if len(sys.argv) != 4:
    raise SystemExit(
        "Usage: python create_codex_ssh.py REPOSITORY_URL SSH_KEY OUTPUT_DIR"
    )
repository, identity_arg, directory_arg = sys.argv[1:]
identity = Path(identity_arg).expanduser().resolve()
if not identity.is_file():
    raise SystemExit("SSH private key not found. Generate the key pair first.")
public_key = Path(str(identity) + ".pub").read_bytes()
api_key = os.environ["OPENAI_API_KEY"]
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY first.")
directory = Path(directory_arg).expanduser().resolve()
directory.mkdir(mode=0o700, parents=True, exist_ok=False)
sandbox = Sandbox.run(
    auth=AuthStrategy.WANDB,
    container_image="node:22-bookworm",
    resources={"cpu": "2", "memory": "4Gi"},
    max_lifetime_seconds=8 * 3600,
    services=[
        Service(
            port=8443,
            name="ssh",
            visibility="public",
            endpoint=Endpoint(kind="tls_passthrough"),
        )
    ],
)


def run(command):
    result = sandbox.exec(command, timeout_seconds=900).result()
    if result.returncode != 0:
        details = (result.stdout + result.stderr).replace(api_key, "[REDACTED]")
        raise RuntimeError(details[-3000:])
    return result.stdout


try:
    (directory / "sandbox-id").write_text(sandbox.sandbox_id + "\n")
    print("Sandbox created; ID saved in", directory / "sandbox-id", flush=True)
    sandbox.wait()
    address = next(
        item.address for item in sandbox.service_addresses if item.port == 8443
    )
    host, port = address.rsplit(":", 1)
    run(
        [
            "bash",
            "-c",
            """
set -eu
apt-get update -qq
DEBIAN_FRONTEND=noninteractive apt-get install -y -qq openssh-server stunnel4 bubblewrap
npm install --global @openai/codex@0.154.0
useradd --create-home --shell /bin/bash agent
usermod --password '*' agent
install -d -m 700 -o agent -g agent /home/agent/.ssh
install -d -m 755 /run/sshd /workspace
mkdir -p /etc/codex-ssh
chmod 700 /etc/codex-ssh
""",
        ]
    )
    sandbox.write_file("/home/agent/.ssh/authorized_keys", public_key).result()
    run(
        [
            "bash",
            "-c",
            "chmod 600 /home/agent/.ssh/authorized_keys; chown agent:agent /home/agent/.ssh/authorized_keys",
        ]
    )
    run(["git", "clone", "--", repository, "/workspace/project"])
    run(["chown", "-R", "agent:agent", "/workspace/project"])
    login = sandbox.exec(
        ["runuser", "-u", "agent", "--", "codex", "login", "--with-api-key"],
        stdin=True,
        timeout_seconds=60,
    )
    login.stdin.write((api_key + "\n").encode()).result()
    login.stdin.close().result()
    if login.result().returncode != 0:
        raise RuntimeError("Codex authentication failed.")
    run(
        [
            "openssl",
            "req",
            "-x509",
            "-newkey",
            "rsa:2048",
            "-noenc",
            "-keyout",
            "/etc/codex-ssh/tls.key",
            "-out",
            "/etc/codex-ssh/tls.crt",
            "-days",
            "1",
            "-subj",
            "/CN=codex-sandbox",
            "-addext",
            f"subjectAltName=DNS:{host}",
        ]
    )
    sandbox.write_file(
        "/etc/codex-ssh/sshd_config",
        b"""Port 2222
ListenAddress 127.0.0.1
HostKey /etc/ssh/ssh_host_ed25519_key
AuthorizedKeysFile .ssh/authorized_keys
PubkeyAuthentication yes
PasswordAuthentication no
KbdInteractiveAuthentication no
PermitRootLogin no
AllowUsers agent
AllowTcpForwarding local
X11Forwarding no
Subsystem sftp internal-sftp
""",
    ).result()
    sandbox.write_file(
        "/etc/codex-ssh/stunnel.conf",
        b"""pid = /run/codex-stunnel.pid
output = /var/log/codex-stunnel.log
[ssh]
accept = 0.0.0.0:8443
connect = 127.0.0.1:2222
cert = /etc/codex-ssh/tls.crt
key = /etc/codex-ssh/tls.key
""",
    ).result()
    run(["chmod", "600", "/etc/codex-ssh/tls.key"])
    run(
        [
            "/usr/sbin/sshd",
            "-f",
            "/etc/codex-ssh/sshd_config",
            "-E",
            "/var/log/codex-sshd.log",
        ]
    )
    run(["stunnel", "/etc/codex-ssh/stunnel.conf"])
    certificate = directory / "tls.crt"
    certificate.write_bytes(sandbox.read_file("/etc/codex-ssh/tls.crt").result())
    host_key = (
        sandbox.read_file("/etc/ssh/ssh_host_ed25519_key.pub").result().decode().strip()
    )
    known_hosts = directory / "known_hosts"
    known_hosts.write_text("cw-codex " + host_key + "\n")
    forwarder = directory / "tls_proxy.py"
    forwarder.write_text('''"""Forward SSH through verified TLS, with backpressure in both directions."""
import asyncio
import ssl
import sys

async def main():
    host, port, certificate = sys.argv[1:]
    context = ssl.create_default_context(cafile=certificate)
    remote_reader, remote_writer = await asyncio.open_connection(
        host, int(port), ssl=context, server_hostname=host
    )
    local_reader = asyncio.StreamReader()
    transport, _ = await asyncio.get_running_loop().connect_read_pipe(
        lambda: asyncio.StreamReaderProtocol(local_reader), sys.stdin.buffer
    )
    async def upload():
        while data := await local_reader.read(65536):
            remote_writer.write(data)
            await remote_writer.drain()
    def write_output(data):
        sys.stdout.buffer.write(data)
        sys.stdout.buffer.flush()
    async def download():
        while data := await remote_reader.read(65536):
            await asyncio.to_thread(write_output, data)
    tasks = [asyncio.create_task(upload()), asyncio.create_task(download())]
    try:
        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for task in done:
            task.result()
    finally:
        for task in tasks:
            task.cancel()
        transport.close()
        remote_writer.close()
        await remote_writer.wait_closed()

try:
    asyncio.run(main())
except (OSError, ssl.SSLError) as error:
    print(f'TLS proxy: {error}', file=sys.stderr)
    sys.exit(1)
''')
    proxy = shlex.join(
        [sys.executable, str(forwarder), "%h", "%p", str(certificate)]
    )
    config = directory / "ssh_config"
    config.write_text(f'''Host cw-codex
    HostName {host}
    Port {port}
    User agent
    IdentityFile "{identity}"
    IdentitiesOnly yes
    HostKeyAlias cw-codex
    UserKnownHostsFile "{known_hosts}"
    StrictHostKeyChecking yes
    ServerAliveInterval 30
    ProxyCommand {proxy}
''')
    print("SSH config:", config, flush=True)
except BaseException as error:
    try:
        sandbox.stop().result()
    except BaseException:
        error.add_note(
            f"Automatic cleanup failed. Stop sandbox {sandbox.sandbox_id} manually."
        )
    raise