Coverage for agentos/security/sandbox.py: 44%

210 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-08 01:44 +0800

1""" 

2AgentOS Sandbox — Secure Code Execution Sandbox 

3━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 

4 

5Production-grade code execution sandbox with: 

6 

7 - Resource limits (CPU time, memory, disk, network) 

8 - Timeout enforcement 

9 - Process isolation (subprocess-based) 

10 - Dangerous import/module blacklisting 

11 - Output size limits 

12 - Audit logging for every execution 

13 

14Architecture: 

15 SandboxPolicy → defines resource limits and restrictions 

16 SandboxExecutor → executes code in isolated subprocess 

17 SandboxResult → immutable execution result 

18""" 

19 

20from __future__ import annotations 

21 

22import builtins 

23import os 

24import resource 

25import subprocess 

26import sys 

27import tempfile 

28import time 

29import traceback 

30from dataclasses import dataclass, field 

31from enum import StrEnum 

32from typing import Any 

33 

34# --------------------------------------------------------------------------- 

35# Sandbox Policy 

36# --------------------------------------------------------------------------- 

37 

38 

39class SecurityLevel(StrEnum): 

40 """Security level determines which restrictions apply.""" 

41 

42 RESTRICTED = "restricted" # Max restrictions, no network/disk 

43 STANDARD = "standard" # Reasonable limits for most use cases 

44 RELAXED = "relaxed" # Minimal restrictions, trusted code only 

45 

46 

47# Dangerous modules and builtins to block 

48DANGEROUS_MODULES: set[str] = { 

49 "os", 

50 "subprocess", 

51 "shlex", 

52 "sys", 

53 "ctypes", 

54 "socket", 

55 "http", 

56 "requests", 

57 "urllib", 

58 "multiprocessing", 

59 "threading", 

60 "concurrent.futures", 

61 "importlib", 

62 "pkgutil", 

63 "pkg_resources", 

64 "pickle", 

65 "marshal", 

66 "shelve", 

67 "pathlib", 

68 "shutil", 

69 "glob", 

70 "fnmatch", 

71 "signal", 

72 "atexit", 

73 "code", 

74 "codeop", 

75 "compileall", 

76 "pty", 

77 "fcntl", 

78 "tty", 

79 "termios", 

80 "webbrowser", 

81} 

82 

83DANGEROUS_BUILTINS: set[str] = { 

84 "__import__", 

85 "compile", 

86 "eval", 

87 "exec", 

88 "open", 

89 "input", 

90 "breakpoint", 

91 "globals", 

92 "locals", 

93 "vars", 

94 "getattr", 

95 "setattr", 

96 "delattr", 

97 "type", 

98 "issubclass", 

99 "isinstance", 

100 "memoryview", 

101 "bytearray", 

102} 

103 

104 

105@dataclass 

106class SandboxPolicy: 

107 """Resource limits and restrictions for sandboxed execution.""" 

108 

109 # Resource limits 

110 max_cpu_time_seconds: float = 30.0 

111 max_memory_mb: int = 256 

112 max_output_size_bytes: int = 10 * 1024 * 1024 # 10 MB 

113 max_files_open: int = 50 

114 max_processes: int = 1 

115 network_enabled: bool = False 

116 disk_write_enabled: bool = False 

117 

118 # Execution constraints 

119 timeout_seconds: float = 60.0 

120 max_iterations: int = 10_000_000 # Safety net for infinite loops 

121 

122 # Security 

123 block_dangerous_modules: bool = True 

124 block_dangerous_builtins: bool = True 

125 allowed_modules: set[str] = field(default_factory=set) 

126 allowed_imports: set[str] = field(default_factory=set) 

127 

128 # Audit 

129 enable_audit: bool = True 

130 

131 @classmethod 

132 def for_level(cls, level: SecurityLevel) -> SandboxPolicy: 

133 """Create a policy for a given security level.""" 

134 if level == SecurityLevel.RESTRICTED: 

