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

209 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-06 10:59 +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 subprocess 

25import sys 

26import tempfile 

27import time 

28import traceback 

29from dataclasses import dataclass, field 

30from enum import Enum 

31from typing import Any, Dict, List, Optional, Set 

32 

33 

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

35# Sandbox Policy 

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

37 

38 

39class SecurityLevel(str, Enum): 

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

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

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

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

44 

45 

46# Dangerous modules and builtins to block 

47DANGEROUS_MODULES: Set[str] = { 

48 "os", "subprocess", "shlex", "sys", "ctypes", 

49 "socket", "http", "requests", "urllib", 

50 "multiprocessing", "threading", "concurrent.futures", 

51 "importlib", "pkgutil", "pkg_resources", 

52 "pickle", "marshal", "shelve", 

53 "pathlib", "shutil", "glob", "fnmatch", 

54 "signal", "atexit", 

55 "code", "codeop", "compileall", 

56 "pty", "fcntl", "tty", "termios", 

57 "webbrowser", 

58} 

59 

60DANGEROUS_BUILTINS: Set[str] = { 

61 "__import__", "compile", "eval", "exec", "open", 

62 "input", "breakpoint", 

63 "globals", "locals", "vars", 

64 "getattr", "setattr", "delattr", 

65 "type", "issubclass", "isinstance", 

66 "memoryview", "bytearray", 

67} 

68 

69 

70@dataclass 

71class SandboxPolicy: 

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

73 

74 # Resource limits 

75 max_cpu_time_seconds: float = 30.0 

76 max_memory_mb: int = 256 

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

78 max_files_open: int = 50 

79 max_processes: int = 1 

80 network_enabled: bool = False 

81 disk_write_enabled: bool = False 

82 

83 # Execution constraints 

84 timeout_seconds: float = 60.0 

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

86 

87 # Security 

88 block_dangerous_modules: bool = True 

89 block_dangerous_builtins: bool = True 

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

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

92 

93 # Audit 

94 enable_audit: bool = True 

95 

96 @classmethod 

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

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

99 if level == SecurityLevel.RESTRICTED: 

100 return cls( 

101 max_cpu_time_seconds=5.0, 

102 max_memory_mb=64, 

103 max_output_size_bytes=1 * 1024 * 1024, 

104 network_enabled=False, 

105 disk_write_enabled=False, 

106 timeout_seconds=10.0, 

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

108 ) 

109 elif level == SecurityLevel.RELAXED: 

110 return cls( 

111 max_cpu_time_seconds=120.0, 

112 max_memory_mb=1024, 

113 max_output_size_bytes=100 * 1024 * 1024, 

114 network_enabled=True, 

115 disk_write_enabled=True, 

116 timeout_seconds=300.0, 

117 block_dangerous_modules=False, 

118 ) 

119 else: # STANDARD 

120 return cls() 

121 

122 

123# --------------------------------------------------------------------------- 

124# Sandbox Result 

125# --------------------------------------------------------------------------- 

126 

127 

128@dataclass(frozen=True) 

129class SandboxResult: 

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

131 exit_code: int 

132 stdout: str 

133 stderr: str 

134 execution_time_ms: float 

135 peak_memory_mb: float 

136 was_killed: bool = False 

137 kill_reason: str = "" 

138 error: Optional[str] = None 

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

140 

141 @property 

142 def success(self) -> bool: 

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

144 

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

146 return { 

147 "exit_code": self.exit_code, 

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

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

150 "execution_time_ms": self.execution_time_ms, 

151 "peak_memory_mb": self.peak_memory_mb, 

152 "was_killed": self.was_killed, 

153 "kill_reason": self.kill_reason, 

154 "success": self.success, 

155 } 

156 

157 

158# --------------------------------------------------------------------------- 

159# Audit Logger 

160# --------------------------------------------------------------------------- 

161 

162 

163class SandboxAuditLogger: 

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

165 

166 def __init__(self): 

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

168 

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

170 self._events.append({ 

171 "event": event_type, 

172 "timestamp": time.time(), 

173 **details, 

174 }) 

175 

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

177 return self._events[-limit:] 

178 

179 def clear(self) -> None: 

