diff --git a/examples/sandbox-backend/.env.example b/examples/sandbox-backend/.env.example new file mode 100644 index 000000000..6f18ad3c0 --- /dev/null +++ b/examples/sandbox-backend/.env.example @@ -0,0 +1,9 @@ +# CubeSandbox native SDK configuration (see sdk/python/cubesandbox/_config.py) +CUBE_API_URL=http://127.0.0.1:3000 +CUBE_TEMPLATE_ID=tpl-xxxxxxxxxxxxxxxx + +# Optional overrides +# CUBE_SANDBOX_USER=root # user envd runs commands as inside the sandbox +# CUBE_SANDBOX_TIMEOUT=1800 # sandbox TTL in seconds, refreshed on every reconnect +# CUBE_EXEC_TIMEOUT=120 # default per-command timeout in seconds +# CUBE_HOOK_STATE_DIR=~/.cache/cubesandbox-hook # where session->sandbox_id mappings live diff --git a/examples/sandbox-backend/.gitignore b/examples/sandbox-backend/.gitignore new file mode 100644 index 000000000..cff554307 --- /dev/null +++ b/examples/sandbox-backend/.gitignore @@ -0,0 +1,3 @@ +.env +__pycache__/ +*.pyc diff --git a/examples/sandbox-backend/DESIGN.md b/examples/sandbox-backend/DESIGN.md new file mode 100644 index 000000000..9d9fa23d0 --- /dev/null +++ b/examples/sandbox-backend/DESIGN.md @@ -0,0 +1,206 @@ +# CubeSandbox Hook 方案设计记录 + +> 状态：已实现，见同目录下 `cubesandbox_exec.py` / `cubesandbox_rewrite.py` / `install.sh` / `README.md` +> 日期：2026-07-05（设计）/ 2026-07-06（实现） +> 分支：`feat/sandbox-backend` + +## 相对本文档的实现变更 + +实现时对"待解决问题"一节做了如下处理，供后续review参考： + +1. **`cd` 不再特殊放行**：不判断命令是否是 `cd`，而是让 `cd`/`export` 都统一走 + `cubesandbox-exec`。执行时在沙箱内用每会话随机、仅所有者可读写的状态目录 + （`/tmp/.cubesandbox-state-`）持久化 cwd 和已导出的 + 环境变量：每次执行前 `source`/`cd` 恢复，执行后写回。这样多次独立的 + `commands.run()` 调用之间也能保持连续的 shell 状态，不需要在 hook 里猜测哪些 + 命令是"纯 shell builtin"。 +2. **文件系统一致性**：创建沙箱时，用 `host-mount`（见 `../host-mount/`）把 hook + 输入里的 `cwd`（Claude Code 项目目录）以只读方式挂载到沙箱内的同一路径， + 使 Bash 能读取与 Read/Write/Edit 在 host 上相同的项目文件，同时不能通过挂载修改 host。如果该路径 + 不在 CubeMaster 的 `allowed_host_mount_prefixes` 白名单内，会捕获 `ApiError` 并 + 降级为不挂载的普通沙箱（打印警告，不中断执行）。 +3. **沙箱生命周期**：一个 Claude Code session（hook 输入里的 `session_id`）对应 + 一个沙箱，用本地状态文件（`~/.cache/cubesandbox-hook/.json`）缓存 + `sandbox_id`，后续调用 `Sandbox.connect()` 复用；提供 `--reset` 手动销毁。 +4. **性能**：未做量化测试；每次调用仍是一次新的 Python 进程 + 一次 + `commands.run()` HTTP 往返，冷启动成本已通过会话级复用消除，但比原生本地 Bash + 仍有网络往返开销，待实测补充数据。 + +## 背景 + +目前 `examples/sandboxed-claude/` 已实现了两种 Claude Code 集成方式： + +1. **Headless 执行**：在沙箱内部运行 Claude Code 本身 +2. **MCP Server**：Claude Code 作为 MCP client，通过 `sandbox_run_code` / `sandbox_run_command` 等工具显式调用沙箱 + +MCP 方案的问题是：Claude Code **必须主动选择**沙箱工具，如果 AI 选择直接用 Bash 工具就能绕过沙箱。 + +## 灵感来源：RTK (Rust Token Killer) + +RTK 是一个用 Rust 写的 CLI 代理，通过 Claude Code 的 **PreToolUse Hook** 透明拦截所有 Bash 命令，改写后再交还 Claude Code 执行。AI 完全不知道命令被改写过。 + +### RTK 的三件套 + +``` +~/.claude/ +├── settings.json ← 注册 PreToolUse hook +├── hooks/ +│ └── rtk-rewrite.sh ← 拦截脚本 +├── RTK.md ← 元命令说明 +└── CLAUDE.md ← 追加 @RTK.md 引用 +``` + +### RTK 的数据流 + +``` +Claude Code 要执行 "git status" + │ + ▼ +PreToolUse hook 触发 → rtk-rewrite.sh + │ + ▼ +rtk rewrite "git status" → 返回 "rtk git status" + │ + ▼ +返回 JSON: {"updatedInput": {"command": "rtk git status"}} + │ + ▼ +Claude Code 执行改写后的命令（输出被 rtk 压缩 ~80%） +``` + +### rtk-rewrite.sh 核心逻辑 + +```bash +INPUT=$(cat) # 从 stdin 读 PreToolUse JSON +COMMAND=$(提取 tool_input.command) + +# 已有 rtk 前缀的命令直接放行 +[[ "$COMMAND" == rtk\ * ]] && exit 0 + +# 调用 rtk Rust 二进制做改写 +REWRITTEN=$(rtk rewrite "$COMMAND") + +# 返回 updatedInput 给 Claude Code +cat < --mount "npm test" + ▼ +Claude Code executes the rewritten command with its normal Bash tool + │ + ▼ +cubesandbox-exec (cubesandbox_exec.py) + │ reuses (or creates) one sandbox per Claude Code session + ▼ +CubeAPI ──► CubeMaster ──► Cubelet ──► KVM MicroVM + │ + stdout/stderr/exit code + returned to Claude Code +``` + +The AI never sees `cubesandbox-exec` in its own reasoning -- it issued +`npm test`, and it gets back exactly the output/exit code `npm test` would +have produced, except it ran inside a disposable MicroVM instead of on your +machine. + +## Quick Start + +```bash +cd examples/sandbox-backend +cp .env.example .env +# edit .env: set CUBE_API_URL and CUBE_TEMPLATE_ID (see cubemastercli tpl list) + +./install.sh +``` + +`install.sh` will: +1. Copy the two scripts into `~/.claude/hooks/` +2. Install a `cubesandbox-exec` shim on `$PATH` (default `~/.local/bin`) +3. Register the `PreToolUse` hook in `~/.claude/settings.json` (idempotent) +4. `pip install` the `cubesandbox` SDK + `python-dotenv` + +Restart Claude Code, then try: + +``` +> Run `ls -la` and tell me what's here +``` + +Claude Code will show a normal Bash tool call and normal-looking output -- +but it actually ran inside a CubeSandbox MicroVM. Verify with: + +```bash +cubesandbox-exec --session "hostname; cat /proc/1/cgroup | head -1" +``` + +To remove everything: + +```bash +./install.sh --uninstall +``` + +## Filesystem consistency + +Claude Code's `Read` / `Write` / `Edit` tools operate on the **host** +filesystem (wherever `claude` was launched), while `Bash` commands now run +**inside the sandbox**. Without extra care those are two different +filesystems and `cd $PROJECT && cat file.py` inside the sandbox would fail. + +This example solves it with CubeSandbox's [`host-mount`](../host-mount) +extension: the first time a sandbox is created for a session, +`cubesandbox_exec.py` bind-mounts the project directory (the hook's `cwd` +field) **read-only at the same path** inside the sandbox. Sandbox-side Bash +commands can safely read the same source files as host-side tools, but host +tools remain the sole writer for project files. + +This requires the project directory to be under one of CubeMaster's +`allowed_host_mount_prefixes` (default: `/data/shared/`): + +```yaml +# CubeMaster config +extra_conf: + allowed_host_mount_prefixes: + - "/data/shared/" + - "/home/you/projects/" # add your project root +``` + +If the path isn't allowed, `cubesandbox-exec` **does not fail** -- it logs a +warning and falls back to a sandbox with its own isolated filesystem. Bash +commands still run safely sandboxed, they just won't see files created via +Read/Write/Edit (and vice versa). For a fully consistent view, add your +project root to `allowed_host_mount_prefixes` or keep it under +`/data/shared/`. + +## Session & shell state + +- **One sandbox per Claude Code session.** `cubesandbox_rewrite.py` passes + `session_id` from the hook payload; `cubesandbox_exec.py` caches the + resulting `sandbox_id` in `~/.cache/cubesandbox-hook/.json` + and reconnects to it on every subsequent Bash call, so you don't pay a + cold-start cost per command. +- **`cd` / `export` persist across calls.** Each `commands.run()` invocation + is otherwise stateless, so before running your command the wrapper + restores `$PWD` and exported vars from a randomly named, owner-only + directory inside the sandbox's `/tmp`, and updates them afterwards. From + Claude Code's point of view this behaves like one + continuous shell session, the same as its native (non-sandboxed) Bash + tool. +- **Manual reset.** `cubesandbox-exec --reset --session ` kills the + cached sandbox and clears its state, forcing a clean sandbox on the next + call (handy between unrelated tasks in the same Claude Code session). + +## Files + +``` +sandbox-backend/ +├── DESIGN.md # Design notes / rationale (RTK-inspired) +├── cubesandbox_exec.py # Sandbox exec CLI (session cache, state persistence, host-mount) +├── cubesandbox_rewrite.py # PreToolUse hook: rewrites Bash tool_input.command +├── sandbox_exec.py # Standalone CLI (e2b SDK, opt-in) +├── mcp_server.py # MCP server (opt-in sandbox tools for Claude Code) +├── install.sh # Installs hooks + settings.json registration (+ --uninstall) +├── test_cubesandbox_exec.py # Session cache and sandbox lifecycle tests +├── test_cubesandbox_rewrite.py # Hook command-rewrite safety tests +├── test_mcp_server.py # MCP framing and tool dispatch tests +├── requirements.txt +├── .env.example +└── README.md # This file +``` + +## Troubleshooting + +| Symptom | Likely cause | Fix | +|---|---|---| +| Bash output unchanged, no sandbox created | Hook not registered / Claude Code not restarted | Check `~/.claude/settings.json`, restart `claude` | +| `CUBE_TEMPLATE_ID is not set` | `.env` missing/not loaded | `cp .env.example .env` and fill in values; `cubesandbox-exec` loads `.env` from its own directory | +| `hostPath ... is not within an allowed mount prefix` (warning) | Project dir not under `allowed_host_mount_prefixes` | Add it to CubeMaster config, or ignore -- sandbox still works, just without shared FS | +| Commands run but `cd`/`export` don't stick between calls | The session's private state directory under `/tmp` is not writable in the template image | Ensure the template's `/tmp` is writable by `CUBE_SANDBOX_USER` | +| Every command spins up a brand-new sandbox | `session_id` missing from hook payload, or `~/.cache/cubesandbox-hook` not writable | Check hook stdin payload has `session_id`; check `CUBE_HOOK_STATE_DIR` permissions | +| `SandboxNotFoundError` / stale sandbox | Sandbox TTL (`CUBE_SANDBOX_TIMEOUT`) expired between calls | Increase `CUBE_SANDBOX_TIMEOUT`, or accept the automatic recreate (some cwd/env state is lost) | + +## Related + +- [`../sandboxed-claude/`](../sandboxed-claude/) -- run Claude Code itself + inside a sandbox (pattern A) +- [`../host-mount/`](../host-mount/) -- the host-mount mechanism used here +- Claude Code Hooks reference: https://code.claude.com/docs/en/hooks diff --git a/examples/sandbox-backend/conftest.py b/examples/sandbox-backend/conftest.py new file mode 100644 index 000000000..e8a3bfbdd --- /dev/null +++ b/examples/sandbox-backend/conftest.py @@ -0,0 +1,22 @@ +# Copyright (c) 2024 Tencent Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Shared pytest fixtures for the sandbox-backend example tests.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +EXAMPLE_DIR = Path(__file__).resolve().parent +if str(EXAMPLE_DIR) not in sys.path: + sys.path.insert(0, str(EXAMPLE_DIR)) + + +@pytest.fixture +def tmp_state_dir(tmp_path, monkeypatch): + """Provide an isolated temp directory for cubesandbox_exec state files.""" + state_dir = tmp_path / "hook-state" + monkeypatch.setenv("CUBE_HOOK_STATE_DIR", str(state_dir)) + return state_dir diff --git a/examples/sandbox-backend/cubesandbox_exec.py b/examples/sandbox-backend/cubesandbox_exec.py new file mode 100644 index 000000000..762648bbb --- /dev/null +++ b/examples/sandbox-backend/cubesandbox_exec.py @@ -0,0 +1,287 @@ +#!/usr/bin/env python3 +""" +CubeSandbox exec backend for the PreToolUse hook. + +Runs a single shell command inside a CubeSandbox MicroVM and behaves like +``ssh host ''`` / ``docker exec bash -c ''``: + +- stdout / stderr are passed through byte-for-byte +- the command's real exit code is propagated via ``sys.exit()`` +- one sandbox is created per Claude Code **session** and reused for every + subsequent call (see ``--session``), so repeated invocations don't pay + a cold-start cost +- the sandbox's cwd and exported environment variables are persisted + *inside the sandbox filesystem* between calls, so a bare ``cd foo`` or + ``export X=1`` issued by one Bash tool call is visible to the next one + -- exactly like Claude Code's own persistent host shell session. + +This script is normally invoked by ``cubesandbox_rewrite.py`` (the +PreToolUse hook), not by a human, but it is fully usable standalone: + + cubesandbox-exec "npm test" + cubesandbox-exec --session my-session "git status" + cubesandbox-exec --mount /data/shared/myproject "cd /data/shared/myproject && ls" + cubesandbox-exec --reset --session my-session +""" + +from __future__ import annotations + +import argparse +import fcntl +import json +import os +import secrets +import sys +from pathlib import Path + +# Stdlib-only at module level so that `--help` / `--reset` can at least +# parse args without the optional `cubesandbox` / `python-dotenv` deps +# being installed. The heavy imports happen in _bootstrap() below. +_SCRIPT_DIR = Path(__file__).resolve().parent + +# Populated by _bootstrap() (called from main() after argparse). +ApiError = Config = CubeSandboxError = Sandbox = SandboxNotFoundError = None # type: ignore[assignment] +TEMPLATE_ID = "" +SANDBOX_USER = "root" +SANDBOX_TTL = 1800 +STATE_DIR = Path(os.getenv("CUBE_HOOK_STATE_DIR", os.path.expanduser("~/.cache/cubesandbox-hook"))) +DEFAULT_SESSION = "default" + +# One lock file per STATE_DIR; fcntl.flock serializes concurrent +# cubesandbox-exec processes that race to create a sandbox for the same +# brand-new session (Claude Code emits parallel Bash tool calls). +_LOCK_FILE = STATE_DIR / ".create.lock" + + +def _bootstrap() -> None: + """Load .env (from the script's own dir, not the process CWD) and import + the cubesandbox SDK + read config. Called lazily from main() so that + argparse --help works even before deps are installed.""" + global ApiError, Config, CubeSandboxError, Sandbox, SandboxNotFoundError + global TEMPLATE_ID, SANDBOX_USER, SANDBOX_TTL + + try: + from dotenv import load_dotenv + except ImportError: + load_dotenv = None # type: ignore[assignment] + + if load_dotenv: + load_dotenv(_SCRIPT_DIR / ".env") + load_dotenv() # optional: let CWD .env override for dev/debug + + from cubesandbox import ( # noqa: E402 + ApiError as _ApiError, + Config as _Config, + CubeSandboxError as _CubeSandboxError, + Sandbox as _Sandbox, + SandboxNotFoundError as _SandboxNotFoundError, + ) + ApiError, Config, CubeSandboxError, Sandbox, SandboxNotFoundError = ( + _ApiError, _Config, _CubeSandboxError, _Sandbox, _SandboxNotFoundError, + ) + + TEMPLATE_ID = os.getenv("CUBE_TEMPLATE_ID", "") + SANDBOX_USER = os.getenv("CUBE_SANDBOX_USER", "root") + SANDBOX_TTL = int(os.getenv("CUBE_SANDBOX_TIMEOUT", "1800")) + + +# ── Session -> sandbox_id persistence (on the host running this script) ─ + +def _state_path(session_id: str) -> Path: + safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in session_id) or DEFAULT_SESSION + return STATE_DIR / f"{safe}.json" + + +def _load_state(session_id: str) -> dict: + path = _state_path(session_id) + if path.exists(): + try: + return json.loads(path.read_text()) + except (json.JSONDecodeError, OSError): + pass + return {} + + +def _save_state(session_id: str, state: dict) -> None: + STATE_DIR.mkdir(parents=True, exist_ok=True) + _state_path(session_id).write_text(json.dumps(state)) + + +def _clear_state(session_id: str) -> None: + path = _state_path(session_id) + if path.exists(): + path.unlink() + + +# ── Sandbox lifecycle ─────────────────────────────────────────────────── + +def get_sandbox(session_id: str, mount: str | None) -> Sandbox: + """Reuse the sandbox already created for this Claude Code session, or + create a fresh one on first use / after it has expired or been killed. + + The create+save critical section is guarded by a file lock so that + concurrent cubesandbox-exec processes (Claude Code emits parallel Bash + tool calls) don't each create their own sandbox and orphan the loser. + """ + # Fast path: reuse an existing cached sandbox without locking. + state = _load_state(session_id) + sandbox_id = state.get("sandbox_id") + if sandbox_id: + try: + return _connect(sandbox_id) + except SandboxNotFoundError: + pass # sandbox is gone -- fall through and create a new one + except CubeSandboxError as e: + print(f"[cubesandbox-exec] warning: reconnect failed ({e}); creating a new sandbox", file=sys.stderr) + + if not TEMPLATE_ID: + print("[cubesandbox-exec] error: CUBE_TEMPLATE_ID is not set", file=sys.stderr) + sys.exit(127) + + # Slow path: create under an exclusive lock so only one process wins. + STATE_DIR.mkdir(parents=True, exist_ok=True) + with open(_LOCK_FILE, "w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + # Re-read state inside the lock -- another process may have just + # created the sandbox while we were waiting. + state = _load_state(session_id) + sandbox_id = state.get("sandbox_id") + if sandbox_id: + try: + return _connect(sandbox_id) + except SandboxNotFoundError: + pass + except CubeSandboxError as e: + print(f"[cubesandbox-exec] warning: reconnect failed ({e}); creating a new sandbox", file=sys.stderr) + + sandbox = _create_sandbox(mount) + _save_state(session_id, { + "sandbox_id": sandbox.sandbox_id, + "mount": mount, + "state_token": secrets.token_urlsafe(24), + }) + return sandbox + + +def _connect(sandbox_id: str) -> Sandbox: + """Reconnect to an existing sandbox, refreshing its TTL to + ``SANDBOX_TTL`` (Sandbox.connect's default uses Config.timeout=300, + which would ignore the user's CUBE_SANDBOX_TIMEOUT setting).""" + return Sandbox.connect(sandbox_id, config=Config(timeout=SANDBOX_TTL)) + + +def _create_sandbox(mount: str | None) -> Sandbox: + """Create a sandbox, optionally bind-mounting *mount* (Claude Code's + project cwd) read-only at the same path. Falls back to a plain sandbox if the host-mount is rejected + (e.g. the path isn't under CubeMaster's ``allowed_host_mount_prefixes``). + """ + if mount: + try: + return Sandbox.create( + TEMPLATE_ID, + timeout=SANDBOX_TTL, + metadata={ + "host-mount": json.dumps([ + {"hostPath": mount, "mountPath": mount, "readOnly": True}, + ]) + }, + ) + except ApiError as e: + print( + f"[cubesandbox-exec] warning: host-mount of {mount!r} was rejected ({e}); " + "falling back to a sandbox without a shared filesystem. Bash commands will " + "run in isolation and won't see files created via the Read/Write/Edit tools. " + "See README.md 'Filesystem consistency' section.", + file=sys.stderr, + ) + return Sandbox.create(TEMPLATE_ID, timeout=SANDBOX_TTL) + + +# ── Command execution ─────────────────────────────────────────────────── + +def run(command: str, session_id: str, timeout: float | None, mount: str | None) -> int: + sandbox = get_sandbox(session_id, mount) + state = _load_state(session_id) + state_token = state.get("state_token") + if not isinstance(state_token, str): + print("[cubesandbox-exec] error: sandbox state is missing", file=sys.stderr) + return 1 + + default_cwd = mount or "$HOME" + state_dir = f"/tmp/.cubesandbox-state-{state_token}" + cwd_file = f"{state_dir}/cwd" + env_file = f"{state_dir}/env" + wrapped = ( + f"umask 077; mkdir -p -- {state_dir}; " + f"[ -d {state_dir} ] && [ ! -L {state_dir} ] || exit 1; " + f"[ -f {env_file} ] && [ ! -L {env_file} ] && source {env_file} >/dev/null 2>&1; " + f'cd "$(cat {cwd_file} 2>/dev/null || echo {default_cwd})" 2>/dev/null; ' + f"{command}\n" + "__CBX_STATUS__=$?; " + f"pwd > {cwd_file} 2>/dev/null; " + f"export -p > {env_file} 2>/dev/null; " + "exit $__CBX_STATUS__" + ) + + try: + result = sandbox.commands.run(wrapped, timeout=timeout, user=SANDBOX_USER) + except CubeSandboxError as e: + print(f"[cubesandbox-exec] error: {e}", file=sys.stderr) + return 1 + + if result.stdout: + sys.stdout.write(result.stdout) + if result.stderr: + sys.stderr.write(result.stderr) + return result.exit_code + + +def reset(session_id: str) -> None: + state = _load_state(session_id) + sandbox_id = state.get("sandbox_id") + if sandbox_id: + try: + _connect(sandbox_id).kill() + except CubeSandboxError: + pass + _clear_state(session_id) + print(f"[cubesandbox-exec] sandbox for session {session_id!r} destroyed") + + +# ── CLI ────────────────────────────────────────────────────────────────── + +def main() -> None: + parser = argparse.ArgumentParser(description="Run a shell command inside a cached CubeSandbox MicroVM") + parser.add_argument("command", nargs="?", help="Shell command to run (passed through to bash -c)") + parser.add_argument( + "--session", default=DEFAULT_SESSION, + help="Session key used to cache/reuse the sandbox (typically Claude Code's session_id)", + ) + parser.add_argument( + "--timeout", type=float, default=float(os.getenv("CUBE_EXEC_TIMEOUT", "120")), + help="Command timeout in seconds (default 120)", + ) + parser.add_argument( + "--mount", default=None, + help="Host directory to bind-mount read-only into the sandbox at the same path " + "(only used the first time a sandbox is created for --session). Must be under " + "CubeMaster's allowed_host_mount_prefixes.", + ) + parser.add_argument("--reset", action="store_true", help="Destroy the cached sandbox for this session and exit") + args = parser.parse_args() + + if args.reset: + _bootstrap() + reset(args.session) + return + + if not args.command: + parser.print_help() + sys.exit(2) + + _bootstrap() + sys.exit(run(args.command, args.session, args.timeout, args.mount)) + + +if __name__ == "__main__": + main() diff --git a/examples/sandbox-backend/cubesandbox_rewrite.py b/examples/sandbox-backend/cubesandbox_rewrite.py new file mode 100644 index 000000000..a75eb9dfb --- /dev/null +++ b/examples/sandbox-backend/cubesandbox_rewrite.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +""" +PreToolUse hook: transparently redirect every Bash command Claude Code runs +into an isolated CubeSandbox MicroVM. + +Registered (by install.sh) in ~/.claude/settings.json as: + +{ + "hooks": { + "PreToolUse": [ + { + "matcher": "Bash", + "hooks": [ + {"type": "command", "command": "~/.claude/hooks/cubesandbox_rewrite.py"} + ] + } + ] + } +} + +For every Bash tool call, Claude Code pipes a JSON payload like: + +{ + "session_id": "abc123", + "cwd": "/home/user/myproject", + "hook_event_name": "PreToolUse", + "tool_name": "Bash", + "tool_input": {"command": "npm test", "timeout": 120000} +} + +This script rewrites tool_input.command to: + + cubesandbox-exec --session abc123 --mount /home/user/myproject --timeout 120.000 'npm test' + +and returns it via `updatedInput`, so Claude Code executes the *rewritten* +command with its normal Bash tool -- the AI never sees the sandbox layer. +""" + +from __future__ import annotations + +import json +import shlex +import sys + +EXEC_BIN = "cubesandbox-exec" + + +def _has_unquoted_newline(command: str) -> bool: + """Return True if *command* contains a \\n or \\r outside of quotes. + + Bash treats unquoted newlines as command separators, so a command + like ``cubesandbox-exec\\nrm -rf /`` would execute ``rm -rf /`` on + the host — we must never let such a command bypass rewriting. + """ + in_single = in_double = in_ansi_c = False + index = 0 + while index < len(command): + ch = command[index] + if in_ansi_c: + if ch == "\\" and index + 1 < len(command): + if command[index + 1] in "nr": + return True + index += 2 + continue + if ch == "'": + in_ansi_c = False + elif in_single: + if ch == "'": + in_single = False + elif in_double: + if ch == "\\" and index + 1 < len(command): + index += 2 + continue + if ch == '"': + in_double = False + else: + if ch == "\\" and index + 1 < len(command): + index += 2 + continue + if ch == "$" and index + 1 < len(command) and command[index + 1] == "'": + in_ansi_c = True + index += 2 + continue + if ch == "'": + in_single = True + elif ch == '"': + in_double = True + elif ch in "\n\r": + return True + index += 1 + return False + + +def _already_sandboxed(command: str) -> bool: + """Return True if *command* is already a cubesandbox-exec invocation. + + Rejects commands containing unquoted newlines (see + ``_has_unquoted_newline``) so that multi-line injection like + ``cubesandbox-exec\\nrm -rf /`` is NOT mistaken for an already-sandboxed + call. Uses ``shlex.split()`` to correctly tokenize single-line commands, + and matches the binary by basename so that full-path invocations (e.g. + ``/home/user/.local/bin/cubesandbox-exec``) are also recognized, + preventing recursive wrapping. + """ + if _has_unquoted_newline(command): + return False + try: + tokens = shlex.split(command) + except ValueError: + return False + if not tokens: + return False + first = tokens[0] + return first == EXEC_BIN or first.rsplit("/", 1)[-1] == EXEC_BIN + + +def main() -> None: + try: + payload = json.load(sys.stdin) + except json.JSONDecodeError: + sys.exit(0) # malformed input -- don't block Claude Code + + if payload.get("tool_name") != "Bash": + sys.exit(0) + + tool_input = payload.get("tool_input") or {} + command = tool_input.get("command") + if not isinstance(command, str) or not command or _already_sandboxed(command): + sys.exit(0) + + session_id = payload.get("session_id") or "default" + cwd = payload.get("cwd") + + rewritten = [EXEC_BIN, "--session", session_id] + if cwd: + rewritten += ["--mount", cwd] + + timeout_ms = tool_input.get("timeout") + if isinstance(timeout_ms, (int, float)) and timeout_ms > 0: + rewritten += ["--timeout", f"{timeout_ms / 1000:.3f}"] + + rewritten.append(command) + new_command = " ".join(shlex.quote(a) for a in rewritten) + + updated_input = dict(tool_input) + updated_input["command"] = new_command + + output = { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": updated_input, + } + } + print(json.dumps(output)) + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/examples/sandbox-backend/install.sh b/examples/sandbox-backend/install.sh new file mode 100644 index 000000000..506aded78 --- /dev/null +++ b/examples/sandbox-backend/install.sh @@ -0,0 +1,141 @@ +#!/usr/bin/env bash +# Install the CubeSandbox PreToolUse hook for Claude Code. +# +# What this does: +# 1. Copies cubesandbox_exec.py / cubesandbox_rewrite.py into ~/.claude/hooks/ +# 2. Installs a `cubesandbox-exec` shim on your PATH (~/.local/bin by default) +# 3. Registers the PreToolUse hook in ~/.claude/settings.json (idempotent -- +# safe to re-run, never clobbers your other settings/hooks) +# +# Usage: +# ./install.sh # install for the current user (~/.claude) +# CLAUDE_DIR=/path ./install.sh # install into a different Claude config dir +# ./install.sh --uninstall # remove the hook registration + shim + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CLAUDE_DIR="${CLAUDE_DIR:-$HOME/.claude}" +HOOKS_DIR="$CLAUDE_DIR/hooks" +SETTINGS_FILE="$CLAUDE_DIR/settings.json" +BIN_DIR="${CUBE_BIN_DIR:-$HOME/.local/bin}" + +REWRITE_HOOK_PATH="$HOOKS_DIR/cubesandbox_rewrite.py" +EXEC_SCRIPT_PATH="$HOOKS_DIR/cubesandbox_exec.py" +SHIM_PATH="$BIN_DIR/cubesandbox-exec" + +# ── settings.json merge helper (pure python, no jq dependency) ───────── +merge_hook() { +python3 - "$SETTINGS_FILE" "$REWRITE_HOOK_PATH" <<'PYEOF' +import json +import sys +from pathlib import Path + +settings_path, hook_path = sys.argv[1], sys.argv[2] +path = Path(settings_path) +data = {} +if path.exists() and path.stat().st_size > 0: + data = json.loads(path.read_text()) + +hooks = data.setdefault("hooks", {}) +pretool = hooks.setdefault("PreToolUse", []) + +entry = {"type": "command", "command": hook_path} +for group in pretool: + if group.get("matcher") == "Bash": + for h in group.get("hooks", []): + if h.get("command") == hook_path: + print("Hook already registered, skipping.") + sys.exit(0) + group.setdefault("hooks", []).append(entry) + break +else: + pretool.append({"matcher": "Bash", "hooks": [entry]}) + +path.parent.mkdir(parents=True, exist_ok=True) +path.write_text(json.dumps(data, indent=2) + "\n") +print(f"Registered PreToolUse hook in {settings_path}") +PYEOF +} + +unmerge_hook() { +python3 - "$SETTINGS_FILE" "$REWRITE_HOOK_PATH" <<'PYEOF' +import json +import sys +from pathlib import Path + +settings_path, hook_path = sys.argv[1], sys.argv[2] +path = Path(settings_path) +if not path.exists(): + sys.exit(0) + +data = json.loads(path.read_text()) +pretool = data.get("hooks", {}).get("PreToolUse", []) +for group in pretool: + if group.get("matcher") == "Bash": + group["hooks"] = [h for h in group.get("hooks", []) if h.get("command") != hook_path] +pretool[:] = [g for g in pretool if g.get("hooks")] + +path.write_text(json.dumps(data, indent=2) + "\n") +print(f"Removed PreToolUse hook from {settings_path}") +PYEOF +} + +if [[ "${1:-}" == "--uninstall" ]]; then + if command -v python3 >/dev/null 2>&1; then + unmerge_hook + else + echo "python3 not found; skipping settings.json cleanup." >&2 + fi + rm -f "$REWRITE_HOOK_PATH" "$EXEC_SCRIPT_PATH" "$HOOKS_DIR/.env" "$SHIM_PATH" + echo "Uninstalled. (Sandboxes created by previous sessions are not auto-killed --" + echo "run 'cubesandbox-exec --reset --session ' or clean up via CubeMaster.)" + exit 0 +fi + +command -v python3 >/dev/null 2>&1 || { echo "python3 is required" >&2; exit 1; } + +echo "Installing CubeSandbox hook backend for Claude Code..." + +mkdir -p "$HOOKS_DIR" "$BIN_DIR" + +install -m 0755 "$SCRIPT_DIR/cubesandbox_rewrite.py" "$REWRITE_HOOK_PATH" +install -m 0755 "$SCRIPT_DIR/cubesandbox_exec.py" "$EXEC_SCRIPT_PATH" +# Copy .env so that cubesandbox_exec.py can load CUBE_TEMPLATE_ID at runtime +if [[ -f "$SCRIPT_DIR/.env" ]]; then + cp "$SCRIPT_DIR/.env" "$HOOKS_DIR/.env" + echo "Copied .env to $HOOKS_DIR/.env" +fi + +cat > "$SHIM_PATH" <&2 + exit 1 +fi + +# Back up settings.json before mutating it so a botched merge is recoverable. +if [[ -f "$SETTINGS_FILE" ]]; then + cp "$SETTINGS_FILE" "$SETTINGS_FILE.bak.$(date +%s)" +fi + +merge_hook + +echo +echo "Done." +echo +if [[ ":$PATH:" != *":$BIN_DIR:"* ]]; then + echo "WARNING: $BIN_DIR is not on your PATH. Add this to your shell profile:" + echo " export PATH=\"$BIN_DIR:\$PATH\"" + echo +fi +if [[ ! -f "$SCRIPT_DIR/.env" ]]; then + echo "Next step: cp .env.example .env and fill in CUBE_TEMPLATE_ID / CUBE_API_URL," + echo "then restart Claude Code. Every Bash command it runs will now execute inside" + echo "an isolated CubeSandbox MicroVM." +fi diff --git a/examples/sandbox-backend/mcp_server.py b/examples/sandbox-backend/mcp_server.py new file mode 100644 index 000000000..ea7248b45 --- /dev/null +++ b/examples/sandbox-backend/mcp_server.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +""" +CubeSandbox MCP Server for Claude Code. + +Provides tools that let Claude Code automatically execute untrusted code +in isolated MicroVM sandboxes instead of directly on the host. + +Add to ~/.claude/mcp.json or project .claude/mcp.json: + +{ + "mcpServers": { + "cubesandbox": { + "command": "python3", + "args": [ + "/path/to/CubeSandbox/examples/sandbox-backend/mcp_server.py" + ], + "env": { + "E2B_API_URL": "http://127.0.0.1:3000", + "E2B_API_KEY": "e2b_000000", + "CUBE_TEMPLATE_ID": "tpl-c703537d5106496790d44702" + } + } + } +} +""" + +import atexit +import json +import os +import shlex +import sys +import traceback +# ── Configuration ────────────────────────────────────────────────────── +E2B_API_URL = os.getenv("E2B_API_URL", "http://127.0.0.1:3000") +E2B_API_KEY = os.getenv("E2B_API_KEY", "e2b_000000") +TEMPLATE_ID = os.getenv("CUBE_TEMPLATE_ID", "") +SANDBOX_TTL = int(os.getenv("CUBE_SANDBOX_TIMEOUT", "600")) + +_sandbox = None + + +def _cleanup_sandbox(): + """Kill the cached sandbox on process exit to avoid orphans.""" + global _sandbox + if _sandbox is not None: + try: + _sandbox.kill() + except Exception: + pass + _sandbox = None + + +def get_sandbox(timeout=SANDBOX_TTL): + """Lazy-create and reuse a sandbox, refreshing its TTL on each call. + + Without TTL refresh the sandbox expires after ``SANDBOX_TTL`` seconds + of inactivity, causing the next tool call to fail with + ``SandboxNotFoundError``. An ``atexit`` handler ensures the sandbox + is killed when the MCP server process exits. + """ + global _sandbox + if _sandbox is None: + from e2b_code_interpreter import Sandbox + _sandbox = Sandbox.create(TEMPLATE_ID, timeout=timeout) + atexit.register(_cleanup_sandbox) + else: + try: + _sandbox.set_timeout(timeout) + except Exception: + # Sandbox may have expired or been killed — create a fresh one. + from e2b_code_interpreter import Sandbox + _sandbox = Sandbox.create(TEMPLATE_ID, timeout=timeout) + return _sandbox + + +def run_command(cmd, timeout=120): + """Run a command inside the sandbox.""" + from e2b.sandbox.commands.command_handle import CommandExitException + try: + result = get_sandbox().commands.run(cmd, timeout=timeout) + return {"exit_code": result.exit_code, "stdout": result.stdout.strip(), "stderr": result.stderr.strip()} + except CommandExitException as e: + return {"exit_code": getattr(e, "exit_code", 1), "stdout": "", "stderr": str(e)} + + +# ── MCP Tool Definitions ────────────────────────────────────────────── + +TOOLS = [ + { + "name": "sandbox_run_code", + "description": "Execute Python code in an isolated CubeSandbox MicroVM. Use this for running untrusted code, testing generated scripts, or installing packages safely.", + "inputSchema": { + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "Python code to execute in the sandbox", + }, + "timeout": { + "type": "integer", + "description": "Execution timeout in seconds (default 120)", + "default": 120, + }, + }, + "required": ["code"], + }, + }, + { + "name": "sandbox_run_command", + "description": "Execute a shell command in an isolated CubeSandbox MicroVM. Use this to safely test shell commands, explore file systems, or run build tools without affecting the host.", + "inputSchema": { + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "Shell command to execute in the sandbox", + }, + "timeout": { + "type": "integer", + "description": "Execution timeout in seconds (default 120)", + "default": 120, + }, + }, + "required": ["command"], + }, + }, + { + "name": "sandbox_write_file", + "description": "Write content to a file inside the sandbox.", + "inputSchema": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Absolute path inside the sandbox (e.g., /tmp/script.py)", + }, + "content": { + "type": "string", + "description": "File content to write", + }, + }, + "required": ["path", "content"], + }, + }, + { + "name": "sandbox_read_file", + "description": "Read a file's content from inside the sandbox.", + "inputSchema": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Absolute path inside the sandbox", + }, + }, + "required": ["path"], + }, + }, + { + "name": "sandbox_reset", + "description": "Destroy the current sandbox and create a fresh one. Use this between unrelated tasks to get a clean environment.", + "inputSchema": { + "type": "object", + "properties": {}, + }, + }, +] + + +# ── MCP Protocol Handler ────────────────────────────────────────────── + +def handle_request(request): + method = request.get("method", "") + req_id = request.get("id") + + if method == "initialize": + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": { + "name": "cubesandbox-mcp", + "version": "1.0.0", + }, + }, + } + + elif method == "tools/list": + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"tools": TOOLS}, + } + + elif method == "tools/call": + tool_name = request["params"]["name"] + arguments = request["params"].get("arguments", {}) + + try: + if tool_name == "sandbox_run_code": + code = arguments["code"] + timeout = arguments.get("timeout", 120) + result = run_command(f"python3 -c {shlex.quote(code)}", timeout=timeout) + text = result["stdout"] or result["stderr"] or "(no output)" + + elif tool_name == "sandbox_run_command": + cmd = arguments["command"] + timeout = arguments.get("timeout", 120) + result = run_command(cmd, timeout=timeout) + text = result["stdout"] or result["stderr"] or "(no output)" + + elif tool_name == "sandbox_write_file": + path = arguments["path"] + content = arguments["content"] + get_sandbox().files.write(path, content) + text = f"Written {len(content)} bytes to {path}" + + elif tool_name == "sandbox_read_file": + path = arguments["path"] + content = get_sandbox().files.read(path) + text = content if isinstance(content, str) else content.decode("utf-8", errors="replace") + + elif tool_name == "sandbox_reset": + global _sandbox + if _sandbox: + _sandbox.kill() + _sandbox = None + text = "Sandbox destroyed. A new one will be created on next use." + + else: + text = f"Unknown tool: {tool_name}" + + except Exception as e: + traceback.print_exc(file=sys.stderr) + text = f"Error: {e}" + + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": text}], + }, + } + + elif method == "notifications/initialized": + return None # No response for notifications + + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"}, + } + + +# ── Main Loop ───────────────────────────────────────────────────────── + +def _read_mcp_message(): + """Read one MCP message from stdin using Content-Length framing. + + Returns the parsed JSON request, ``None`` for a malformed (skippable) + message, or raises ``EOFError`` when stdin is exhausted so the caller + can break out of the read loop instead of spinning. + """ + content_length = None + read_any = False + for line in sys.stdin: + read_any = True + line = line.rstrip("\r\n") + if not line: + break # empty line = end of headers + if line.startswith("Content-Length:"): + try: + content_length = int(line.split(":", 1)[1].strip()) + except ValueError: + print(f"[cubesandbox-mcp] warning: malformed Content-Length header: {line!r}", file=sys.stderr) + return None + if not read_any: + raise EOFError # stdin exhausted — client disconnected + if content_length is None: + return None + body = sys.stdin.read(content_length) + try: + return json.loads(body) + except json.JSONDecodeError as e: + print(f"[cubesandbox-mcp] warning: malformed JSON body: {e}", file=sys.stderr) + return None + + +def main(): + """Run the MCP server on stdio (JSON-RPC).""" + from dotenv import load_dotenv + load_dotenv() + while True: + try: + request = _read_mcp_message() + except EOFError: + break + if request is None: + continue + response = handle_request(request) + if response is not None: + body = json.dumps(response) + sys.stdout.write(f"Content-Length: {len(body)}\r\n\r\n{body}") + sys.stdout.flush() + + +if __name__ == "__main__": + main() diff --git a/examples/sandbox-backend/requirements.txt b/examples/sandbox-backend/requirements.txt new file mode 100644 index 000000000..27c89bca7 --- /dev/null +++ b/examples/sandbox-backend/requirements.txt @@ -0,0 +1,3 @@ +cubesandbox +e2b-code-interpreter +python-dotenv diff --git a/examples/sandbox-backend/sandbox_exec.py b/examples/sandbox-backend/sandbox_exec.py new file mode 100644 index 000000000..330a5dc1d --- /dev/null +++ b/examples/sandbox-backend/sandbox_exec.py @@ -0,0 +1,137 @@ +""" +Sandbox execution backend for Claude Code. + +Claude Code runs on the HOST and calls this script to execute +untrusted code inside an isolated CubeSandbox MicroVM. + +Usage from Claude Code: + python sandbox_exec.py --code "print(1+1)" + python sandbox_exec.py --file /path/to/script.py + python sandbox_exec.py --cmd "ls -la /etc" + python sandbox_exec.py --pip "requests" --code "import requests; ..." +""" + +import argparse +import os +import shlex +import sys +from e2b_code_interpreter import Sandbox +from e2b.sandbox.commands.command_handle import CommandExitException + +TEMPLATE_ID = os.getenv("CUBE_TEMPLATE_ID", "") +E2B_API_URL = os.getenv("E2B_API_URL", "http://127.0.0.1:3000") +E2B_API_KEY = os.getenv("E2B_API_KEY", "e2b_000000") + +# Cache: keep one sandbox alive for reuse within a single process. +# NOTE: Each CLI invocation (``python sandbox_exec.py ...``) is a separate +# Python process, so this cache only helps when sandbox_exec is imported +# and called repeatedly as a module (e.g. by mcp_server.py). For +# cross-process sandbox reuse — where the same sandbox survives across +# many CLI calls — use ``cubesandbox_exec.py``, which persists sandbox +# IDs to disk via a file-based session cache. +_sandbox = None + + +def _get_sandbox(timeout=300): + """Get or create a reusable sandbox.""" + global _sandbox + if _sandbox is None: + _sandbox = Sandbox.create(TEMPLATE_ID, timeout=timeout) + return _sandbox + + +def _run(sandbox, cmd, timeout=120): + """Run a command in the sandbox.""" + try: + return sandbox.commands.run(cmd, timeout=timeout) + except CommandExitException as e: + return e + + +def exec_code(code, pip_packages=None, timeout=120): + """Execute Python code in the sandbox.""" + sandbox = _get_sandbox(timeout + 60) + + if pip_packages: + r = _run(sandbox, f"pip install {' '.join(shlex.quote(p) for p in pip_packages)}", timeout=60) + if r.exit_code != 0: + return f"[pip error] {r.stderr}" + + result = _run(sandbox, f"python3 -c {shlex.quote(code)}", timeout=timeout) + if result.exit_code == 0: + return result.stdout + else: + return f"[error] {result.stderr}" + + +def exec_file(filepath, timeout=120): + """Copy a local file into the sandbox and execute it.""" + sandbox = _get_sandbox(timeout + 60) + with open(filepath, "r") as f: + content = f.read() + + sandbox.files.write("/tmp/script.py", content) + result = _run(sandbox, "python3 /tmp/script.py", timeout=timeout) + if result.exit_code == 0: + return result.stdout + else: + return f"[error] {result.stderr}" + + +def exec_cmd(command, timeout=120): + """Execute an arbitrary shell command in the sandbox.""" + sandbox = _get_sandbox(timeout + 60) + result = _run(sandbox, command, timeout=timeout) + if result.exit_code == 0: + return result.stdout + else: + return f"[error] {result.stderr}" + + +def cleanup(): + """Destroy the cached sandbox.""" + global _sandbox + if _sandbox: + _sandbox.kill() + _sandbox = None + + +def main(): + from dotenv import load_dotenv + load_dotenv() + global TEMPLATE_ID, E2B_API_URL, E2B_API_KEY + TEMPLATE_ID = os.getenv("CUBE_TEMPLATE_ID", TEMPLATE_ID) + E2B_API_URL = os.getenv("E2B_API_URL", E2B_API_URL) + E2B_API_KEY = os.getenv("E2B_API_KEY", E2B_API_KEY) + parser = argparse.ArgumentParser( + description="Execute code in a CubeSandbox sandbox" + ) + parser.add_argument("--code", help="Python code to execute") + parser.add_argument("--file", help="Local Python file to execute in sandbox") + parser.add_argument("--cmd", help="Shell command to execute") + parser.add_argument("--pip", nargs="+", help="Pip packages to install first") + parser.add_argument("--timeout", type=int, default=120, help="Execution timeout") + parser.add_argument("--keep-alive", action="store_true", + help="Keep sandbox alive after execution") + args = parser.parse_args() + + if not TEMPLATE_ID: + print("Error: CUBE_TEMPLATE_ID not set in .env") + sys.exit(1) + + try: + if args.code: + print(exec_code(args.code, args.pip, args.timeout)) + elif args.file: + print(exec_file(args.file, args.timeout)) + elif args.cmd: + print(exec_cmd(args.cmd, args.timeout)) + else: + parser.print_help() + finally: + if not args.keep_alive: + cleanup() + + +if __name__ == "__main__": + main() diff --git a/examples/sandbox-backend/test_cubesandbox_exec.py b/examples/sandbox-backend/test_cubesandbox_exec.py new file mode 100644 index 000000000..989fd0b26 --- /dev/null +++ b/examples/sandbox-backend/test_cubesandbox_exec.py @@ -0,0 +1,164 @@ +# Copyright (c) 2024 Tencent Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for cubesandbox_exec.py — session state management. + +Tests the file-based session→sandbox_id persistence without needing +a live CubeSandbox deployment. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +import importlib +import os + + +@pytest.fixture +def exec_module(tmp_state_dir): + """Import/reload cubesandbox_exec with an isolated STATE_DIR and + mock SDK symbols (so _bootstrap() is not needed).""" + import cubesandbox_exec + importlib.reload(cubesandbox_exec) + # Populate the SDK symbols that _bootstrap() would normally set, + # so state-management and get_sandbox tests work without a live SDK. + cubesandbox_exec.Sandbox = MagicMock() + cubesandbox_exec.Config = MagicMock() + cubesandbox_exec.SandboxNotFoundError = type( + "SandboxNotFoundError", (Exception,), {}) + cubesandbox_exec.CubeSandboxError = type( + "CubeSandboxError", (Exception,), {}) + cubesandbox_exec.ApiError = type("ApiError", (Exception,), {}) + cubesandbox_exec.TEMPLATE_ID = "tpl-test" + return cubesandbox_exec + + +# ── _state_path ─────────────────────────────────────────────────────── + + +class TestStatePath: + def test_alnum_session_unchanged(self, exec_module, tmp_state_dir): + p = exec_module._state_path("abc123") + assert p == tmp_state_dir / "abc123.json" + + def test_dashes_and_underscores_preserved(self, exec_module, tmp_state_dir): + p = exec_module._state_path("session-id_42") + assert p.name == "session-id_42.json" + + def test_special_chars_sanitized(self, exec_module, tmp_state_dir): + p = exec_module._state_path("user/session@host") + # / and @ → _ + assert p.name == "user_session_host.json" + + def test_empty_session_defaults(self, exec_module, tmp_state_dir): + p = exec_module._state_path("") + assert p.name == f"{exec_module.DEFAULT_SESSION}.json" + + def test_dot_dot_does_not_escape(self, exec_module, tmp_state_dir): + """Path traversal must not escape STATE_DIR.""" + p = exec_module._state_path("../../../etc/passwd") + assert p.parent == tmp_state_dir + assert ".." not in p.parts[len(tmp_state_dir.parts):] + + +# ── _save_state / _load_state round-trip ───────────────────────────── + + +class TestStateRoundTrip: + def test_save_then_load(self, exec_module, tmp_state_dir): + state = {"sandbox_id": "sb-abc", "mount": "/data/project"} + exec_module._save_state("test-session", state) + loaded = exec_module._load_state("test-session") + assert loaded == state + + def test_load_missing_returns_empty(self, exec_module, tmp_state_dir): + assert exec_module._load_state("nonexistent") == {} + + def test_load_corrupted_json_returns_empty(self, exec_module, tmp_state_dir): + # Write garbage to the state file + path = exec_module._state_path("corrupted") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("not valid json {{{") + assert exec_module._load_state("corrupted") == {} + + def test_save_creates_state_dir(self, exec_module, tmp_state_dir): + assert not tmp_state_dir.exists() + exec_module._save_session_state = exec_module._save_state # alias + exec_module._save_state("s1", {"sandbox_id": "x"}) + assert tmp_state_dir.exists() + + def test_save_overwrites(self, exec_module, tmp_state_dir): + exec_module._save_state("s1", {"sandbox_id": "old"}) + exec_module._save_state("s1", {"sandbox_id": "new"}) + assert exec_module._load_state("s1")["sandbox_id"] == "new" + + +# ── _clear_state ────────────────────────────────────────────────────── + + +class TestClearState: + def test_clear_removes_file(self, exec_module, tmp_state_dir): + exec_module._save_state("s1", {"sandbox_id": "x"}) + assert exec_module._state_path("s1").exists() + exec_module._clear_state("s1") + assert not exec_module._state_path("s1").exists() + + def test_clear_missing_is_noop(self, exec_module, tmp_state_dir): + # Must not raise + exec_module._clear_state("never-existed") + + def test_clear_then_load_returns_empty(self, exec_module, tmp_state_dir): + exec_module._save_state("s1", {"sandbox_id": "x"}) + exec_module._clear_state("s1") + assert exec_module._load_state("s1") == {} + + +# ── get_sandbox: state-driven reuse ────────────────────────────────── + + +class TestGetSandboxReuse: + def test_reuses_cached_sandbox(self, exec_module, tmp_state_dir): + """When a sandbox_id is cached, get_sandbox reconnects instead + of creating a new one.""" + exec_module._save_state("sess1", {"sandbox_id": "sb-existing", "mount": None}) + + mock_sb = MagicMock() + exec_module.Sandbox.connect.return_value = mock_sb + + sb = exec_module.get_sandbox("sess1", mount=None) + assert sb is mock_sb + exec_module.Sandbox.connect.assert_called_once() + exec_module.Sandbox.create.assert_not_called() + + def test_creates_new_when_cache_missing(self, exec_module, tmp_state_dir): + mock_sb = MagicMock() + mock_sb.sandbox_id = "sb-new" + exec_module.Sandbox.connect.side_effect = exec_module.SandboxNotFoundError() + exec_module.Sandbox.create.return_value = mock_sb + + sb = exec_module.get_sandbox("fresh-session", mount=None) + assert sb is mock_sb + exec_module.Sandbox.create.assert_called_once() + saved = exec_module._load_state("fresh-session") + assert saved["sandbox_id"] == "sb-new" + assert saved["state_token"] + + def test_recreates_when_sandbox_expired(self, exec_module, tmp_state_dir): + """If the cached sandbox is gone (SandboxNotFoundError), a new + one is created and the state is updated.""" + exec_module._save_state("sess1", {"sandbox_id": "sb-dead", "mount": None}) + + mock_sb = MagicMock() + mock_sb.sandbox_id = "sb-reborn" + exec_module.Sandbox.connect.side_effect = exec_module.SandboxNotFoundError() + exec_module.Sandbox.create.return_value = mock_sb + + sb = exec_module.get_sandbox("sess1", mount=None) + assert sb is mock_sb + saved = exec_module._load_state("sess1") + assert saved["sandbox_id"] == "sb-reborn" + assert saved["state_token"] diff --git a/examples/sandbox-backend/test_cubesandbox_rewrite.py b/examples/sandbox-backend/test_cubesandbox_rewrite.py new file mode 100644 index 000000000..741e5c749 --- /dev/null +++ b/examples/sandbox-backend/test_cubesandbox_rewrite.py @@ -0,0 +1,251 @@ +# Copyright (c) 2024 Tencent Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for cubesandbox_rewrite.py — the PreToolUse hook. + +Security-critical: verifies that the sandbox-escape vulnerability +(multi-line newline injection) and the full-path recursion issue are +fixed, and that command rewriting is correct. +""" + +from __future__ import annotations + +import io +import json +import sys + +import pytest + +import cubesandbox_rewrite as hook + +EXEC_BIN = hook.EXEC_BIN + + +# ── _already_sandboxed ──────────────────────────────────────────────── + + +class TestAlreadySandboxed: + """The _already_sandboxed() guard must be injection-proof.""" + + @pytest.mark.parametrize("cmd", [ + EXEC_BIN, + f"{EXEC_BIN} --session abc 'npm test'", + f"{EXEC_BIN} --session s --mount /tmp 'ls'", + ]) + def test_short_name_recognized(self, cmd): + assert hook._already_sandboxed(cmd) is True + + @pytest.mark.parametrize("cmd", [ + "/home/user/.local/bin/cubesandbox-exec 'npm test'", + "/usr/local/bin/cubesandbox-exec --session x 'git status'", + "./cubesandbox-exec 'echo hi'", + ]) + def test_full_path_recognized(self, cmd): + """Full-path invocations must also be detected to prevent recursion.""" + assert hook._already_sandboxed(cmd) is True + + @pytest.mark.parametrize("cmd", [ + "npm test", + "git status", + "python3 -c 'print(1+1)'", + "ls -la", + "", + " ", + ]) + def test_normal_commands_not_matched(self, cmd): + assert hook._already_sandboxed(cmd) is False + + # ── Security: newline injection ──────────────────────────────────── + + @pytest.mark.parametrize("cmd", [ + # A literal newline after the binary name must NOT match — the + # second line would execute on the host, bypassing the sandbox. + f"{EXEC_BIN}\nrm -rf /", + f"{EXEC_BIN} \nrm -rf /", + # Carriage-return variant + f"{EXEC_BIN}\rrm -rf /", + # CRLF + f"{EXEC_BIN}\r\necho pwned", + ]) + def test_newline_injection_blocked(self, cmd): + """The critical security fix: multi-line commands must NOT be + mistaken for already-sandboxed calls.""" + assert hook._already_sandboxed(cmd) is False + + def test_tab_is_not_injection(self): + """A tab is ordinary whitespace in bash — ``cubesandbox-exec\\trm`` + is a single command (cubesandbox-exec with arg 'rm'), so it IS + already sandboxed and should return True.""" + assert hook._already_sandboxed(f"{EXEC_BIN}\trm -rf /") is True + + def test_quoted_newline_in_argument(self): + """A newline *inside* a quoted argument is legitimate and should + not cause a false negative — shlex handles this correctly.""" + # shlex.split keeps the newline inside the single-quoted argument, + # so tokens[0] is still EXEC_BIN. + cmd = f"{EXEC_BIN} 'echo hello\\nworld'" + assert hook._already_sandboxed(cmd) is True + + @pytest.mark.parametrize("cmd", [ + f"{EXEC_BIN}$'\\\\n'rm -rf /", + f"{EXEC_BIN}$'\\\\r'echo pwned", + ]) + def test_ansi_c_newline_injection_blocked(self, cmd): + assert hook._already_sandboxed(cmd) is False + + def test_escaped_double_quote_keeps_newline_quoted(self): + cmd = f'{EXEC_BIN} "echo \\"quoted\\"\\ntext"' + assert hook._already_sandboxed(cmd) is True + + @pytest.mark.parametrize("cmd", [ + f"{EXEC_BIN}extra 'cmd'", # prefix match, not exact + f"not-{EXEC_BIN} 'cmd'", # different binary entirely + f"{EXEC_BIN}_backup 'cmd'", # suffix variant + ]) + def test_lookalike_names_not_matched(self, cmd): + assert hook._already_sandboxed(cmd) is False + + def test_malformed_shlex(self): + """Unbalanced quotes → shlex.error → return False (safe default).""" + assert hook._already_sandboxed(f"{EXEC_BIN} 'unbalanced") is False + + def test_empty_string(self): + assert hook._already_sandboxed("") is False + + def test_only_whitespace(self): + assert hook._already_sandboxed(" \t\n ") is False + + +# ── main() — end-to-end hook behaviour ──────────────────────────────── + + +def _run_hook(payload: dict) -> tuple[str | None, int]: + """Feed *payload* to hook.main() via stdin, return (stdout, exit_code).""" + stdin_buf = io.StringIO(json.dumps(payload)) + stdout_buf = io.StringIO() + old_stdin, old_stdout = sys.stdin, sys.stdout + sys.stdin, sys.stdout = stdin_buf, stdout_buf + try: + hook.main() + exit_code = 0 + except SystemExit as e: + exit_code = e.code if isinstance(e.code, int) else 0 + finally: + sys.stdin, sys.stdout = old_stdin, old_stdout + out = stdout_buf.getvalue() + return (out if out else None), exit_code + + +class TestMainRewrite: + def test_basic_bash_command_rewritten(self): + payload = { + "session_id": "abc123", + "cwd": "/home/user/project", + "tool_name": "Bash", + "tool_input": {"command": "npm test", "timeout": 120000}, + } + out, code = _run_hook(payload) + assert code == 0 + assert out is not None + result = json.loads(out) + assert result["hookSpecificOutput"]["permissionDecision"] == "allow" + rewritten = result["hookSpecificOutput"]["updatedInput"]["command"] + assert rewritten.startswith(f"{EXEC_BIN} ") + assert "--session" in rewritten and "abc123" in rewritten + assert "--mount" in rewritten and "/home/user/project" in rewritten + assert "--timeout" in rewritten and "120.000" in rewritten + assert "npm test" in rewritten + + def test_non_bash_tool_passes_through(self): + payload = { + "session_id": "abc123", + "tool_name": "Read", + "tool_input": {"file_path": "/etc/hosts"}, + } + out, code = _run_hook(payload) + assert code == 0 + assert out is None # hook exits silently for non-Bash tools + + def test_already_sandboxed_passes_through(self): + payload = { + "session_id": "abc123", + "tool_name": "Bash", + "tool_input": {"command": f"{EXEC_BIN} --session x 'ls'"}, + } + out, code = _run_hook(payload) + assert code == 0 + assert out is None # no rewrite — let it execute as-is + + def test_malformed_json_exits_silently(self): + stdin_buf = io.StringIO("not json {{{") + old_stdin = sys.stdin + sys.stdin = stdin_buf + try: + with pytest.raises(SystemExit) as exc_info: + hook.main() + assert exc_info.value.code == 0 + finally: + sys.stdin = old_stdin + + def test_missing_command_passes_through(self): + payload = { + "session_id": "abc123", + "tool_name": "Bash", + "tool_input": {}, + } + out, code = _run_hook(payload) + assert code == 0 + assert out is None + + def test_missing_session_defaults(self): + payload = { + "tool_name": "Bash", + "tool_input": {"command": "echo hi"}, + } + out, code = _run_hook(payload) + assert code == 0 + assert out is not None + result = json.loads(out) + rewritten = result["hookSpecificOutput"]["updatedInput"]["command"] + assert "--session" in rewritten and "default" in rewritten + + def test_no_cwd_omits_mount(self): + payload = { + "session_id": "abc123", + "tool_name": "Bash", + "tool_input": {"command": "echo hi"}, + } + out, code = _run_hook(payload) + assert code == 0 + result = json.loads(out) + rewritten = result["hookSpecificOutput"]["updatedInput"]["command"] + assert "--mount" not in rewritten + + def test_zero_timeout_omitted(self): + payload = { + "session_id": "abc123", + "tool_name": "Bash", + "tool_input": {"command": "echo hi", "timeout": 0}, + } + out, code = _run_hook(payload) + assert code == 0 + result = json.loads(out) + rewritten = result["hookSpecificOutput"]["updatedInput"]["command"] + assert "--timeout" not in rewritten + + def test_newline_injection_not_rewritten(self): + """A command that tries to inject a newline after cubesandbox-exec + must be treated as a *normal* command and wrapped (not skipped).""" + malicious = f"{EXEC_BIN}\nrm -rf /" + payload = { + "session_id": "abc123", + "tool_name": "Bash", + "tool_input": {"command": malicious}, + } + out, code = _run_hook(payload) + assert code == 0 + assert out is not None + result = json.loads(out) + rewritten = result["hookSpecificOutput"]["updatedInput"]["command"] + # The malicious command should be wrapped inside cubesandbox-exec + # as a single quoted argument, not passed through as-is. + assert rewritten.startswith(f"{EXEC_BIN} --session") diff --git a/examples/sandbox-backend/test_mcp_server.py b/examples/sandbox-backend/test_mcp_server.py new file mode 100644 index 000000000..ef6e4cea0 --- /dev/null +++ b/examples/sandbox-backend/test_mcp_server.py @@ -0,0 +1,327 @@ +# Copyright (c) 2024 Tencent Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for mcp_server.py — MCP protocol handling. + +Focuses on the fixes applied per PR #765 review: + - TTL refresh on get_sandbox() and atexit cleanup + - EOF detection in _read_mcp_message() (prevents infinite loop) + - Error logging instead of silent swallowing +""" + +from __future__ import annotations + +import io +import json +import sys +from unittest.mock import MagicMock, patch + +import pytest + +import mcp_server + + +# ── _read_mcp_message ───────────────────────────────────────────────── + + +def _make_stdin(messages: list[str]) -> io.StringIO: + """Build a stdin buffer containing one or more Content-Length framed messages.""" + buf = io.StringIO() + for msg in messages: + body = msg if isinstance(msg, str) else json.dumps(msg) + buf.write(f"Content-Length: {len(body)}\r\n\r\n{body}") + buf.seek(0) + return buf + + +class TestReadMcpMessage: + def test_valid_message(self): + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "initialize"}) + old_stdin = sys.stdin + sys.stdin = _make_stdin([body]) + try: + result = mcp_server._read_mcp_message() + assert result == {"jsonrpc": "2.0", "id": 1, "method": "initialize"} + finally: + sys.stdin = old_stdin + + def test_eof_raises_eoferror(self): + """When stdin is exhausted, _read_mcp_message must raise EOFError + so main() can break — NOT return None (which would cause an + infinite tight loop).""" + old_stdin = sys.stdin + sys.stdin = io.StringIO("") # immediate EOF + try: + with pytest.raises(EOFError): + mcp_server._read_mcp_message() + finally: + sys.stdin = old_stdin + + def test_malformed_content_length_returns_none(self, capsys): + old_stdin = sys.stdin + sys.stdin = io.StringIO("Content-Length: abc\r\n\r\n{}") + try: + result = mcp_server._read_mcp_message() + assert result is None + finally: + sys.stdin = old_stdin + captured = capsys.readouterr() + assert "malformed Content-Length" in captured.err + + def test_malformed_json_returns_none(self, capsys): + old_stdin = sys.stdin + sys.stdin = _make_stdin(["not valid json"]) + try: + result = mcp_server._read_mcp_message() + assert result is None + finally: + sys.stdin = old_stdin + captured = capsys.readouterr() + assert "malformed JSON" in captured.err + + def test_missing_content_length_returns_none(self): + old_stdin = sys.stdin + sys.stdin = io.StringIO("Some-Header: value\r\n\r\n") + try: + result = mcp_server._read_mcp_message() + assert result is None + finally: + sys.stdin = old_stdin + + def test_multiple_messages_sequential(self): + msg1 = json.dumps({"id": 1, "method": "initialize"}) + msg2 = json.dumps({"id": 2, "method": "tools/list"}) + old_stdin = sys.stdin + sys.stdin = _make_stdin([msg1, msg2]) + try: + r1 = mcp_server._read_mcp_message() + r2 = mcp_server._read_mcp_message() + assert r1["id"] == 1 + assert r2["id"] == 2 + finally: + sys.stdin = old_stdin + + +# ── handle_request ──────────────────────────────────────────────────── + + +class TestHandleRequest: + def test_initialize(self): + resp = mcp_server.handle_request({"jsonrpc": "2.0", "id": 1, "method": "initialize"}) + assert resp["id"] == 1 + assert resp["result"]["serverInfo"]["name"] == "cubesandbox-mcp" + + def test_tools_list(self): + resp = mcp_server.handle_request({"jsonrpc": "2.0", "id": 2, "method": "tools/list"}) + tool_names = [t["name"] for t in resp["result"]["tools"]] + assert "sandbox_run_command" in tool_names + assert "sandbox_run_code" in tool_names + assert "sandbox_write_file" in tool_names + assert "sandbox_read_file" in tool_names + assert "sandbox_reset" in tool_names + + def test_unknown_method(self): + resp = mcp_server.handle_request({"jsonrpc": "2.0", "id": 3, "method": "foo/bar"}) + assert "error" in resp + assert resp["error"]["code"] == -32601 + + def test_notifications_initialized_returns_none(self): + resp = mcp_server.handle_request({"method": "notifications/initialized"}) + assert resp is None + + @patch("mcp_server.get_sandbox") + def test_sandbox_run_command(self, mock_get): + mock_sb = MagicMock() + mock_result = MagicMock(exit_code=0, stdout="hello\n", stderr="") + mock_sb.commands.run.return_value = mock_result + mock_get.return_value = mock_sb + + resp = mcp_server.handle_request({ + "jsonrpc": "2.0", "id": 10, + "method": "tools/call", + "params": {"name": "sandbox_run_command", "arguments": {"command": "echo hello"}}, + }) + text = resp["result"]["content"][0]["text"] + assert "hello" in text + + @patch("mcp_server.get_sandbox") + def test_sandbox_run_code(self, mock_get): + mock_sb = MagicMock() + mock_result = MagicMock(exit_code=0, stdout="42\n", stderr="") + mock_sb.commands.run.return_value = mock_result + mock_get.return_value = mock_sb + + resp = mcp_server.handle_request({ + "jsonrpc": "2.0", "id": 11, + "method": "tools/call", + "params": {"name": "sandbox_run_code", "arguments": {"code": "print(1+1)"}}, + }) + text = resp["result"]["content"][0]["text"] + assert "42" in text + + @patch("mcp_server.get_sandbox") + def test_sandbox_write_file(self, mock_get): + mock_sb = MagicMock() + mock_get.return_value = mock_sb + + resp = mcp_server.handle_request({ + "jsonrpc": "2.0", "id": 12, + "method": "tools/call", + "params": {"name": "sandbox_write_file", + "arguments": {"path": "/tmp/x.py", "content": "print(1)"}}, + }) + text = resp["result"]["content"][0]["text"] + assert "Written" in text + mock_sb.files.write.assert_called_once_with("/tmp/x.py", "print(1)") + + @patch("mcp_server.get_sandbox") + def test_sandbox_read_file(self, mock_get): + mock_sb = MagicMock() + mock_sb.files.read.return_value = "file contents" + mock_get.return_value = mock_sb + + resp = mcp_server.handle_request({ + "jsonrpc": "2.0", "id": 13, + "method": "tools/call", + "params": {"name": "sandbox_read_file", "arguments": {"path": "/tmp/x.py"}}, + }) + text = resp["result"]["content"][0]["text"] + assert text == "file contents" + + @patch("mcp_server.get_sandbox") + def test_sandbox_reset(self, mock_get): + mock_sb = MagicMock() + mcp_server._sandbox = mock_sb # pretend a sandbox exists + + resp = mcp_server.handle_request({ + "jsonrpc": "2.0", "id": 14, + "method": "tools/call", + "params": {"name": "sandbox_reset", "arguments": {}}, + }) + text = resp["result"]["content"][0]["text"] + assert "destroyed" in text.lower() + mock_sb.kill.assert_called_once() + assert mcp_server._sandbox is None + + @patch("mcp_server.get_sandbox") + def test_unknown_tool_returns_error_text(self, mock_get): + resp = mcp_server.handle_request({ + "jsonrpc": "2.0", "id": 15, + "method": "tools/call", + "params": {"name": "nonexistent", "arguments": {}}, + }) + text = resp["result"]["content"][0]["text"] + assert "Unknown tool" in text + + @patch("mcp_server.get_sandbox") + def test_exception_logged(self, mock_get, capsys): + """handle_request must log the traceback, not silently swallow it.""" + mock_get.side_effect = RuntimeError("boom") + resp = mcp_server.handle_request({ + "jsonrpc": "2.0", "id": 16, + "method": "tools/call", + "params": {"name": "sandbox_run_command", "arguments": {"command": "ls"}}, + }) + text = resp["result"]["content"][0]["text"] + assert "Error" in text + captured = capsys.readouterr() + assert "Traceback" in captured.err + + +# ── get_sandbox / _cleanup_sandbox ──────────────────────────────────── + + +class TestSandboxLifecycle: + def setup_method(self): + mcp_server._sandbox = None + + def teardown_method(self): + mcp_server._sandbox = None + + @patch("mcp_server.atexit") + @patch("e2b_code_interpreter.Sandbox") + def test_creates_sandbox_and_registers_cleanup(self, mock_sb_cls, mock_atexit): + mock_sb = MagicMock() + mock_sb_cls.create.return_value = mock_sb + + sb = mcp_server.get_sandbox(timeout=600) + assert sb is mock_sb + mock_sb_cls.create.assert_called_once() + # atexit.register must be called so the sandbox is killed on exit + mock_atexit.register.assert_called_once_with(mcp_server._cleanup_sandbox) + + @patch("mcp_server.atexit") + @patch("e2b_code_interpreter.Sandbox") + def test_ttl_refreshed_on_reuse(self, mock_sb_cls, mock_atexit): + mock_sb = MagicMock() + mock_sb_cls.create.return_value = mock_sb + + # First call creates the sandbox + mcp_server.get_sandbox(timeout=600) + # Second call should refresh TTL, not create a new one + mcp_server.get_sandbox(timeout=900) + mock_sb.set_timeout.assert_called_once_with(900) + assert mock_sb_cls.create.call_count == 1 + + @patch("mcp_server.atexit") + @patch("e2b_code_interpreter.Sandbox") + def test_recreates_on_ttl_refresh_failure(self, mock_sb_cls, mock_atexit): + mock_sb = MagicMock() + mock_sb.set_timeout.side_effect = RuntimeError("expired") + mock_sb_cls.create.return_value = mock_sb + + mcp_server.get_sandbox(timeout=600) + assert mock_sb_cls.create.call_count == 1 + # Second call: set_timeout fails → create a new sandbox + mcp_server.get_sandbox(timeout=600) + assert mock_sb_cls.create.call_count == 2 + + def test_cleanup_kills_sandbox(self): + mock_sb = MagicMock() + mcp_server._sandbox = mock_sb + mcp_server._cleanup_sandbox() + mock_sb.kill.assert_called_once() + assert mcp_server._sandbox is None + + def test_cleanup_noop_when_no_sandbox(self): + mcp_server._sandbox = None + mcp_server._cleanup_sandbox() # must not raise + assert mcp_server._sandbox is None + + def test_cleanup_swallows_kill_errors(self): + mock_sb = MagicMock() + mock_sb.kill.side_effect = RuntimeError("already dead") + mcp_server._sandbox = mock_sb + mcp_server._cleanup_sandbox() # must not raise + assert mcp_server._sandbox is None + + +# ── main() integration ──────────────────────────────────────────────── + + +class TestMain: + def test_eof_exits_cleanly(self): + """main() must break on EOF, not spin forever.""" + old_stdin, old_stdout = sys.stdin, sys.stdout + sys.stdin = io.StringIO("") # immediate EOF + sys.stdout = io.StringIO() + try: + mcp_server.main() # should return normally + finally: + sys.stdin, sys.stdout = old_stdin, old_stdout + + @patch("mcp_server.handle_request") + def test_valid_request_produces_response(self, mock_handle): + mock_handle.return_value = {"jsonrpc": "2.0", "id": 1, "result": {"ok": True}} + body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "initialize"}) + old_stdin, old_stdout = sys.stdin, sys.stdout + sys.stdin = _make_stdin([body]) + out_buf = io.StringIO() + sys.stdout = out_buf + try: + mcp_server.main() + finally: + sys.stdin, sys.stdout = old_stdin, old_stdout + output = out_buf.getvalue() + assert "Content-Length:" in output + parsed = json.loads(output.split("\r\n\r\n", 1)[1]) + assert parsed["result"]["ok"] is True diff --git a/examples/sandbox-backend/test_sandbox_exec.py b/examples/sandbox-backend/test_sandbox_exec.py new file mode 100644 index 000000000..6cdb7add6 --- /dev/null +++ b/examples/sandbox-backend/test_sandbox_exec.py @@ -0,0 +1,222 @@ +# Copyright (c) 2024 Tencent Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for sandbox_exec.py — the standalone CLI. + +Tests argument parsing, the _run wrapper, and exec_* functions +with a mock sandbox (no live deployment needed). +""" + +from __future__ import annotations + +import sys +from unittest.mock import MagicMock, patch + +import pytest + +import sandbox_exec + + +# ── _run wrapper ────────────────────────────────────────────────────── + + +class TestRunWrapper: + def test_run_returns_result_on_success(self): + sb = MagicMock() + expected = MagicMock(exit_code=0, stdout="ok", stderr="") + sb.commands.run.return_value = expected + result = sandbox_exec._run(sb, "echo hi", timeout=10) + assert result is expected + sb.commands.run.assert_called_once_with("echo hi", timeout=10) + + def test_run_returns_exception_on_nonzero_exit(self): + from e2b.sandbox.commands.command_handle import CommandExitException + sb = MagicMock() + exc = CommandExitException(stderr="err", stdout="", exit_code=1, error=None) + sb.commands.run.side_effect = exc + result = sandbox_exec._run(sb, "false", timeout=10) + assert result is exc + + +# ── exec_code ───────────────────────────────────────────────────────── + + +class TestExecCode: + @patch("sandbox_exec._get_sandbox") + def test_exec_code_success(self, mock_get): + sb = MagicMock() + mock_result = MagicMock(exit_code=0, stdout="42\n", stderr="") + sb.commands.run.return_value = mock_result + mock_get.return_value = sb + + out = sandbox_exec.exec_code("print(1+1)") + assert "42" in out + + @patch("sandbox_exec._get_sandbox") + def test_exec_code_failure(self, mock_get): + sb = MagicMock() + mock_result = MagicMock(exit_code=1, stdout="", stderr="SyntaxError") + sb.commands.run.return_value = mock_result + mock_get.return_value = sb + + out = sandbox_exec.exec_code("invalid syntax") + assert "[error]" in out + + @patch("sandbox_exec._get_sandbox") + def test_exec_code_with_pip(self, mock_get): + sb = MagicMock() + pip_result = MagicMock(exit_code=0, stdout="", stderr="") + code_result = MagicMock(exit_code=0, stdout="ok\n", stderr="") + sb.commands.run.side_effect = [pip_result, code_result] + mock_get.return_value = sb + + out = sandbox_exec.exec_code("import requests", pip_packages=["requests"]) + assert "ok" in out + # First call: pip install, second: python3 -c + assert sb.commands.run.call_count == 2 + + @patch("sandbox_exec._get_sandbox") + def test_exec_code_pip_failure(self, mock_get): + sb = MagicMock() + pip_result = MagicMock(exit_code=1, stdout="", stderr="pip error") + sb.commands.run.return_value = pip_result + mock_get.return_value = sb + + out = sandbox_exec.exec_code("import x", pip_packages=["nonexistent-pkg"]) + assert "[pip error]" in out + + +# ── exec_cmd ────────────────────────────────────────────────────────── + + +class TestExecCmd: + @patch("sandbox_exec._get_sandbox") + def test_exec_cmd_success(self, mock_get): + sb = MagicMock() + sb.commands.run.return_value = MagicMock(exit_code=0, stdout="file1\nfile2\n", stderr="") + mock_get.return_value = sb + + out = sandbox_exec.exec_cmd("ls") + assert "file1" in out + + @patch("sandbox_exec._get_sandbox") + def test_exec_cmd_failure(self, mock_get): + sb = MagicMock() + sb.commands.run.return_value = MagicMock(exit_code=127, stdout="", stderr="command not found") + mock_get.return_value = sb + + out = sandbox_exec.exec_cmd("nonexistent-cmd") + assert "[error]" in out + + +# ── exec_file ───────────────────────────────────────────────────────── + + +class TestExecFile: + @patch("sandbox_exec._get_sandbox") + def test_exec_file_success(self, mock_get, tmp_path): + script = tmp_path / "test.py" + script.write_text("print('from file')") + + sb = MagicMock() + sb.commands.run.return_value = MagicMock(exit_code=0, stdout="from file\n", stderr="") + mock_get.return_value = sb + + out = sandbox_exec.exec_file(str(script)) + assert "from file" in out + # File should be written into the sandbox + sb.files.write.assert_called_once() + written_path = sb.files.write.call_args[0][0] + assert written_path == "/tmp/script.py" + + +# ── cleanup ─────────────────────────────────────────────────────────── + + +class TestCleanup: + def test_cleanup_kills_sandbox(self): + mock_sb = MagicMock() + sandbox_exec._sandbox = mock_sb + sandbox_exec.cleanup() + mock_sb.kill.assert_called_once() + assert sandbox_exec._sandbox is None + + def test_cleanup_noop_when_none(self): + sandbox_exec._sandbox = None + sandbox_exec.cleanup() # must not raise + assert sandbox_exec._sandbox is None + + +# ── _get_sandbox ────────────────────────────────────────────────────── + + +class TestGetSandbox: + def setup_method(self): + sandbox_exec._sandbox = None + + def teardown_method(self): + sandbox_exec._sandbox = None + + @patch("sandbox_exec.Sandbox") + def test_creates_sandbox_on_first_call(self, mock_sb_cls): + mock_sb = MagicMock() + mock_sb_cls.create.return_value = mock_sb + sb = sandbox_exec._get_sandbox(timeout=300) + assert sb is mock_sb + mock_sb_cls.create.assert_called_once() + + @patch("sandbox_exec.Sandbox") + def test_reuses_sandbox_on_second_call(self, mock_sb_cls): + mock_sb = MagicMock() + mock_sb_cls.create.return_value = mock_sb + sandbox_exec._get_sandbox(timeout=300) + sandbox_exec._get_sandbox(timeout=300) + assert mock_sb_cls.create.call_count == 1 + + +# ── CLI argument parsing ────────────────────────────────────────────── + + +class TestArgParsing: + @patch("sandbox_exec.exec_code") + @patch("sandbox_exec.cleanup") + @patch("sandbox_exec.TEMPLATE_ID", "tpl-test") + def test_code_arg(self, mock_cleanup, mock_exec): + mock_exec.return_value = "result" + sys.argv = ["sandbox_exec.py", "--code", "print(1)"] + sandbox_exec.main() + mock_exec.assert_called_once() + + @patch("sandbox_exec.exec_cmd") + @patch("sandbox_exec.cleanup") + @patch("sandbox_exec.TEMPLATE_ID", "tpl-test") + def test_cmd_arg(self, mock_cleanup, mock_exec): + mock_exec.return_value = "result" + sys.argv = ["sandbox_exec.py", "--cmd", "ls -la"] + sandbox_exec.main() + mock_exec.assert_called_once() + + @patch("sandbox_exec.exec_file") + @patch("sandbox_exec.cleanup") + @patch("sandbox_exec.TEMPLATE_ID", "tpl-test") + def test_file_arg(self, mock_cleanup, mock_exec): + mock_exec.return_value = "result" + sys.argv = ["sandbox_exec.py", "--file", "/tmp/test.py"] + sandbox_exec.main() + mock_exec.assert_called_once() + + @patch("sandbox_exec.TEMPLATE_ID", "tpl-test") + def test_keep_alive_skips_cleanup(self): + """--keep-alive must skip cleanup() so the sandbox survives.""" + sys.argv = ["sandbox_exec.py", "--code", "print(1)", "--keep-alive"] + with patch("sandbox_exec.exec_code", return_value="ok") as mock_exec, \ + patch("sandbox_exec.cleanup") as mock_cleanup: + sandbox_exec.main() + mock_exec.assert_called_once() + mock_cleanup.assert_not_called() + + @patch("sandbox_exec.TEMPLATE_ID", "") + def test_missing_template_id_exits(self): + sys.argv = ["sandbox_exec.py", "--code", "print(1)"] + with pytest.raises(SystemExit) as exc_info: + sandbox_exec.main() + assert exc_info.value.code == 1