135 return cls( 

136 max_cpu_time_seconds=5.0, 

137 max_memory_mb=64, 

138 max_output_size_bytes=1 * 1024 * 1024, 

139 network_enabled=False, 

140 disk_write_enabled=False, 

141 timeout_seconds=10.0, 

142 allowed_modules={ 

143 "math", 

144 "json", 

145 "datetime", 

146 "collections", 

147 "itertools", 

148 "functools", 

149 }, 

150 ) 

151 elif level == SecurityLevel.RELAXED: 

152 return cls( 

153 max_cpu_time_seconds=120.0, 

154 max_memory_mb=1024, 

155 max_output_size_bytes=100 * 1024 * 1024, 

156 network_enabled=True, 

157 disk_write_enabled=True, 

158 timeout_seconds=300.0, 

159 block_dangerous_modules=False, 

160 ) 

161 else: # STANDARD 

162 return cls() 

163 

164 

165# --------------------------------------------------------------------------- 

166# Sandbox Result 

167# --------------------------------------------------------------------------- 

168 

169 

170@dataclass(frozen=True) 

171class SandboxResult: 

172 """Immutable result of a sandboxed code execution.""" 

173 

174 exit_code: int 

175 stdout: str 

176 stderr: str 

177 execution_time_ms: float 

178 peak_memory_mb: float 

179 was_killed: bool = False 

180 kill_reason: str = "" 

181 error: str | None = None 

182 timestamp: float = field(default_factory=time.time) 

183 

184 @property 

185 def success(self) -> bool: 

186 return self.exit_code == 0 and not self.was_killed 

187 

188 def to_dict(self) -> dict[str, Any]: 

189 return { 

190 "exit_code": self.exit_code, 

191 "stdout": self.stdout[:1000], 

192 "stderr": self.stderr[:1000], 

193 "execution_time_ms": self.execution_time_ms, 

194 "peak_memory_mb": self.peak_memory_mb, 

195 "was_killed": self.was_killed, 

196 "kill_reason": self.kill_reason, 

197 "success": self.success, 

198 } 

199 

200 

201# --------------------------------------------------------------------------- 

202# Audit Logger 

203# --------------------------------------------------------------------------- 

204 

205 

206class SandboxAuditLogger: 

207 """Logs sandbox execution events for security auditing.""" 

208 

209 def __init__(self): 

210 self._events: list[dict[str, Any]] = [] 

211 

212 def log(self, event_type: str, details: dict[str, Any]) -> None: 

213 self._events.append( 

214 { 

215 "event": event_type, 

216 "timestamp": time.time(), 

217 **details, 

218 } 

219 ) 

220 

221 def get_events(self, limit: int = 100) -> list[dict[str, Any]]: 

222 return self._events[-limit:] 

223 

224 def clear(self) -> None: 

225 self._events.clear() 

226 

227 

228# --------------------------------------------------------------------------- 

229# Sandbox Executor 

230# --------------------------------------------------------------------------- 

231 

232 

233class SandboxExecutor: 

234 """ 

235 Execute Python code in an isolated sandbox with resource limits. 

236 

237 Uses subprocess isolation with resource limits set via RLIMIT. 

238 """ 

239 

240 def __init__( 

241 self, 

242 policy: SandboxPolicy | None = None, 

243 audit: bool = True, 

244 ): 

245 self._policy = policy or SandboxPolicy() 

246 self._auditor = SandboxAuditLogger() if audit else None 

247 

248 @property 

249 def policy(self) -> SandboxPolicy: 

250 return self._policy 

251 

252 def execute(self, code: str, globals_dict: dict | None = None) -> SandboxResult: 

253 """ 

254 Execute code in sandbox. 

255 

256 Args: 

257 code: Python code string to execute 

258 globals_dict: Optional global namespace (restricted) 

259 

260 Returns: 

261 SandboxResult with execution details 

262 """ 

263 audit_id = f"sandbox_{int(time.time() * 1000)}" 

264 

265 if self._auditor: 