180 self._events.clear() 

181 

182 

183# --------------------------------------------------------------------------- 

184# Sandbox Executor 

185# --------------------------------------------------------------------------- 

186 

187 

188class SandboxExecutor: 

189 """ 

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

191 

192 Uses subprocess isolation with resource limits set via RLIMIT. 

193 """ 

194 

195 def __init__( 

196 self, 

197 policy: Optional[SandboxPolicy] = None, 

198 audit: bool = True, 

199 ): 

200 self._policy = policy or SandboxPolicy() 

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

202 

203 @property 

204 def policy(self) -> SandboxPolicy: 

205 return self._policy 

206 

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

208 """ 

209 Execute code in sandbox. 

210 

211 Args: 

212 code: Python code string to execute 

213 globals_dict: Optional global namespace (restricted) 

214 

215 Returns: 

216 SandboxResult with execution details 

217 """ 

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

219 

220 if self._auditor: 

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

222 "audit_id": audit_id, 

223 "code_length": len(code), 

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

225 }) 

226 

227 start = time.time() 

228 peak_memory = 0.0 

229 

230 try: 

231 # Build restricted globals 

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

233 

234 # Write code to temp file for subprocess isolation 

235 with tempfile.NamedTemporaryFile( 

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

237 dir=tempfile.gettempdir() 

238 ) as f: 

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

240 script_path = f.name 

241 

242 try: 

243 result = subprocess.run( 

244 [sys.executable, script_path], 

245 capture_output=True, 

246 text=True, 

247 timeout=self._policy.timeout_seconds, 

248 env=self._build_sandbox_env(), 

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

250 ) 

251 

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

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

254 

255 sandbox_result = SandboxResult( 

256 exit_code=result.returncode, 

257 stdout=stdout, 

258 stderr=stderr, 

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

260 peak_memory_mb=peak_memory, 

261 ) 

262 

263 except subprocess.TimeoutExpired: 

264 sandbox_result = SandboxResult( 

265 exit_code=-1, 

266 stdout="", 

267 stderr="", 

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

269 peak_memory_mb=peak_memory, 

270 was_killed=True, 

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

272 error="Execution timed out", 

273 ) 

274 finally: 

275 try: 

276 os.unlink(script_path) 

277 except OSError: 

278 pass 

279 

280 except Exception as e: 

281 sandbox_result = SandboxResult( 

282 exit_code=-1, 

283 stdout="", 

284 stderr=traceback.format_exc(), 

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

286 peak_memory_mb=peak_memory, 

287 error=str(e), 

288 ) 

289 

290 if self._auditor: 

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

292 "audit_id": audit_id, 

293 "success": sandbox_result.success, 

294 "execution_time_ms": sandbox_result.execution_time_ms, 

295 "exit_code": sandbox_result.exit_code, 

296 }) 

297 

298 return sandbox_result 

299 

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

301 """ 

302 Execute code synchronously in-process with restricted namespace. 

303 

304 Faster than execute() but slightly less isolated. 

305 """ 

306 start = time.time() 

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

308 

309 if self._auditor: 

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

311 

312 try: 

313 restricted_globals = self._build_restricted_globals({}) 

314 

315 exec(code, restricted_globals) 

316 

317 # Propagate results back to caller namespace 

318 if namespace is not None: 

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

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

321 

322 result = SandboxResult( 

323 exit_code=0, 

324 stdout="", 

325 stderr="", 

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

327 peak_memory_mb=0.0, 

328 ) 

329 except Exception as e: 

330 result = SandboxResult( 

331 exit_code=-1, 

332 stdout="", 

333 stderr=traceback.format_exc(), 

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

335 peak_memory_mb=0.0, 

336 error=str(e), 

337 ) 

338 

339 if self._auditor: 

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

341 "audit_id": audit_id, 

342 "success": result.success, 

343 }) 

344 

345 return result 

346 

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

348 """Build a restricted global namespace.""" 

349 safe_builtins = { 

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

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

352 } 

353 

354 globals_dict = {"__builtins__": safe_builtins} 

355 

356 # Add allowed modules 

357 for mod_name in self._policy.allowed_modules: 

358 try: 

359 mod = __import__(mod_name) 

