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

210 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-05 20:52 +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 os 

23import resource 

24import signal 

25import subprocess 

26import sys 

27import tempfile 

28import time 

29import traceback 

30from dataclasses import dataclass, field 

31from enum import Enum 

32from typing import Any, Callable, Dict, List, Optional, Set 

33 

34 

35# --------------------------------------------------------------------------- 

36# Sandbox Policy 

37# --------------------------------------------------------------------------- 

38 

39 

40class SecurityLevel(str, Enum): 

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

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", "subprocess", "shlex", "sys", "ctypes", 

50 "socket", "http", "requests", "urllib", 

51 "multiprocessing", "threading", "concurrent.futures", 

52 "importlib", "pkgutil", "pkg_resources", 

53 "pickle", "marshal", "shelve", 

54 "pathlib", "shutil", "glob", "fnmatch", 

55 "signal", "atexit", 

56 "code", "codeop", "compileall", 

57 "pty", "fcntl", "tty", "termios", 

58 "webbrowser", 

59} 

60 

61DANGEROUS_BUILTINS: Set[str] = { 

62 "__import__", "compile", "eval", "exec", "open", 

63 "input", "breakpoint", 

64 "globals", "locals", "vars", 

65 "getattr", "setattr", "delattr", 

66 "type", "issubclass", "isinstance", 

67 "memoryview", "bytearray", 

68} 

69 

70 

71@dataclass 

72class SandboxPolicy: 

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

74 

75 # Resource limits 

76 max_cpu_time_seconds: float = 30.0 

77 max_memory_mb: int = 256 

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

79 max_files_open: int = 50 

80 max_processes: int = 1 

81 network_enabled: bool = False 

82 disk_write_enabled: bool = False 

83 

84 # Execution constraints 

85 timeout_seconds: float = 60.0 

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

87 

88 # Security 

89 block_dangerous_modules: bool = True 

90 block_dangerous_builtins: bool = True 

91 allowed_modules: Set[str] = field(default_factory=set) 

92 allowed_imports: Set[str] = field(default_factory=set) 

93 

94 # Audit 

95 enable_audit: bool = True 

96 

97 @classmethod 

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

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

100 if level == SecurityLevel.RESTRICTED: 

101 return cls( 

102 max_cpu_time_seconds=5.0, 

103 max_memory_mb=64, 

104 max_output_size_bytes=1 * 1024 * 1024, 

105 network_enabled=False, 

106 disk_write_enabled=False, 

107 timeout_seconds=10.0, 

108 allowed_modules={"math", "json", "datetime", "collections", "itertools", "functools"}, 

109 ) 

110 elif level == SecurityLevel.RELAXED: 

111 return cls( 

112 max_cpu_time_seconds=120.0, 

113 max_memory_mb=1024, 

114 max_output_size_bytes=100 * 1024 * 1024, 

115 network_enabled=True, 

116 disk_write_enabled=True, 

117 timeout_seconds=300.0, 

118 block_dangerous_modules=False, 

119 ) 

120 else: # STANDARD 

121 return cls() 

122 

123 

124# --------------------------------------------------------------------------- 

125# Sandbox Result 

126# --------------------------------------------------------------------------- 

127 

128 

129@dataclass(frozen=True) 

130class SandboxResult: 

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

132 exit_code: int 

133 stdout: str 

134 stderr: str 

135 execution_time_ms: float 

136 peak_memory_mb: float 

137 was_killed: bool = False 

138 kill_reason: str = "" 

139 error: Optional[str] = None 

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

141 

142 @property 

143 def success(self) -> bool: 

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

145 

146 def to_dict(self) -> Dict[str, Any]: 

147 return { 

148 "exit_code": self.exit_code, 

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

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

151 "execution_time_ms": self.execution_time_ms, 

152 "peak_memory_mb": self.peak_memory_mb, 

153 "was_killed": self.was_killed, 

154 "kill_reason": self.kill_reason, 

155 "success": self.success, 

156 } 

157 

158 

159# --------------------------------------------------------------------------- 

160# Audit Logger 

161# --------------------------------------------------------------------------- 

162 

163 

164class SandboxAuditLogger: 

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

166 

167 def __init__(self): 

168 self._events: List[Dict[str, Any]] = [] 

169 

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

171 self._events.append({ 

172 "event": event_type, 

173 "timestamp": time.time(), 

174 **details, 

175 }) 

176 

177 def get_events(self, limit: int = 100) -> List[Dict[str, Any]]: 

178 return self._events[-limit:] 

179 

180 def clear(self) -> None: 

181 self._events.clear() 

182 

183 

184# --------------------------------------------------------------------------- 

185# Sandbox Executor 

186# --------------------------------------------------------------------------- 

187 

188 

189class SandboxExecutor: 

190 """ 

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

192 

193 Uses subprocess isolation with resource limits set via RLIMIT. 

194 """ 

