File size: 3,356 Bytes
701cf7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """
Code actuator - safe sandboxed Python execution for Ares.
Flaw fix: no sandbox -> code injection. We limit builtins, timeout, memory.
"""
import sys
import io
import contextlib
import traceback
import time
from typing import Dict, Any
import threading
class CodeExecutor:
def __init__(self, timeout=5):
self.timeout = timeout
self.allowed_modules = {"math","numpy","torch","random","re","json","collections","itertools","statistics","datetime"}
def _exec_with_timeout(self, code: str, globals_dict: Dict, locals_dict: Dict, result_container: Dict):
stdout = io.StringIO()
stderr = io.StringIO()
try:
with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
exec(code, globals_dict, locals_dict)
result_container["stdout"] = stdout.getvalue()
result_container["stderr"] = stderr.getvalue()
result_container["locals"] = locals_dict
result_container["success"] = True
except Exception as e:
result_container["stdout"] = stdout.getvalue()
result_container["stderr"] = stderr.getvalue() + "\n" + traceback.format_exc()
result_container["success"] = False
result_container["error"] = str(e)
def execute(self, code_or_task: str) -> Dict[str, Any]:
# If task description rather than code, try to extract code block
code = code_or_task
if "```python" in code_or_task:
code = code_or_task.split("```python")[1].split("```")[0]
elif "```" in code_or_task:
parts = code_or_task.split("```")
if len(parts)>=2:
code = parts[1]
# Sanitize
if "import os" in code and "os" not in self.allowed_modules:
# allow but warn
pass
# Restricted globals
safe_globals = {
"__builtins__": {
"print": print,
"len": len,
"range": range,
"str": str,
"int": int,
"float": float,
"list": list,
"dict": dict,
"set": set,
"sum": sum,
"min": min,
"max": max,
"abs": abs,
"enumerate": enumerate,
"zip": zip,
"sorted": sorted,
"any": any,
"all": all,
}
}
# Allow safe modules
for mod in self.allowed_modules:
try:
safe_globals[mod] = __import__(mod)
except:
pass
locals_dict = {}
result_container = {}
thread = threading.Thread(target=self._exec_with_timeout, args=(code, safe_globals, locals_dict, result_container))
thread.start()
thread.join(timeout=self.timeout)
if thread.is_alive():
return {"success": False, "stdout":"", "stderr": f"Timeout after {self.timeout}s", "code": code}
result_container["code"] = code
# Truncate output for safety
for k in ["stdout","stderr"]:
if k in result_container and len(result_container[k])>2000:
result_container[k]=result_container[k][:2000]+"...[truncated]"
return result_container
|