360 globals_dict[mod_name] = mod 

361 except ImportError: 

362 pass 

363 

364 globals_dict.update(extra) 

365 return globals_dict 

366 

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

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

369 env = { 

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

371 "HOME": tempfile.gettempdir(), 

372 "TMPDIR": tempfile.gettempdir(), 

373 "PYTHONUNBUFFERED": "1", 

374 "PYTHONDONTWRITEBYTECODE": "1", 

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

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

377 } 

378 return env 

379 

380 def _set_resource_limits(self) -> None: 

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

382 try: 

383 # CPU time limit 

384 cpu_seconds = int(self._policy.max_cpu_time_seconds) 

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

386 

387 # Memory limit 

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

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

390 

391 # File descriptors 

392 resource.setrlimit( 

393 resource.RLIMIT_NOFILE, 

394 (self._policy.max_files_open, self._policy.max_files_open), 

395 ) 

396 

397 # Processes 

398 resource.setrlimit( 

399 resource.RLIMIT_NPROC, 

400 (self._policy.max_processes, self._policy.max_processes), 

401 ) 

402 except (ValueError, resource.error): 

403 pass 

404 

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

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

407 wrapper = f''' 

408import sys 

409 

410# Blacklist dangerous modules via import hook 

411_dangerous_modules = {repr(DANGEROUS_MODULES)} 

412 

413class _RestrictedFinder: 

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

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

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

417 return None 

418 

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

420 

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

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

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

424 del sys.modules[_m] 

425 

426# Execute user code 

427{code} 

428''' 

429 return wrapper 

430 

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

432 if self._auditor: 

433 return self._auditor.get_events() 

434 return [] 

435 

436 def clear_audit(self) -> None: 

437 if self._auditor: 

438 self._auditor.clear() 

439 

440 

441# --------------------------------------------------------------------------- 

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

443# --------------------------------------------------------------------------- 

444 

445# RiskLevel = SecurityLevel (alias) 

446RiskLevel = SecurityLevel 

447 

448 

449@dataclass 

450class SafetyReport: 

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

452 risk_level: RiskLevel = RiskLevel.STANDARD 

453 is_safe: bool = True 

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

455 recommendation: str = "" 

456 score: float = 1.0 

457 

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

459 return { 

460 "risk_level": self.risk_level.value, 

461 "is_safe": self.is_safe, 

462 "findings": self.findings, 

463 "recommendation": self.recommendation, 

464 "score": self.score, 

465 } 

466 

467 

468class LLMSafetyAnalyzer: 

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

470 

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

472 self._policy = policy or SandboxPolicy() 

473 

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

475 findings: List[str] = [] 

476 is_safe = True 

477 score = 1.0 

478 

479 for mod in DANGEROUS_MODULES: 

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

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

482 is_safe = False 

483 score = max(0.0, score - 0.15) 

484 

485 for builtin_name in DANGEROUS_BUILTINS: 

486 if builtin_name in code: 

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

488 is_safe = False 

489 score = max(0.0, score - 0.1) 

490 

491 if not is_safe: 

492 return SafetyReport( 

493 risk_level=RiskLevel.RESTRICTED, 

494 is_safe=False, 

495 findings=findings, 

496 recommendation="Code contains potentially dangerous operations", 

497 score=score, 

498 ) 

499 return SafetyReport( 

500 risk_level=RiskLevel.STANDARD, 

501 is_safe=True, 

502 findings=findings, 

503 score=score, 

504 ) 

505 

506 

507class Sandbox: 

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

509 

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

511 self._policy = policy or SandboxPolicy() 

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

513 

514 @property 

515 def policy(self) -> SandboxPolicy: 

516 return self._policy 

517 

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

519 return self._executor.execute(code) 

520 

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

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

523 

524 

525class SandboxManager: 

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

527 

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

529 self._default_policy = default_policy or SandboxPolicy() 

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

531 

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

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

534 self._sandboxes[name] = sb 

535 return sb 

536 

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

538 return self._sandboxes.get(name) 

539 

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

541 if name in self._sandboxes: 

542 del self._sandboxes[name] 

543 return True 

544 return False 

545 

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

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

548 

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

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