""" title: "Secure Code Executor" author: Wes Caldwell description: "Executes Python and Shell code in isolated environments with safety controls" required_open_webui_version: "0.5.0" requirements: version: "1.3" license: "MIT" """ import asyncio import os import re import shlex import subprocess import tempfile import logging import resource from typing import Optional, Dict, Callable, Awaitable, List from pydantic import BaseModel, Field logging.basicConfig(level=logging.INFO) class Tools: class Valves(BaseModel): max_timeout: int = Field( default=30, ge=1, le=300, description="Max execution time in seconds" ) max_memory: int = Field( default=256, ge=64, le=4096, description="Max memory in MB" ) allowed_commands: List[str] = Field( default=["ls", "cat", "grep", "python", "bash"], description="Allowed shell commands", ) deny_commands: List[str] = Field( default=["shutdown", "rm", "mv", "dd"], description="Always blocked shell commands", ) use_temp_env: bool = Field( default=True, description="Use temporary isolated environments for Python" ) class UserValves(BaseModel): allow_file_ops: bool = Field( default=True, description="Allow file operations in shell commands" ) max_code_length: int = Field( default=2048, ge=256, le=8192, description="Maximum allowed code length" ) def __init__(self): self.valves = self.Valves() self.user_valves = self.UserValves() async def _execute_process( self, command: List[str], env: Dict[str, str], timeout: int, __event_emitter__: Optional[Callable[[dict], Awaitable[None]]] = None, ) -> str: """Executes process safely with resource limits""" output = [] process = None try: def set_limits(): """Sets CPU & memory limits for execution""" resource.setrlimit( resource.RLIMIT_AS, ( self.valves.max_memory * 1024 * 1024, self.valves.max_memory * 1024 * 1024, ), ) resource.setrlimit(resource.RLIMIT_CPU, (timeout, timeout)) process = await asyncio.create_subprocess_exec( *command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, env=env, preexec_fn=set_limits, ) async def read_output(): while True: line = await process.stdout.readline() if not line: break decoded = line.decode().strip() output.append(decoded) if __event_emitter__: await __event_emitter__( { "type": "status", "data": {"description": decoded, "done": False}, } ) await asyncio.wait_for(read_output(), timeout=timeout) await process.wait() except asyncio.TimeoutError: if process: process.terminate() await process.wait() return "⏰ Execution timed out" except Exception as e: return f"🚨 Error: {str(e)}" exit_code = process.returncode result = "\n".join(output) return result + f"\n💥 Exit Code: {exit_code}" if exit_code != 0 else result async def execute_python_code( self, code: str, __event_emitter__: Optional[Callable[[dict], Awaitable[None]]] = None, ) -> str: """Executes Python code in an isolated environment""" if len(code) > self.user_valves.max_code_length: return "❌ Code exceeds maximum allowed length" env = os.environ.copy() env["PYTHONUNBUFFERED"] = "1" if self.valves.use_temp_env: with tempfile.TemporaryDirectory() as tmp_dir: env_path = os.path.join(tmp_dir, "venv") subprocess.run(["python3", "-m", "venv", env_path], check=True) python_exec = os.path.join(env_path, "bin", "python") return await self._execute_process( [python_exec, "-c", code], env, self.valves.max_timeout, __event_emitter__, ) else: return await self._execute_process( ["python3", "-u", "-c", code], env, self.valves.max_timeout, __event_emitter__, ) @staticmethod def extract_code(content: str) -> Optional[str]: """Extracts code from markdown-style blocks""" match = re.search(r"```(?:python|sh)(.*?)```", content, re.DOTALL) return match.group(1).strip() if match else None # Example Usage if __name__ == "__main__": import asyncio tools = Tools() async def main(): python_code = """ print("Hello, Secure Execution!") """ result = await tools.execute_python_code(python_code) print("Python Execution Result:\n", result) shell_command = "ls -la" shell_result = await tools.execute_shell_command(shell_command) print("Shell Execution Result:\n", shell_result) asyncio.run(main())