133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
# Copyright (c) 2025 Patrick Motsch
|
|
# All rights reserved.
|
|
"""Sandboxed code execution for the AI agent executeCode tool."""
|
|
|
|
import logging
|
|
import sys
|
|
import io
|
|
import traceback
|
|
from typing import Dict, Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SANDBOX_ALLOWED_MODULES = {
|
|
"math", "statistics", "json", "csv", "re", "datetime", "time",
|
|
"collections", "itertools", "functools", "decimal", "fractions",
|
|
"random", "string", "textwrap", "operator", "copy",
|
|
}
|
|
|
|
_PYTHON_BLOCKED_BUILTINS = {
|
|
"open", "exec", "eval", "compile", "__import__", "globals", "locals",
|
|
"getattr", "setattr", "delattr", "breakpoint", "exit", "quit",
|
|
"input", "memoryview",
|
|
}
|
|
|
|
_MAX_EXECUTION_TIME_S = 30
|
|
_MAX_OUTPUT_CHARS = 50000
|
|
|
|
|
|
_RESTRICTED_IO = None
|
|
|
|
def _getRestrictedIo():
|
|
"""Return a restricted ``io`` module exposing only StringIO/BytesIO."""
|
|
global _RESTRICTED_IO
|
|
if _RESTRICTED_IO is None:
|
|
import types
|
|
m = types.ModuleType("io")
|
|
m.StringIO = io.StringIO
|
|
m.BytesIO = io.BytesIO
|
|
_RESTRICTED_IO = m
|
|
return _RESTRICTED_IO
|
|
|
|
|
|
def _safeImport(name, *args, **kwargs):
|
|
"""Restricted import that only allows whitelisted modules."""
|
|
if name == "io":
|
|
return _getRestrictedIo()
|
|
if name not in SANDBOX_ALLOWED_MODULES:
|
|
raise ImportError(f"Module '{name}' is not allowed. Permitted: io (StringIO/BytesIO only), {', '.join(sorted(SANDBOX_ALLOWED_MODULES))}")
|
|
return __builtins__["__import__"](name, *args, **kwargs) if isinstance(__builtins__, dict) else __import__(name, *args, **kwargs)
|
|
|
|
|
|
def _buildRestrictedGlobals() -> Dict[str, Any]:
|
|
"""Build a restricted globals dict for exec()."""
|
|
import builtins
|
|
safeBuiltins = {}
|
|
for name in dir(builtins):
|
|
if name.startswith("_"):
|
|
continue
|
|
if name in _PYTHON_BLOCKED_BUILTINS:
|
|
continue
|
|
safeBuiltins[name] = getattr(builtins, name)
|
|
|
|
safeBuiltins["__import__"] = _safeImport
|
|
safeBuiltins["__name__"] = "__sandbox__"
|
|
safeBuiltins["__builtins__"] = safeBuiltins
|
|
|
|
for modName in SANDBOX_ALLOWED_MODULES:
|
|
try:
|
|
safeBuiltins[modName] = __import__(modName)
|
|
except ImportError:
|
|
pass
|
|
|
|
return {"__builtins__": safeBuiltins}
|
|
|
|
|
|
def _makeReadFile(services):
|
|
"""Create a readFile(fileId) closure bound to the current services context."""
|
|
def readFile(fileId: str) -> str:
|
|
mgmt = getattr(services, 'interfaceDbComponent', None) if services else None
|
|
if not mgmt:
|
|
raise RuntimeError("readFile: no file store available in this session")
|
|
data = mgmt.getFileData(str(fileId))
|
|
if data is None:
|
|
raise FileNotFoundError(f"File '{fileId}' not found in workspace")
|
|
return data.decode("utf-8")
|
|
return readFile
|
|
|
|
|
|
async def executePython(code: str, *, services=None) -> Dict[str, Any]:
|
|
"""Execute Python code in a restricted sandbox. Returns {success, output, error}."""
|
|
import asyncio
|
|
|
|
def _run():
|
|
restrictedGlobals = _buildRestrictedGlobals()
|
|
if services:
|
|
restrictedGlobals["__builtins__"]["readFile"] = _makeReadFile(services)
|
|
capturedOutput = io.StringIO()
|
|
oldStdout = sys.stdout
|
|
oldStderr = sys.stderr
|
|
|
|
try:
|
|
sys.stdout = capturedOutput
|
|
sys.stderr = capturedOutput
|
|
|
|
# Do not use signal.SIGALRM here: _run executes inside a thread-pool worker
|
|
# (asyncio.run_in_executor). signal.signal only works on the main thread.
|
|
# Wall-clock limit is enforced by asyncio.wait_for around run_in_executor.
|
|
|
|
exec(compile(code, "<sandbox>", "exec"), restrictedGlobals)
|
|
|
|
output = capturedOutput.getvalue()
|
|
if len(output) > _MAX_OUTPUT_CHARS:
|
|
output = output[:_MAX_OUTPUT_CHARS] + f"\n... (truncated at {_MAX_OUTPUT_CHARS} chars)"
|
|
return {"success": True, "output": output}
|
|
|
|
except TimeoutError:
|
|
return {"success": False, "error": f"Execution timed out after {_MAX_EXECUTION_TIME_S}s"}
|
|
except Exception as e:
|
|
tb = traceback.format_exc()
|
|
return {"success": False, "error": f"{type(e).__name__}: {e}", "traceback": tb}
|
|
finally:
|
|
sys.stdout = oldStdout
|
|
sys.stderr = oldStderr
|
|
|
|
loop = asyncio.get_event_loop()
|
|
try:
|
|
result = await asyncio.wait_for(
|
|
loop.run_in_executor(None, _run),
|
|
timeout=float(_MAX_EXECUTION_TIME_S) + 5.0,
|
|
)
|
|
return result
|
|
except asyncio.TimeoutError:
|
|
return {"success": False, "error": f"Execution timed out after {_MAX_EXECUTION_TIME_S}s"}
|