195 

196 def __init__( 

197 self, 

198 policy: Optional[SandboxPolicy] = None, 

199 audit: bool = True, 

200 ): 

201 self._policy = policy or SandboxPolicy() 

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

203 

204 @property 

205 def policy(self) -> SandboxPolicy: 

206 return self._policy 

207 

208 def execute(self, code: str, globals_dict: Optional[Dict] = None) -> SandboxResult: 

209 """ 

210 Execute code in sandbox. 

211 

212 Args: 

213 code: Python code string to execute 

214 globals_dict: Optional global namespace (restricted) 

215 

216 Returns: 

217 SandboxResult with execution details 

218 """ 

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

220 

221 if self._auditor: 

222 self._auditor.log("execute", { 

223 "audit_id": audit_id, 

224 "code_length": len(code), 

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

226 }) 

227 

228 start = time.time() 

229 peak_memory = 0.0 

230 

231 try: 

232 # Build restricted globals 

233 restricted_globals = self._build_restricted_globals(globals_dict or {}) 

234 

235 # Write code to temp file for subprocess isolation 

236 with tempfile.NamedTemporaryFile( 

237 mode="w", suffix=".py", delete=False, prefix="sandbox_", 

238 dir=tempfile.gettempdir() 

239 ) as f: 

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

241 script_path = f.name 

242 

243 try: 

244 result = subprocess.run( 

245 [sys.executable, script_path], 

246 capture_output=True, 

247 text=True, 

248 timeout=self._policy.timeout_seconds, 

249 env=self._build_sandbox_env(), 

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

251 ) 

252 

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

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

255 

256 sandbox_result = SandboxResult( 

257 exit_code=result.returncode, 

258 stdout=stdout, 

259 stderr=stderr, 

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

261 peak_memory_mb=peak_memory, 

262 ) 

263 

264 except subprocess.TimeoutExpired: 

265 sandbox_result = SandboxResult( 

266 exit_code=-1, 

267 stdout="", 

268 stderr="", 

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

270 peak_memory_mb=peak_memory, 

271 was_killed=True, 

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

273 error="Execution timed out", 

274 ) 

275 finally: 

276 try: 

277 os.unlink(script_path) 

278 except OSError: 

279 pass 

280 

281 except Exception as e: 

282 sandbox_result = SandboxResult( 

283 exit_code=-1, 

284 stdout="", 

285 stderr=traceback.format_exc(), 

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

287 peak_memory_mb=peak_memory, 

288 error=str(e), 

289 ) 

290 

291 if self._auditor: 

292 self._auditor.log("complete", { 

293 "audit_id": audit_id, 

294 "success": sandbox_result.success, 

295 "execution_time_ms": sandbox_result.execution_time_ms, 

296 "exit_code": sandbox_result.exit_code, 

297 }) 

298 

299 return sandbox_result 

300 

301 def execute_sync(self, code: str, namespace: Optional[Dict] = None) -> SandboxResult: 

302 """ 

303 Execute code synchronously in-process with restricted namespace. 

304 

305 Faster than execute() but slightly less isolated. 

306 """ 

307 start = time.time() 

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

309 

310 if self._auditor: 

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

312 

313 try: 

314 restricted_globals = self._build_restricted_globals({}) 

315 

316 exec(code, restricted_globals) 

317 

318 # Propagate results back to caller namespace 

319 if namespace is not None: 

320 namespace.update({k: v for k, v in restricted_globals.items() 

321 if not k.startswith('__') and k not in self._policy.allowed_modules}) 

322 

323 result = SandboxResult( 

324 exit_code=0, 

325 stdout="", 

326 stderr="", 

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

328 peak_memory_mb=0.0, 

329 ) 

330 except Exception as e: 

331 result = SandboxResult( 

332 exit_code=-1, 

333 stdout="", 

334 stderr=traceback.format_exc(), 

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

336 peak_memory_mb=0.0, 

337 error=str(e), 

338 ) 

339 

340 if self._auditor: 

341 self._auditor.log("complete_sync", { 

342 "audit_id": audit_id, 

343 "success": result.success, 

344 }) 

345 

346 return result 

347 

348 def _build_restricted_globals(self, extra: Dict) -> Dict: 

349 """Build a restricted global namespace.""" 

350 safe_builtins = { 

351 k: v for k, v in __builtins__.items() 

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

353 } 

354 

355 globals_dict = {"__builtins__": safe_builtins} 

356 

357 # Add allowed modules 

358 for mod_name in self._policy.allowed_modules: 

359 try: 

360 mod = __import__(mod_name) 

361 globals_dict[mod_name] = mod 

362 except ImportError: 

363 pass 

364 

365 globals_dict.update(extra) 

366 return globals_dict 

367 

368 def _build_sandbox_env(self) -> Dict[str, str]: 

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