266 self._auditor.log( 

267 "execute", 

268 { 

269 "audit_id": audit_id, 

270 "code_length": len(code), 

271 "policy_level": self._policy.network_enabled and "relaxed" or "restricted", 

272 }, 

273 ) 

274 

275 start = time.time() 

276 peak_memory = 0.0 

277 

278 try: 

279 # Build restricted globals 

280 self._build_restricted_globals(globals_dict or {}) 

281 

282 # Write code to temp file for subprocess isolation 

283 with tempfile.NamedTemporaryFile( 

284 mode="w", suffix=".py", delete=False, prefix="sandbox_", dir=tempfile.gettempdir() 

285 ) as f: 

286 f.write(self._wrap_code_with_limits(code)) 

287 script_path = f.name 

288 

289 try: 

290 result = subprocess.run( 

291 [sys.executable, script_path], 

292 capture_output=True, 

293 text=True, 

294 timeout=self._policy.timeout_seconds, 

295 env=self._build_sandbox_env(), 

296 preexec_fn=self._set_resource_limits if os.name != "nt" else None, 

297 ) 

298 

299 stdout = result.stdout[: self._policy.max_output_size_bytes] 

300 stderr = result.stderr[: self._policy.max_output_size_bytes] 

301 

302 sandbox_result = SandboxResult( 

303 exit_code=result.returncode, 

304 stdout=stdout, 

305 stderr=stderr, 

306 execution_time_ms=(time.time() - start) * 1000, 

307 peak_memory_mb=peak_memory, 

308 ) 

309 

310 except subprocess.TimeoutExpired: 

311 sandbox_result = SandboxResult( 

312 exit_code=-1, 

313 stdout="", 

314 stderr="", 

315 execution_time_ms=(time.time() - start) * 1000, 

316 peak_memory_mb=peak_memory, 

317 was_killed=True, 

318 kill_reason=f"Timeout: exceeded {self._policy.timeout_seconds}s", 

319 error="Execution timed out", 

320 ) 

321 finally: 

322 try: 

323 os.unlink(script_path) 

324 except OSError: 

325 pass 

326 

327 except Exception as e: 

328 sandbox_result = SandboxResult( 

329 exit_code=-1, 

330 stdout="", 

331 stderr=traceback.format_exc(), 

332 execution_time_ms=(time.time() - start) * 1000, 

333 peak_memory_mb=peak_memory, 

334 error=str(e), 

335 ) 

336 

337 if self._auditor: 

338 self._auditor.log( 

339 "complete", 

340 { 

341 "audit_id": audit_id, 

342 "success": sandbox_result.success, 

343 "execution_time_ms": sandbox_result.execution_time_ms, 

344 "exit_code": sandbox_result.exit_code, 

345 }, 

346 ) 

347 

348 return sandbox_result 

349 

350 def execute_sync(self, code: str, namespace: dict | None = None) -> SandboxResult: 

351 """ 

352 Execute code synchronously in-process with restricted namespace. 

353 

354 Faster than execute() but slightly less isolated. 

355 """ 

356 start = time.time() 

357 audit_id = f"sync_{int(start * 1000)}" 

358 

359 if self._auditor: 

360 self._auditor.log("execute_sync", {"audit_id": audit_id, "code_length": len(code)}) 

361 

362 try: 

363 restricted_globals = self._build_restricted_globals({}) 

364 

365 exec(code, restricted_globals) 

366 

367 # Propagate results back to caller namespace 

368 if namespace is not None: 

369 namespace.update( 

370 { 

371 k: v 

372 for k, v in restricted_globals.items() 

373 if not k.startswith("__") and k not in self._policy.allowed_modules 

374 } 

375 ) 

376 

377 result = SandboxResult( 

378 exit_code=0, 

379 stdout="", 

380 stderr="", 

381 execution_time_ms=(time.time() - start) * 1000, 

382 peak_memory_mb=0.0, 

383 ) 

384 except Exception as e: 

