Spaces:
Paused
Paused
File size: 4,590 Bytes
b9f94e1 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """Python execution tool implementation."""
from __future__ import annotations
import asyncio
import contextlib
import logging
import os
import sys
import tempfile
from pathlib import Path
from typing import Any
from hermes.tools.base.tool import BaseTool, ToolSchema
logger = logging.getLogger(__name__)
class PythonExecutionTool(BaseTool):
"""Tool for executing Python code with safety measures."""
def __init__(self, timeout: int = 30, max_output_size: int = 10000, max_memory_mb: int = 256) -> None:
super().__init__()
self.timeout = timeout
self.max_output_size = max_output_size
self.max_memory_mb = max_memory_mb
def _define_schema(self) -> ToolSchema:
return ToolSchema(
name="python_executor",
description="Execute Python code in an isolated environment with resource limits",
parameters={
"code": {
"type": "string",
"description": "Python code to execute",
},
"timeout": {
"type": "integer",
"description": "Execution timeout in seconds",
"default": 30,
},
},
required=["code"],
category="execution",
tags=["python", "execute", "code"],
)
async def execute(self, **kwargs: Any) -> dict[str, Any]:
"""Execute Python code with isolation."""
code = kwargs["code"]
timeout = kwargs.get("timeout", self.timeout)
restricted_code = self._inject_safety(code)
with tempfile.NamedTemporaryFile(
mode="w", suffix=".py", delete=False, encoding="utf-8"
) as f:
f.write(restricted_code)
temp_path = f.name
try:
env = os.environ.copy()
sensitive_keys = [k for k in env if any(word in k.upper() for word in
["API_KEY", "APIKEY", "SECRET", "TOKEN", "PASSWORD", "CREDENTIAL", "AUTH"])]
for k in sensitive_keys:
del env[k]
process = await asyncio.create_subprocess_exec(
sys.executable,
"-I",
temp_path,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
cwd=tempfile.gettempdir(),
)
try:
stdout, stderr = await asyncio.wait_for(
process.communicate(), timeout=timeout
)
except TimeoutError:
process.kill()
await process.wait()
return {
"success": False,
"error": f"Execution timed out after {timeout}s",
"stdout": "",
"stderr": "",
"return_code": -1,
}
stdout_str = stdout.decode("utf-8", errors="replace")[:self.max_output_size]
stderr_str = stderr.decode("utf-8", errors="replace")[:self.max_output_size]
return {
"success": process.returncode == 0,
"stdout": stdout_str,
"stderr": stderr_str,
"return_code": process.returncode,
}
except Exception as e:
logger.error(f"Python execution error: {e}")
return {
"success": False,
"error": str(e),
"stdout": "",
"stderr": "",
"return_code": -1,
}
finally:
with contextlib.suppress(Exception):
Path(temp_path).unlink(missing_ok=True)
def _inject_safety(self, code: str) -> str:
"""Inject safety restrictions into code — blocks imports and subprocess access."""
banner = """import sys as _sys, builtins as _builtins
# Block __import__ — prevents all module imports
_builtins.__import__ = None
# Nullify dangerous modules so even if __import__ is somehow recovered they fail
for _mod in ['os', 'subprocess', 'shutil', 'socket', 'ctypes', 'pickle',
'marshal', 'compile', 'codeop', 'code', 'inspect',
'importlib', 'pkgutil', 'mmap', 'webbrowser', 'antigravity',
'ctypes', '_ctypes', '_winapi', 'posix', 'grp', 'pwd']:
_sys.modules[_mod] = None
del _sys, _builtins
"""
return banner + code
|