370 env = { 

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

372 "HOME": tempfile.gettempdir(), 

373 "TMPDIR": tempfile.gettempdir(), 

374 "PYTHONUNBUFFERED": "1", 

375 "PYTHONDONTWRITEBYTECODE": "1", 

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

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

378 } 

379 return env 

380 

381 def _set_resource_limits(self) -> None: 

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

383 try: 

384 # CPU time limit 

385 cpu_seconds = int(self._policy.max_cpu_time_seconds) 

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

387 

388 # Memory limit 

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

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

391 

392 # File descriptors 

393 resource.setrlimit( 

394 resource.RLIMIT_NOFILE, 

395 (self._policy.max_files_open, self._policy.max_files_open), 

396 ) 

397 

398 # Processes 

399 resource.setrlimit( 

400 resource.RLIMIT_NPROC, 

401 (self._policy.max_processes, self._policy.max_processes), 

402 ) 

403 except (ValueError, resource.error): 

404 pass 

405 

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

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

408 wrapper = f''' 

409import sys 

410 

411# Blacklist dangerous modules via import hook 

412_dangerous_modules = {repr(DANGEROUS_MODULES)} 

413 

414class _RestrictedFinder: 

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

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

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

418 return None 

419 

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

421 

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

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

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

425 del sys.modules[_m] 

426 

427# Execute user code 

428{code} 

429''' 

430 return wrapper 

431 

432 def get_audit_log(self) -> List[Dict[str, Any]]: 

433 if self._auditor: 

434 return self._auditor.get_events() 

435 return [] 

436 

437 def clear_audit(self) -> None: 

438 if self._auditor: 

439 self._auditor.clear() 

440 

441 

442# --------------------------------------------------------------------------- 

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

444# --------------------------------------------------------------------------- 

445 

446# RiskLevel = SecurityLevel (alias) 

447RiskLevel = SecurityLevel 

448 

449 

450@dataclass 

451class SafetyReport: 

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

453 risk_level: RiskLevel = RiskLevel.STANDARD 

454 is_safe: bool = True 

455 findings: List[str] = field(default_factory=list) 

456 recommendation: str = "" 

457 score: float = 1.0 

458 

459 def to_dict(self) -> Dict[str, Any]: 

460 return { 

461 "risk_level": self.risk_level.value, 

462 "is_safe": self.is_safe, 

463 "findings": self.findings, 

464 "recommendation": self.recommendation, 

465 "score": self.score, 

466 } 

467 

468 

469class LLMSafetyAnalyzer: 

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

471 

472 def __init__(self, policy: Optional[SandboxPolicy] = None): 

473 self._policy = policy or SandboxPolicy() 

474 

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

476 findings: List[str] = [] 

477 is_safe = True 

478 score = 1.0 

479 

480 for mod in DANGEROUS_MODULES: 

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

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

483 is_safe = False 

484 score = max(0.0, score - 0.15) 

485 

486 for builtin_name in DANGEROUS_BUILTINS: 

487 if builtin_name in code: 

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

489 is_safe = False 

490 score = max(0.0, score - 0.1) 

491 

492 if not is_safe: 

493 return SafetyReport( 

494 risk_level=RiskLevel.RESTRICTED, 

495 is_safe=False, 

496 findings=findings, 

497 recommendation="Code contains potentially dangerous operations", 

498 score=score, 

499 ) 

500 return SafetyReport( 

501 risk_level=RiskLevel.STANDARD, 

502 is_safe=True, 

503 findings=findings, 

504 score=score, 

505 ) 

506 

507 

508class Sandbox: 

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

510 

511 def __init__(self, policy: Optional[SandboxPolicy] = None): 

512 self._policy = policy or SandboxPolicy() 

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

514 

515 @property 

516 def policy(self) -> SandboxPolicy: 

517 return self._policy 

518 

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

520 return self._executor.execute(code) 

521 

522 def run_sync(self, code: str, namespace: Optional[Dict] = None) -> SandboxResult: 

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

524 

525 

526class SandboxManager: 

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

528 

529 def __init__(self, default_policy: Optional[SandboxPolicy] = None): 

530 self._default_policy = default_policy or SandboxPolicy() 

531 self._sandboxes: Dict[str, Sandbox] = {} 

532 

533 def create(self, name: str, policy: Optional[SandboxPolicy] = None) -> Sandbox: 

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

535 self._sandboxes[name] = sb 

536 return sb 

537 

538 def get(self, name: str) -> Optional[Sandbox]: 

539 return self._sandboxes.get(name) 

540 

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

542 if name in self._sandboxes: 

543 del self._sandboxes[name] 

544 return True 

545 return False 

546 

547 def list(self) -> List[str]: 

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

549 

550 def execute_all(self, code: str) -> Dict[str, SandboxResult]: 

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