385 result = SandboxResult( 

386 exit_code=-1, 

387 stdout="", 

388 stderr=traceback.format_exc(), 

389 execution_time_ms=(time.time() - start) * 1000, 

390 peak_memory_mb=0.0, 

391 error=str(e), 

392 ) 

393 

394 if self._auditor: 

395 self._auditor.log( 

396 "complete_sync", 

397 { 

398 "audit_id": audit_id, 

399 "success": result.success, 

400 }, 

401 ) 

402 

403 return result 

404 

405 def _build_restricted_globals(self, extra: dict) -> dict: 

406 """Build a restricted global namespace.""" 

407 safe_builtins = { 

408 k: v 

409 for k, v in __builtins__.items() 

410 if k not in (self._policy.block_dangerous_builtins and DANGEROUS_BUILTINS or set()) 

411 } 

412 

413 globals_dict = {"__builtins__": safe_builtins} 

414 

415 # Add allowed modules 

416 for mod_name in self._policy.allowed_modules: 

417 try: 

418 mod = __import__(mod_name) 

419 globals_dict[mod_name] = mod 

420 except ImportError: 

421 pass 

422 

423 globals_dict.update(extra) 

424 return globals_dict 

425 

426 def _build_sandbox_env(self) -> dict[str, str]: 

427 """Build a restricted environment for subprocess.""" 

428 env = { 

429 "PATH": "/usr/bin:/bin:/usr/local/bin", 

430 "HOME": tempfile.gettempdir(), 

431 "TMPDIR": tempfile.gettempdir(), 

432 "PYTHONUNBUFFERED": "1", 

433 "PYTHONDONTWRITEBYTECODE": "1", 

434 "SANDBOX_MAX_CPU": str(self._policy.max_cpu_time_seconds), 

435 "SANDBOX_MAX_MEMORY_MB": str(self._policy.max_memory_mb), 

436 } 

437 return env 

438 

439 def _set_resource_limits(self) -> None: 

440 """Set OS-level resource limits (Unix only).""" 

441 try: 

442 # CPU time limit 

443 cpu_seconds = int(self._policy.max_cpu_time_seconds) 

444 resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds + 5)) 

445 

446 # Memory limit 

447 mem_bytes = self._policy.max_memory_mb * 1024 * 1024 

448 resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes)) 

449 

450 # File descriptors 

451 resource.setrlimit( 

452 resource.RLIMIT_NOFILE, 

453 (self._policy.max_files_open, self._policy.max_files_open), 

454 ) 

455 

456 # Processes 

457 resource.setrlimit( 

458 resource.RLIMIT_NPROC, 

459 (self._policy.max_processes, self._policy.max_processes), 

460 ) 

461 except (OSError, ValueError): 

462 pass 

463 

464 def _wrap_code_with_limits(self, code: str) -> str: 

465 """Wrap user code with safety limits.""" 

466 wrapper = f""" 

467import sys 

468 

469# Blacklist dangerous modules via import hook 

470_dangerous_modules = {repr(DANGEROUS_MODULES)} 

471 

472class _RestrictedFinder: 

473 def find_spec(self, fullname, path, target=None): 

474 if fullname in _dangerous_modules or fullname.startswith(tuple(m + '.' for m in _dangerous_modules)): 

475 raise ImportError(f"Module '{{fullname}}' is restricted in sandbox") 

476 return None 

477 

478sys.meta_path.insert(0, _RestrictedFinder()) 

479 

480# Purge pre-cached dangerous modules from sys.modules 

481for _m in list(sys.modules.keys()): 

482 if _m in _dangerous_modules or any(_m.startswith(_d + '.') for _d in _dangerous_modules): 

483 del sys.modules[_m] 

484 

485# Execute user code 

486{code} 

487""" 

488 return wrapper 

489 

490 def get_audit_log(self) -> list[dict[str, Any]]: 

491 if self._auditor: 

492 return self._auditor.get_events() 

493 return [] 

494 

495 def clear_audit(self) -> None: 

496 if self._auditor: 

497 self._auditor.clear() 

498 

499 

500# --------------------------------------------------------------------------- 

