Ultimate MCP Coding Platform

"""Enhanced secure code execution tool with advanced features.""" from __future__ import annotations import asyncio import hashlib import json import logging # The resource module is POSIX-only; guard imports so Windows agents degrade gracefully. try: import resource # type: ignore except ModuleNotFoundError: # pragma: no cover - Windows environments resource = None # type: ignore[assignment] import sys import time import uuid from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path from tempfile import TemporaryDirectory from typing import Any, AsyncGenerator from pydantic import BaseModel, Field from ..database.neo4j_client_enhanced import EnhancedNeo4jClient from ..utils.enhanced_security import SecurityContext, SecurityLevel, ensure_safe_python logger = logging.getLogger(__name__) @dataclass class ExecutionLimits: """Resource limits for code execution.""" max_memory_mb: int = 128 max_cpu_time_seconds: int = 30 max_file_size_mb: int = 10 max_processes: int = 1 max_open_files: int = 64 @dataclass class ExecutionMetrics: """Execution performance metrics.""" cpu_time: float memory_peak_mb: float io_operations: int network_calls: int = 0 class ExecutionRequest(BaseModel): """Enhanced execution request with security context.""" code: str = Field(..., description="Code snippet to execute") language: str = Field(default="python", description="Programming language") timeout_seconds: float = Field(default=8.0, ge=0.5, le=60.0) limits: dict[str, Any] = Field(default_factory=dict, description="Resource limits") environment: dict[str, str] = Field(default_factory=dict, description="Environment variables") cache_enabled: bool = Field(default=True, description="Enable result caching") debug_mode: bool = Field(default=False, description="Enable debug output") class ExecutionResponse(BaseModel): """Enhanced execution response with metrics.""" id: str return_code: int stdout: str stderr: str duration_seconds: float metrics: dict[str, Any] = Field(default_factory=dict) cached: bool = False security_warnings: list[str] = Field(default_factory=list) class ExecutionCache: """In-memory cache for execution results.""" def __init__(self, max_size: int = 1000, ttl_seconds: int = 3600): self.max_size = max_size self.ttl_seconds = ttl_seconds self.cache: dict[str, tuple[ExecutionResponse, float]] = {} def _generate_key(self, request: ExecutionRequest) -> str: """Generate cache key from request.""" content = f"{request.code}:{request.language}:{request.timeout_seconds}" return hashlib.sha256(content.encode()).hexdigest()[:16] def get(self, request: ExecutionRequest) -> ExecutionResponse | None: """Get cached result if available and valid.""" if not request.cache_enabled: return None key = self._generate_key(request) if key in self.cache: result, timestamp = self.cache[key] if time.time() - timestamp < self.ttl_seconds: result.cached = True return result else: del self.cache[key] return None def set(self, request: ExecutionRequest, response: ExecutionResponse) -> None: """Cache execution result.""" if not request.cache_enabled or response.return_code != 0: return key = self._generate_key(request) # Evict oldest entries if cache is full if len(self.cache) >= self.max_size: oldest_key = min(self.cache.keys(), key=lambda k: self.cache[k][1]) del self.cache[oldest_key] self.cache[key] = (response, time.time()) class EnhancedExecutionTool: """Enhanced execution tool with advanced security and monitoring.""" def __init__(self, neo4j: EnhancedNeo4jClient): self.neo4j = neo4j self.cache = ExecutionCache() self.supported_languages = {"python", "javascript", "bash"} @staticmethod def _snapshot_usage(): if resource is None: return None try: return resource.getrusage(resource.RUSAGE_CHILDREN) except Exception: # pragma: no cover - platform-specific failure return None @staticmethod def _compute_usage_metrics(start_usage, language: str) -> dict[str, Any]: metrics: dict[str, Any] = {"language": language} if resource is None or start_usage is None: return metrics try: end_usage = resource.getrusage(resource.RUSAGE_CHILDREN) except Exception: # pragma: no cover - platform-specific failure return metrics def _delta(attr: str) -> float: return getattr(end_usage, attr, 0.0) - getattr(start_usage, attr, 0.0) cpu_time = max(0.0, _delta("ru_utime") + _delta("ru_stime")) io_ops = max(0, int(_delta("ru_inblock") + _delta("ru_oublock"))) peak_raw = max(0.0, end_usage.ru_maxrss - start_usage.ru_maxrss) if sys.platform == "darwin": memory_peak_mb = peak_raw / (1024 * 1024) else: memory_peak_mb = peak_raw / 1024 metrics.update( { "cpu_time": cpu_time, "memory_peak_mb": round(memory_peak_mb, 3), "io_operations": io_ops, } ) return metrics async def run( self, request: ExecutionRequest, security_context: SecurityContext | None = None ) -> ExecutionResponse: """Execute code with enhanced security and monitoring.""" # Check cache first cached_result = self.cache.get(request) if cached_result: logger.info(f"Returning cached result for execution {cached_result.id}") return cached_result # Validate security context if security_context and security_context.security_level == SecurityLevel.PUBLIC: if request.language != "python" or len(request.code) > 1000: raise PermissionError("Public access limited to Python code under 1000 characters") # Validate language support if request.language not in self.supported_languages: raise ValueError(f"Unsupported language: {request.language}") # Enhanced security validation security_warnings = [] if resource is None: security_warnings.append( "POSIX resource module unavailable; OS-level execution limits not enforced" ) if request.language == "python": try: ensure_safe_python(request.code, max_complexity=200) except Exception as e: security_warnings.append(str(e)) if security_context and security_context.security_level == SecurityLevel.PUBLIC: raise # Execute with appropriate handler if request.language == "python": result = await self._execute_python(request, security_warnings) elif request.language == "javascript": result = await self._execute_javascript(request, security_warnings) elif request.language == "bash": result = await self._execute_bash(request, security_warnings) else: raise ValueError(f"Execution handler not implemented for {request.language}") # Cache successful results self.cache.set(request, result) # Persist to database await self._persist_result(result, security_context) return result async def _execute_python( self, request: ExecutionRequest, security_warnings: list[str] ) -> ExecutionResponse: """Execute Python code with resource limits.""" limits = ExecutionLimits(**request.limits) async with self._create_sandbox() as sandbox_path: script_path = sandbox_path / "script.py" script_path.write_text(request.code, encoding="utf-8") # Prepare execution environment env = { "PYTHONPATH": "", "PYTHONDONTWRITEBYTECODE": "1", "PYTHONUNBUFFERED": "1", **request.environment } def _preexec() -> None: """Apply resource limits inside the child process before user code runs.""" if resource is None: return # Memory limit try: mem_limit = limits.max_memory_mb * 1024 * 1024 resource.setrlimit(resource.RLIMIT_AS, (mem_limit, mem_limit)) except (ValueError, OSError): # pragma: no cover - limit not supported pass # CPU time try: cpu_limit = limits.max_cpu_time_seconds resource.setrlimit(resource.RLIMIT_CPU, (cpu_limit, cpu_limit)) except (ValueError, OSError): # pragma: no cover pass # Process and file descriptor ceilings for limit_name, value in ( ("RLIMIT_NPROC", limits.max_processes), ("RLIMIT_NOFILE", limits.max_open_files), ): rlimit = getattr(resource, limit_name, None) if rlimit is None: continue try: resource.setrlimit(rlimit, (value, value)) except (ValueError, OSError): # pragma: no cover continue # Limit file sizes written during execution rlimit_fsize = getattr(resource, "RLIMIT_FSIZE", None) if rlimit_fsize is not None: max_bytes = limits.max_file_size_mb * 1024 * 1024 try: resource.setrlimit(rlimit_fsize, (max_bytes, max_bytes)) except (ValueError, OSError): # pragma: no cover pass cmd = [ sys.executable, "-I", # Isolated mode: ignore PYTHONPATH/user site-packages "-B", # Disable .pyc generation inside the sandbox str(script_path), ] start_time = time.perf_counter() start_rusage = self._snapshot_usage() try: kwargs: dict[str, Any] = { "stdout": asyncio.subprocess.PIPE, "stderr": asyncio.subprocess.PIPE, "cwd": sandbox_path, "env": env, } if resource is not None: kwargs["preexec_fn"] = _preexec # type: ignore[assignment] process = await asyncio.create_subprocess_exec( *cmd, **kwargs, ) stdout, stderr = await asyncio.wait_for( process.communicate(), timeout=request.timeout_seconds ) return_code = process.returncode except asyncio.TimeoutError: if process: process.kill() await process.wait() return_code = -1 stdout = b"" stderr = b"Execution timed out" except Exception as e: return_code = -1 stdout = b"" stderr = str(e).encode() duration = time.perf_counter() - start_time metrics = self._compute_usage_metrics(start_rusage, request.language) metrics["duration_seconds"] = duration metrics["limits_applied"] = limits.__dict__ metrics.setdefault("cpu_time", duration) return ExecutionResponse( id=str(uuid.uuid4()), return_code=return_code, stdout=stdout.decode("utf-8", errors="replace"), stderr=stderr.decode("utf-8", errors="replace"), duration_seconds=duration, metrics=metrics, security_warnings=security_warnings, ) async def _execute_javascript( self, request: ExecutionRequest, security_warnings: list[str] ) -> ExecutionResponse: """Execute JavaScript code using Node.js.""" async with self._create_sandbox() as sandbox_path: script_path = sandbox_path / "script.js" script_path.write_text(request.code, encoding="utf-8") cmd = ["node", "--max-old-space-size=128", str(script_path)] start_time = time.perf_counter() start_rusage = self._snapshot_usage() try: process = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=sandbox_path, ) stdout, stderr = await asyncio.wait_for( process.communicate(), timeout=request.timeout_seconds ) return_code = process.returncode except asyncio.TimeoutError: if process: process.kill() await process.wait() return_code = -1 stdout = b"" stderr = b"Execution timed out" except Exception as e: return_code = -1 stdout = b"" stderr = str(e).encode() duration = time.perf_counter() - start_time metrics = self._compute_usage_metrics(start_rusage, "javascript") metrics["duration_seconds"] = duration metrics.setdefault("cpu_time", duration) return ExecutionResponse( id=str(uuid.uuid4()), return_code=return_code, stdout=stdout.decode("utf-8", errors="replace"), stderr=stderr.decode("utf-8", errors="replace"), duration_seconds=duration, metrics=metrics, security_warnings=security_warnings, ) async def _execute_bash( self, request: ExecutionRequest, security_warnings: list[str] ) -> ExecutionResponse: """Execute Bash script with restrictions.""" # Additional security for bash if any(dangerous in request.code for dangerous in ["rm -rf", "sudo", "curl", "wget"]): security_warnings.append("Potentially dangerous bash commands detected") async with self._create_sandbox() as sandbox_path: script_path = sandbox_path / "script.sh" script_path.write_text(request.code, encoding="utf-8") script_path.chmod(0o755) cmd = ["bash", str(script_path)] start_time = time.perf_counter() start_rusage = self._snapshot_usage() try: process = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=sandbox_path, ) stdout, stderr = await asyncio.wait_for( process.communicate(), timeout=request.timeout_seconds ) return_code = process.returncode except asyncio.TimeoutError: if process: process.kill() await process.wait() return_code = -1 stdout = b"" stderr = b"Execution timed out" except Exception as e: return_code = -1 stdout = b"" stderr = str(e).encode() duration = time.perf_counter() - start_time metrics = self._compute_usage_metrics(start_rusage, "bash") metrics["duration_seconds"] = duration metrics.setdefault("cpu_time", duration) return ExecutionResponse( id=str(uuid.uuid4()), return_code=return_code, stdout=stdout.decode("utf-8", errors="replace"), stderr=stderr.decode("utf-8", errors="replace"), duration_seconds=duration, metrics=metrics, security_warnings=security_warnings, ) @asynccontextmanager async def _create_sandbox(self) -> AsyncGenerator[Path, None]: """Create secure execution sandbox.""" with TemporaryDirectory(prefix="enhanced_mcp_exec_") as tmp_dir: sandbox_path = Path(tmp_dir) # Set restrictive permissions sandbox_path.chmod(0o700) yield sandbox_path async def _persist_result( self, result: ExecutionResponse, security_context: SecurityContext | None ) -> None: """Persist execution result with enhanced metadata.""" await self.neo4j.execute_write_with_retry( """ MERGE (r:ExecutionResult {id: $id}) SET r += { return_code: $return_code, stdout: $stdout, stderr: $stderr, duration_seconds: $duration_seconds, metrics: $metrics, security_warnings: $security_warnings, cached: $cached, user_id: $user_id, security_level: $security_level, created_at: datetime($created_at) } """, { "id": result.id, "return_code": result.return_code, "stdout": result.stdout, "stderr": result.stderr, "duration_seconds": result.duration_seconds, "metrics": json.dumps(result.metrics), "security_warnings": result.security_warnings, "cached": result.cached, "user_id": security_context.user_id if security_context else None, "security_level": ( security_context.security_level.value if security_context else "public" ), "created_at": time.time(), }, ) async def get_execution_stats(self) -> dict[str, Any]: """Get execution statistics.""" stats = await self.neo4j.execute_read_with_retry( """ MATCH (r:ExecutionResult) RETURN count(r) as total_executions, avg(r.duration_seconds) as avg_duration, sum(CASE WHEN r.return_code = 0 THEN 1 ELSE 0 END) as successful_executions, sum(CASE WHEN r.cached THEN 1 ELSE 0 END) as cached_executions """ ) if stats: total_execs = stats[0]["total_executions"] success_rate = ( stats[0]["successful_executions"] / total_execs if total_execs > 0 else 0 ) cache_hit_rate = ( stats[0]["cached_executions"] / total_execs if total_execs > 0 else 0 ) return { "total_executions": total_execs, "average_duration": stats[0]["avg_duration"], "success_rate": success_rate, "cache_hit_rate": cache_hit_rate, "cache_size": len(self.cache.cache), } return {"total_executions": 0, "cache_size": len(self.cache.cache)} __all__ = ["EnhancedExecutionTool", "ExecutionRequest", "ExecutionResponse", "ExecutionLimits"]