Coverage for agentos/security/sandbox_executor.py: 25%
265 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 21:19 +0800
1"""
2AgentOS v1.2.1 — 沙箱执行器。
4基因来源: OpenHands Docker Sandbox + Claude Code subprocess isolation
6提供真正的代码/命令隔离执行能力：
7- Process模式: 子进程隔离（轻量，零依赖）
8- Docker模式: 容器隔离（强隔离，需Docker）
9- 资源限制：内存、CPU、时间、磁盘
10- 文件桥接：自动复制输入文件到沙箱，提取输出文件
11- 与 CodeAgent / ToolOrchestrator 集成
12"""
14from __future__ import annotations
16import asyncio
17import os
18import shutil
19import subprocess
20import tempfile
21import time
22import uuid
23from dataclasses import dataclass, field
24from enum import StrEnum
26# ── 枚举与配置 ───────────────────────────────────
29class SandboxMode(StrEnum):
30 """沙箱模式枚举。"""
32 DOCKER = "docker"
33 PROCESS = "process"
34 NONE = "none" # 直接在当前进程执行（不安全，仅调试用）
37@dataclass
38class SandboxConfig:
39 """沙箱执行配置"""
41 mode: SandboxMode = SandboxMode.PROCESS
42 memory_limit_mb: int = 256
43 cpu_limit: float = 1.0 # CPU 核心数上限
44 timeout_seconds: float = 30.0
45 max_output_bytes: int = 1_000_000 # stdout+stderr 上限
46 network_enabled: bool = False
47 writable_root: bool = False # root 是否可写（Docker模式）
48 docker_image: str = "python:3.11-slim"
49 container_name_prefix: str = "agentos-sandbox-"
50 env_vars: dict[str, str] = field(default_factory=dict)
53# ── 执行结果 ────────────────────────────────────
56@dataclass
57class SandboxResult:
58 """沙箱执行结果"""
60 success: bool
61 exit_code: int = 0
62 stdout: str = ""
63 stderr: str = ""
64 output_files: dict[str, str] = field(default_factory=dict) # 文件名→本地路径
65 duration_ms: float = 0.0
66 truncated: bool = False
67 error: str | None = None
70# ── Process 沙箱 ────────────────────────────────
73class ProcessSandbox:
74 """进程级隔离沙箱。使用 subprocess + 临时目录隔离文件系统。"""
76 def __init__(self, config: SandboxConfig | None = None):
77 self.config = config or SandboxConfig()
78 self._work_dir: str | None = None
80 def setup(self) -> str:
81 """创建隔离的工作目录。返回沙箱目录路径。"""
82 self._work_dir = tempfile.mkdtemp(prefix="agentos-sandbox-")
83 return self._work_dir
85 def copy_in(self, src: str, dst_filename: str | None = None) -> str:
86 """将外部文件复制到沙箱内。返回沙箱内路径。"""
87 if not self._work_dir:
88 self.setup()
89 fname = dst_filename or os.path.basename(src)
90 dst = os.path.join(self._work_dir, fname)
91 if os.path.isfile(src):
92 shutil.copy2(src, dst)
93 elif os.path.isdir(src):
94 shutil.copytree(src, dst, dirs_exist_ok=True)
95 return dst
97 def copy_out(self, sandbox_path: str, local_path: str) -> str:
98 """将沙箱内文件复制到外部。"""
99 os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True)
100 if os.path.isfile(sandbox_path):
101 shutil.copy2(sandbox_path, local_path)
102 return local_path
104 def collect_output_files(self, patterns: list[str] | None = None) -> dict[str, str]:
105 """收集沙箱内生成的文件（按扩展名匹配），复制到本地临时目录。
107 Args:
108 patterns: 文件扩展名或glob模式列表，如 ['.json', '.csv', '.png']。
109 None 则收集所有非目录文件。
111 Returns:
112 {沙箱内文件名: 本地临时路径}
113 """
114 if not self._work_dir:
115 return {}
116 output_dir = tempfile.mkdtemp(prefix="agentos-output-")
117 result: dict[str, str] = {}
118 for root, dirs, files in os.walk(self._work_dir):
119 for fname in files:
120 match = True
121 if patterns:
122 match = any(fname.endswith(p) or fname == p for p in patterns)
123 if match:
124 src = os.path.join(root, fname)
125 rel = os.path.relpath(src, self._work_dir)
126 dst = os.path.join(output_dir, rel)
127 os.makedirs(os.path.dirname(dst), exist_ok=True)
128 shutil.copy2(src, dst)
129 result[fname] = dst
130 return result
132 def execute_code(
133 self,
134 code: str,
135 language: str = "python",
136 input_files: dict[str, str] | None = None,
137 ) -> SandboxResult:
138 """在沙箱中执行代码。
140 Args:
141 code: 代码字符串
142 language: python | bash
143 input_files: {文件名: 外部路径} 输入文件映射
144 """
145 time.monotonic()
146 if not self._work_dir:
147 self.setup()
149 # 复制输入文件
150 if input_files:
151 for fname, src_path in input_files.items():
152 self.copy_in(src_path, fname)
154 if language == "python":
155 script_path = os.path.join(self._work_dir, "_sandbox_script.py")
156 with open(script_path, "w") as f:
157 f.write(code)
158 cmd = [self._get_python(), script_path]
159 elif language == "bash":
160 script_path = os.path.join(self._work_dir, "_sandbox_script.sh")
161 with open(script_path, "w") as f:
162 f.write("#!/bin/bash\nset -e\n" + code)
163 os.chmod(script_path, 0o755)
164 cmd = ["bash", script_path]
165 else:
166 return SandboxResult(success=False, error=f"Unsupported language: {language}")
168 return self._run_subprocess(cmd)
170 def execute_command(self, command: str | list[str]) -> SandboxResult:
171 """在沙箱中执行命令。"""
172 time.monotonic()
173 if not self._work_dir:
174 self.setup()
175 if isinstance(command, str):
176 cmd = ["bash", "-c", command]
177 else:
178 cmd = list(command)
179 return self._run_subprocess(cmd)
181 def _run_subprocess(self, cmd: list[str]) -> SandboxResult:
182 start = time.monotonic()
183 env = os.environ.copy()
184 env.update(self.config.env_vars)
185 # 网络隔离
186 if not self.config.network_enabled:
187 env["http_proxy"] = ""
188 env["https_proxy"] = ""
189 env["HTTP_PROXY"] = ""
190 env["HTTPS_PROXY"] = ""
192 try:
193 proc = subprocess.run(
194 cmd,
195 cwd=self._work_dir,
196 env=env,
197 capture_output=True,
198 timeout=self.config.timeout_seconds,
199 text=True,
200 )
201 stdout = proc.stdout or ""
202 stderr = proc.stderr or ""
203 truncated = False
205 if len(stdout) > self.config.max_output_bytes:
206 stdout = stdout[: self.config.max_output_bytes] + "\n... [stdout truncated]"
207 truncated = True
208 if len(stderr) > self.config.max_output_bytes:
209 stderr = stderr[: self.config.max_output_bytes] + "\n... [stderr truncated]"
210 truncated = True
212 duration = (time.monotonic() - start) * 1000
213 return SandboxResult(
214 success=(proc.returncode == 0),
215 exit_code=proc.returncode,
216 stdout=stdout,
217 stderr=stderr,
218 duration_ms=duration,
219 truncated=truncated,
220 )
221 except subprocess.TimeoutExpired as e:
222 duration = (time.monotonic() - start) * 1000
223 return SandboxResult(
224 success=False,
225 exit_code=-1,
226 stdout=e.stdout or "" if e.stdout else "",
227 stderr=(
228 e.stderr or "Timeout: execution exceeded limit"
229 if e.stderr
230 else "Timeout: execution exceeded limit"
231 ),
232 duration_ms=duration,
233 error=f"Timeout after {self.config.timeout_seconds}s",
234 )
235 except Exception as e:
236 duration = (time.monotonic() - start) * 1000
237 return SandboxResult(success=False, duration_ms=duration, error=str(e))
239 @staticmethod
240 def _get_python() -> str:
241 return shutil.which("python3") or shutil.which("python") or "python3"
243 def cleanup(self):
244 if self._work_dir and os.path.isdir(self._work_dir):
245 shutil.rmtree(self._work_dir, ignore_errors=True)
246 self._work_dir = None
249# ── Docker 沙箱 ────────────────────────────────
252class DockerSandbox:
253 """Docker 容器隔离沙箱。更强的隔离性和可重现性。"""
255 def __init__(self, config: SandboxConfig | None = None):
256 self.config = config or SandboxConfig(mode=SandboxMode.DOCKER)
257 self._container_id: str | None = None
258 self._host_work_dir: str | None = None
260 def setup(self) -> str:
261 """创建并启动 Docker 容器。返回沙箱目录路径。"""
262 self._host_work_dir = tempfile.mkdtemp(prefix="agentos-docker-")
263 container_name = f"{self.config.container_name_prefix}{uuid.uuid4().hex[:8]}"
265 cmd = [
266 "docker",
267 "run",
268 "-d",
269 "--rm",
270 "--name",
271 container_name,
272 f"--memory={self.config.memory_limit_mb}m",
273 f"--cpus={self.config.cpu_limit}",
274 "-v",
275 f"{self._host_work_dir}:/workspace",
276 "-w",
277 "/workspace",
278 ]
279 if not self.config.network_enabled:
280 cmd.append("--network=none")
281 if self.config.writable_root:
282 cmd.append("--read-only=false")
283 else:
284 cmd.append("--read-only")
285 cmd.append("--tmpfs=/tmp:exec")
287 cmd.extend(["sleep", "infinity"])
288 cmd.append(self.config.docker_image)
290 result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
291 if result.returncode != 0:
292 raise RuntimeError(f"Docker setup failed: {result.stderr}")
294 self._container_id = result.stdout.strip()[:12]
295 return self._host_work_dir
297 def copy_in(self, src: str, dst_filename: str | None = None) -> str:
298 if not self._container_id:
299 self.setup()
300 fname = dst_filename or os.path.basename(src)
301 subprocess.run(
302 ["docker", "cp", src, f"{self._container_id}:/workspace/{fname}"],
303 check=True,
304 capture_output=True,
305 timeout=10,
306 )
307 return os.path.join(self._host_work_dir, fname)
309 def execute_code(
310 self,
311 code: str,
312 language: str = "python",
313 input_files: dict[str, str] | None = None,
314 ) -> SandboxResult:
315 if not self._container_id:
316 self.setup()
318 if input_files:
319 for fname, src_path in input_files.items():
320 self.copy_in(src_path, fname)
322 if language == "python":
323 script = "_sandbox_script.py"
324 script_path = os.path.join(self._host_work_dir, script)
325 with open(script_path, "w") as f:
326 f.write(code)
327 cmd = ["docker", "exec", self._container_id, "python3", f"/workspace/{script}"]
328 elif language == "bash":
329 script = "_sandbox_script.sh"
330 script_path = os.path.join(self._host_work_dir, script)
331 with open(script_path, "w") as f:
332 f.write("#!/bin/bash\nset -e\n" + code)
333 subprocess.run(["chmod", "+x", script_path], check=False)
334 cmd = ["docker", "exec", self._container_id, "bash", f"/workspace/{script}"]
335 else:
336 return SandboxResult(success=False, error=f"Unsupported language: {language}")
338 return self._run_docker(cmd)
340 def execute_command(self, command: str | list[str]) -> SandboxResult:
341 if not self._container_id:
342 self.setup()
343 if isinstance(command, str):
344 cmd = ["docker", "exec", self._container_id, "bash", "-c", command]
345 else:
346 cmd = ["docker", "exec", self._container_id] + list(command)
347 return self._run_docker(cmd)
349 def _run_docker(self, cmd: list[str]) -> SandboxResult:
350 start = time.monotonic()
351 try:
352 proc = subprocess.run(
353 cmd,
354 capture_output=True,
355 text=True,
356 timeout=self.config.timeout_seconds,
357 )
358 duration = (time.monotonic() - start) * 1000
359 stdout = proc.stdout or ""
360 stderr = proc.stderr or ""
361 truncated = False
363 if len(stdout) > self.config.max_output_bytes:
364 stdout = stdout[: self.config.max_output_bytes] + "\n... [truncated]"
365 truncated = True
367 return SandboxResult(
368 success=(proc.returncode == 0),
369 exit_code=proc.returncode,
370 stdout=stdout,
371 stderr=stderr,
372 duration_ms=duration,
373 truncated=truncated,
374 )
375 except subprocess.TimeoutExpired:
376 if self._container_id:
377 subprocess.run(["docker", "kill", self._container_id], capture_output=True)
378 return SandboxResult(
379 success=False,
380 exit_code=-1,
381 error=f"Timeout after {self.config.timeout_seconds}s",
382 duration_ms=(time.monotonic() - start) * 1000,
383 )
384 except Exception as e:
385 return SandboxResult(
386 success=False,
387 error=str(e),
388 duration_ms=(time.monotonic() - start) * 1000,
389 )
391 def collect_output_files(self, patterns: list[str] | None = None) -> dict[str, str]:
392 if not self._host_work_dir:
393 return {}
394 output_dir = tempfile.mkdtemp(prefix="agentos-output-")
395 result: dict[str, str] = {}
396 for root, dirs, files in os.walk(self._host_work_dir):
397 for fname in files:
398 if fname.startswith("_sandbox"):
399 continue
400 match = True
401 if patterns:
402 match = any(fname.endswith(p) for p in patterns)
403 if match:
404 src = os.path.join(root, fname)
405 dst = os.path.join(output_dir, fname)
406 shutil.copy2(src, dst)
407 result[fname] = dst
408 return result
410 def cleanup(self):
411 if self._container_id:
412 subprocess.run(["docker", "stop", self._container_id], capture_output=True, timeout=5)
413 self._container_id = None
414 if self._host_work_dir and os.path.isdir(self._host_work_dir):
415 shutil.rmtree(self._host_work_dir, ignore_errors=True)
416 self._host_work_dir = None
419# ── 统一沙箱执行器 ─────────────────────────────
422class SandboxExecutor:
423 """统一沙箱执行器。根据 SandboxConfig.mode 自动选择 Process/Docker。"""
425 def __init__(self, config: SandboxConfig | None = None):
426 self.config = config or SandboxConfig()
427 if self.config.mode == SandboxMode.DOCKER:
428 try:
429 self._sandbox: ProcessSandbox | DockerSandbox = DockerSandbox(self.config)
430 self._sandbox.setup()
431 except Exception:
432 # Docker 不可用时降级到 Process
433 self.config.mode = SandboxMode.PROCESS
434 self._sandbox = ProcessSandbox(self.config)
435 else:
436 self._sandbox = ProcessSandbox(self.config)
438 async def execute_code(
439 self,
440 code: str,
441 language: str = "python",
442 input_files: dict[str, str] | None = None,
443 ) -> SandboxResult:
444 loop = asyncio.get_event_loop()
445 return await loop.run_in_executor(
446 None,
447 self._sandbox.execute_code,
448 code,
449 language,
450 input_files,
451 )
453 async def execute_command(self, command: str | list[str]) -> SandboxResult:
454 loop = asyncio.get_event_loop()
455 return await loop.run_in_executor(None, self._sandbox.execute_command, command)
457 def collect_output_files(self, patterns: list[str] | None = None) -> dict[str, str]:
458 return self._sandbox.collect_output_files(patterns)
460 async def cleanup(self):
461 loop = asyncio.get_event_loop()
462 await loop.run_in_executor(None, self._sandbox.cleanup)
464 def __enter__(self):
465 self._sandbox.setup()
466 return self
468 def __exit__(self, *args):
469 self._sandbox.cleanup()