501# Backward-compatible aliases (v1.x migration) 

502# --------------------------------------------------------------------------- 

503 

504# RiskLevel = SecurityLevel (alias) 

505RiskLevel = SecurityLevel 

506 

507 

508@dataclass 

509class SafetyReport: 

510 """Safety analysis report for code/input evaluation.""" 

511 

512 risk_level: RiskLevel = RiskLevel.STANDARD 

513 is_safe: bool = True 

514 findings: list[str] = field(default_factory=list) 

515 recommendation: str = "" 

516 score: float = 1.0 

517 

518 def to_dict(self) -> dict[str, Any]: 

519 return { 

520 "risk_level": self.risk_level.value, 

521 "is_safe": self.is_safe, 

522 "findings": self.findings, 

523 "recommendation": self.recommendation, 

524 "score": self.score, 

525 } 

526 

527 

528class LLMSafetyAnalyzer: 

529 """LLM-based safety analyzer for code and content review.""" 

530 

531 def __init__(self, policy: SandboxPolicy | None = None): 

532 self._policy = policy or SandboxPolicy() 

533 

534 def analyze(self, code: str) -> SafetyReport: 

535 findings: list[str] = [] 

536 is_safe = True 

537 score = 1.0 

538 

539 for mod in DANGEROUS_MODULES: 

540 if f"import {mod}" in code or f"from {mod}" in code: 

541 findings.append(f"Uses dangerous module: {mod}") 

542 is_safe = False 

543 score = max(0.0, score - 0.15) 

544 

545 for builtin_name in DANGEROUS_BUILTINS: 

546 if builtin_name in code: 

547 findings.append(f"Uses dangerous builtin: {builtin_name}") 

548 is_safe = False 

549 score = max(0.0, score - 0.1) 

550 

551 if not is_safe: 

552 return SafetyReport( 

553 risk_level=RiskLevel.RESTRICTED, 

554 is_safe=False, 

555 findings=findings, 

556 recommendation="Code contains potentially dangerous operations", 

557 score=score, 

558 ) 

559 return SafetyReport( 

560 risk_level=RiskLevel.STANDARD, 

561 is_safe=True, 

562 findings=findings, 

563 score=score, 

564 ) 

565 

566 

567class Sandbox: 

568 """Backward-compatible Sandbox wrapper around SandboxPolicy + SandboxExecutor.""" 

569 

570 def __init__(self, policy: SandboxPolicy | None = None): 

571 self._policy = policy or SandboxPolicy() 

572 self._executor = SandboxExecutor(policy=self._policy) 

573 

574 @property 

575 def policy(self) -> SandboxPolicy: 

576 return self._policy 

577 

578 def run(self, code: str) -> SandboxResult: 

579 return self._executor.execute(code) 

580 

581 def run_sync(self, code: str, namespace: dict | None = None) -> SandboxResult: 

582 return self._executor.execute_sync(code, namespace) 

583 

584 

585class SandboxManager: 

586 """Backward-compatible SandboxManager — manages multiple sandbox instances.""" 

587 

588 def __init__(self, default_policy: SandboxPolicy | None = None): 

589 self._default_policy = default_policy or SandboxPolicy() 

590 self._sandboxes: dict[str, Sandbox] = {} 

591 

592 def create(self, name: str, policy: SandboxPolicy | None = None) -> Sandbox: 

593 sb = Sandbox(policy=policy or self._default_policy) 

594 self._sandboxes[name] = sb 

595 return sb 

596 

597 def get(self, name: str) -> Sandbox | None: 

598 return self._sandboxes.get(name) 

599 

600 def remove(self, name: str) -> bool: 

601 if name in self._sandboxes: 

602 del self._sandboxes[name] 

603 return True 

604 return False 

605 

606 def list(self) -> builtins.list[str]: 

607 return list(self._sandboxes.keys()) 

608 

609 def execute_all(self, code: str) -> dict[str, SandboxResult]: 

610 return {name: sb.run(code) for name, sb in self._sandboxes.items()}