From db75b822f43fa8cb509470942f69f0a6ca5e4f92 Mon Sep 17 00:00:00 2001 From: Hans Aschauer Date: Sat, 7 Feb 2026 10:11:59 +0100 Subject: [PATCH] Create standalone pod_executor package Created a standalone code execution package independent of MCP-Forge: Structure: - pod_executor/security/ - Resource limits, audit protocols, validation - pod_executor/containers/ - Podman client and container management - pod_executor/simple/ - Stateless code executor - pod_executor/jupyter/ - Stateful Jupyter backend with sessions Key changes: - Removed ForgeConfig dependency - all parameters explicit - Audit logger now a protocol with NullAuditLogger/SimpleFileAuditLogger - Validator now a protocol with NoOpValidator/BasicValidator - All imports updated to pod_executor namespace - Audit calls use simple strings instead of enums Benefits: - Standalone package usable without MCP-Forge - Clear separation between execution engine and MCP protocol - Easier testing and development - Reusable in other projects --- src/pod_executor/__init__.py | 82 +++ src/pod_executor/containers/__init__.py | 14 + src/pod_executor/containers/client.py | 157 +++++ src/pod_executor/containers/manager.py | 501 +++++++++++++++ src/pod_executor/jupyter/__init__.py | 14 + src/pod_executor/jupyter/backend.py | 229 +++++++ src/pod_executor/jupyter/kernel.py | 631 +++++++++++++++++++ src/pod_executor/jupyter/sessions.py | 437 +++++++++++++ src/pod_executor/security/__init__.py | 33 + src/pod_executor/security/audit.py | 132 ++++ src/pod_executor/security/resource_limits.py | 150 +++++ src/pod_executor/security/validation.py | 232 +++++++ src/pod_executor/simple/__init__.py | 8 + src/pod_executor/simple/executor.py | 222 +++++++ 14 files changed, 2842 insertions(+) create mode 100644 src/pod_executor/__init__.py create mode 100644 src/pod_executor/containers/__init__.py create mode 100644 src/pod_executor/containers/client.py create mode 100644 src/pod_executor/containers/manager.py create mode 100644 src/pod_executor/jupyter/__init__.py create mode 100644 src/pod_executor/jupyter/backend.py create mode 100644 src/pod_executor/jupyter/kernel.py create mode 100644 src/pod_executor/jupyter/sessions.py create mode 100644 src/pod_executor/security/__init__.py create mode 100644 src/pod_executor/security/audit.py create mode 100644 src/pod_executor/security/resource_limits.py create mode 100644 src/pod_executor/security/validation.py create mode 100644 src/pod_executor/simple/__init__.py create mode 100644 src/pod_executor/simple/executor.py diff --git a/src/pod_executor/__init__.py b/src/pod_executor/__init__.py new file mode 100644 index 0000000..27d3719 --- /dev/null +++ b/src/pod_executor/__init__.py @@ -0,0 +1,82 @@ +""" +Pod Executor - Standalone Python code execution in Podman containers. + +Provides stateless and stateful (Jupyter) code execution backends with +security isolation via Podman containers. + +Examples: + # Simple stateless execution + from pod_executor import CodeExecutor, ResourceLimits + from pod_executor.containers import SecureContainerManager, PodmanClient + from pod_executor.security import NoOpValidator, NullAuditLogger + + client = PodmanClient(socket_path="/run/podman/podman.sock", + validator=NoOpValidator(), + audit_logger=NullAuditLogger()) + container_manager = SecureContainerManager(client, NoOpValidator(), NullAuditLogger()) + limits = ResourceLimits(memory="512m", cpu_quota=100000, storage="1g", timeout=30) + executor = CodeExecutor(container_manager, "python:3.12", limits) + + result = executor.execute("print('Hello World')") + print(result.stdout) + + # Stateful Jupyter execution + from pod_executor import JupyterBackend + + backend = JupyterBackend( + container_manager=container_manager, + image="jupyter/base-notebook", + default_timeout=300 + ) + + result = backend.execute("x = 42", session_id="my-session") + result = backend.execute("print(x * 2)", session_id="my-session") +""" + +from pod_executor.simple.executor import CodeExecutor, ExecutionResult +from pod_executor.jupyter.backend import JupyterBackend +from pod_executor.jupyter.kernel import JupyterKernelManager, KernelError +from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError +from pod_executor.containers.manager import SecureContainerManager, ContainerConfig +from pod_executor.containers.client import PodmanClient, PodmanConnectionError +from pod_executor.security.resource_limits import ResourceLimits +from pod_executor.security.audit import ( + AuditLoggerProtocol, + NullAuditLogger, + SimpleFileAuditLogger, +) +from pod_executor.security.validation import ( + SecurityError, + OperationValidatorProtocol, + NoOpValidator, + BasicValidator, +) + +__version__ = "0.1.0" + +__all__ = [ + # Simple execution + "CodeExecutor", + "ExecutionResult", + # Jupyter execution + "JupyterBackend", + "JupyterKernelManager", + "KernelError", + "SessionManager", + "SessionState", + "SessionError", + # Container management + "SecureContainerManager", + "ContainerConfig", + "PodmanClient", + "PodmanConnectionError", + # Security + "ResourceLimits", + "AuditLoggerProtocol", + "NullAuditLogger", + "SimpleFileAuditLogger", + "SecurityError", + "OperationValidatorProtocol", + "NoOpValidator", + "BasicValidator", +] diff --git a/src/pod_executor/containers/__init__.py b/src/pod_executor/containers/__init__.py new file mode 100644 index 0000000..0649a51 --- /dev/null +++ b/src/pod_executor/containers/__init__.py @@ -0,0 +1,14 @@ +"""Container management for pod_executor.""" + +from pod_executor.containers.client import PodmanClient, PodmanConnectionError +from pod_executor.containers.manager import ( + ContainerConfig, + SecureContainerManager, +) + +__all__ = [ + "PodmanClient", + "PodmanConnectionError", + "ContainerConfig", + "SecureContainerManager", +] diff --git a/src/pod_executor/containers/client.py b/src/pod_executor/containers/client.py new file mode 100644 index 0000000..2ad3e43 --- /dev/null +++ b/src/pod_executor/containers/client.py @@ -0,0 +1,157 @@ +""" +Podman client wrapper with security validation. + +Wraps Podman API with security validation and error handling. +All container operations are validated against security policy. +""" + +import os +from pathlib import Path +from typing import Optional, Union +from podman import PodmanClient as BasePodmanClient + +from pod_executor.security.validation import OperationValidatorProtocol +from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger + + +class PodmanConnectionError(Exception): + """Raised when connection to Podman fails.""" + pass + + +class PodmanClient: + """ + Wrapper around Podman API with security validation. + + All container operations are validated against security policy + before being sent to Podman. Provides lazy connection and + proper error handling. + """ + + def __init__( + self, + socket_path: Union[str, Path], + validator: OperationValidatorProtocol, + audit_logger: Optional[AuditLoggerProtocol] = None + ): + """ + Initialize Podman client wrapper. + + Args: + socket_path: Path to Podman socket + validator: Operation validator for security checks + audit_logger: Optional audit logger (uses NullAuditLogger if None) + """ + self.socket_path = Path(socket_path) + self.validator = validator + self.audit_logger = audit_logger or NullAuditLogger() + self._client: Optional[BasePodmanClient] = None + + def connect(self) -> None: + """ + Connect to Podman via socket. + + Raises: + PodmanConnectionError: If connection fails + """ + try: + # Verify socket exists and is accessible + self.verify_socket_access() + + # Create Podman client with Unix socket + base_url = f"unix://{self.socket_path}" + self._client = BasePodmanClient(base_url=base_url) + + # Test connection with ping (only if client supports it) + if hasattr(self._client, 'ping'): + self._client.ping() + + except PodmanConnectionError: + # Re-raise our own exceptions + raise + except Exception as e: + raise PodmanConnectionError( + f"Failed to connect to Podman at {self.socket_path}: {e}" + ) from e + + def ping(self) -> bool: + """ + Test connection to Podman. + + Returns: + True if connection is healthy + + Raises: + PodmanConnectionError: If not connected or ping fails + """ + if self._client is None: + raise PodmanConnectionError("Not connected to Podman") + + try: + result = self._client.ping() + return result == "OK" or result is True + except Exception as e: + raise PodmanConnectionError(f"Ping failed: {e}") from e + + def disconnect(self) -> None: + """Disconnect from Podman and cleanup.""" + if self._client is not None: + try: + self._client.close() + except Exception: + pass # Ignore errors during cleanup + finally: + self._client = None + + def verify_socket_access(self) -> None: + """ + Verify that socket exists and is accessible. + + Raises: + PodmanConnectionError: If socket is not accessible + """ + if not self.socket_path.exists(): + raise PodmanConnectionError( + f"Podman socket not found: {self.socket_path}" + ) + + if not os.access(self.socket_path, os.R_OK): + raise PodmanConnectionError( + f"Podman socket is not readable: {self.socket_path}" + ) + + def check_api_version(self) -> dict: + """ + Get Podman API version information. + + Returns: + Dictionary with version information + + Raises: + PodmanConnectionError: If not connected + """ + if self._client is None: + raise PodmanConnectionError("Not connected to Podman") + + try: + return self._client.version() + except Exception as e: + raise PodmanConnectionError( + f"Failed to get API version: {e}" + ) from e + + @property + def client(self) -> BasePodmanClient: + """ + Get underlying Podman client (lazy connection). + + Returns: + Connected Podman client + + Raises: + PodmanConnectionError: If connection fails + """ + if self._client is None: + self.connect() + assert self._client is not None # Type narrowing for mypy + return self._client diff --git a/src/pod_executor/containers/manager.py b/src/pod_executor/containers/manager.py new file mode 100644 index 0000000..78872a2 --- /dev/null +++ b/src/pod_executor/containers/manager.py @@ -0,0 +1,501 @@ +"" +Secure container management with security enforcement. + +All container operations are validated against security policy +before being sent to Podman. Provides lifecycle management +with comprehensive audit logging. +"" + +from typing import Optional, Dict, List, Union +from datetime import datetime, timedelta +from pathlib import Path + +from pod_executor.containers.client import PodmanClient +from pod_executor.security.validation import OperationValidatorProtocol, SecurityError +from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger +from pod_executor.security.resource_limits import ResourceLimits + + +class ContainerConfig: + ""Container configuration with security defaults."" + + def __init__( + self", + image: str", + command: Optional[List[str]] = None", + environment: Optional[Dict[str, str]] = None", + volumes: Optional[Dict[str, dict]] = None", + resource_limits: Optional[ResourceLimits] = None", + working_dir: Optional[str] = None", + user: str = "1000:1000", + network_mode: str = "none", + port_bindings: Optional[Dict[str, int]] = None + ): + "" + Initialize container configuration. + + Args: + image: Container image to use + command: Command to run in container + environment: Environment variables + volumes: Volume mounts (host_path -> {bind, mode}) + resource_limits: Resource limits to apply + working_dir: Working directory in container (None to use image default) + user: User to run as (UID:GID) + network_mode: Network mode (none, host, bridge). Default is 'none' for security. + port_bindings: Port mappings for network_mode=host (container_port -> host_port) + "" + self.image = image + self.command = command or [] + self.environment = environment or {} + self.volumes = volumes or {} + self.resource_limits = resource_limits + self.working_dir = working_dir + self.user = user + self.network_mode = network_mode + self.port_bindings = port_bindings or {} + + def to_podman_params(self) -> dict: + "" + Convert to Podman container create parameters. + + Ensures all security requirements are included: + - network_mode: configurable (default 'none' for security) + - read_only: True + - security_opt: ["no-new-privileges"] + - resource limits + - port_bindings: for host networking mode + + Returns: + Dictionary of parameters for Podman containers.create() + "" + params = { + "image": self.image", + "command": self.command if self.command else None", + "environment": self.environment", + "user": self.user", + # Security requirements + "network_mode": self.network_mode", + "read_only": True", + "security_opt": ["no-new-privileges"]", + } + + # Add working_dir only if explicitly set + if self.working_dir is not None: + params["working_dir"] = self.working_dir + + # Add port bindings if using host network mode + # Note: In host mode, ports are directly accessible + # port_bindings are informational for tracking + if self.network_mode == "host" and self.port_bindings: + # With host networking, container uses host's network stack directly + # No explicit port mapping needed, but we track for documentation + pass + + # Add volumes if present + if self.volumes: + params["volumes"] = self.volumes + + # Add resource limits if present + if self.resource_limits: + limit_params = self.resource_limits.to_podman_params() + params.update(limit_params) + # Disable swap to avoid cgroup swap.max issues on some systems + if "mem_limit" in params: + params["memswap_limit"] = -1 # Disable swap + + return params + + +class SecureContainerManager: + ""Manages container lifecycle with security enforcement."" + + def __init__( + self", + podman_client: PodmanClient", + validator: OperationValidatorProtocol", + audit_logger: Optional[AuditLoggerProtocol] = None + ): + "" + Initialize secure container manager. + + Args: + podman_client: Podman client wrapper + validator: Operation validator for security checks + audit_logger: Optional audit logger (uses NullAuditLogger if None) + "" + self.podman = podman_client + self.validator = validator + self.audit_logger = audit_logger or NullAuditLogger() + + def create_container( + self", + config: ContainerConfig", + session_id: Optional[str] = None", + name: Optional[str] = None", + **extra_params + ) -> str: + "" + Create a container with security validation. + + Args: + config: Container configuration + session_id: Session ID for tracking + name: Optional container name + **extra_params: Additional parameters (checked for forbidden values) + + Returns: + Container ID + + Raises: + SecurityError: If configuration violates security policy + "" + # Convert config to Podman parameters + params = config.to_podman_params() + + # Add session label if provided + labels = {} + if session_id: + labels["mcp-forge.session"] = session_id + if labels: + params["labels"] = labels + + if name: + params["name"] = name + + # Merge any extra parameters (will be validated) + params.update(extra_params) + + # Validate against security policy + try: + # Extract image from params for validation + self.validator.validate_container_create( + image=config.image", + params=params", + session_id=session_id + ) + except SecurityError as e: + # Log security violation + self.audit_logger.log(event_type="security.violation", severity="critical", message=f"Security violation: container_create", details={"operation": "container_create", "reason": str(str(e)")}, session_id=session_id + ) + raise + + # Create container + try: + container = self.podman.client.containers.create(**params) + container_id = container.id + + # Register with validator + if session_id: + self.validator.register_session_container(container_id) + + # Log successful creation + self.audit_logger.log(event_type="container.operation", severity="info", operation="create", + container_id=container_id", + image=config.image", + session_id=session_id", + details={ + "name": name", + "command": config.command + } + ) + + return container_id + + except Exception as e: + self.audit_logger.log( + event_type="CONTAINER_CREATE", + severity="ERROR", + message=f"Container creation failed: {e}", + details={ + "image": config.image", + "session_id": session_id", + "error": str(e) + } + ) + raise + + def start_container(self, container_id: str) -> None: + "" + Start a container. + + Args: + container_id: Container ID to start + + Raises: + SecurityError: If container is not a session container + "" + # Verify container is registered (security check) + if container_id not in self.validator.session_containers: + self.audit_logger.log_security_violation( + operation="container_start", + reason=f"Attempted to start unregistered container: {container_id}" + ) + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + container.start() + + self.audit_logger.log(event_type="container.operation", severity="info", operation="start", + container_id=container_id", + image=" # Not available without extra lookup + ) + + except Exception as e: + self.audit_logger.log( + event_type="CONTAINER_START", + severity="ERROR", + message=f"Container start failed: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def stop_container( + self", + container_id: str", + timeout: int = 10 + ) -> None: + "" + Stop a container. + + Args: + container_id: Container ID to stop + timeout: Timeout in seconds + "" + # Verify container is registered + if container_id not in self.validator.session_containers: + self.audit_logger.log_security_violation( + operation="container_stop", + reason=f"Attempted to stop unregistered container: {container_id}" + ) + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + container.stop(timeout=timeout) + + self.audit_logger.log(event_type="container.operation", severity="info", operation="stop", + container_id=container_id", + image="", + details={"timeout": timeout} + ) + + except Exception as e: + self.audit_logger.log( + event_type="CONTAINER_STOP", + severity="ERROR", + message=f"Container stop failed: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def remove_container( + self", + container_id: str", + force: bool = False + ) -> None: + "" + Remove a container. + + Args: + container_id: Container ID to remove + force: Force removal even if running + "" + # Verify container is registered + if container_id not in self.validator.session_containers: + self.audit_logger.log_security_violation( + operation="container_remove", + reason=f"Attempted to remove unregistered container: {container_id}" + ) + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + container.remove(force=force) + + # Unregister from validator + self.validator.unregister_session_container(container_id) + + self.audit_logger.log(event_type="container.operation", severity="info", operation="remove", + container_id=container_id", + image="", + details={"force": force} + ) + + except Exception as e: + self.audit_logger.log( + event_type="CONTAINER_REMOVE", + severity="ERROR", + message=f"Container removal failed: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def get_container_logs( + self", + container_id: str", + tail: int = 100 + ) -> tuple[str, str]: + "" + Get container stdout and stderr logs. + + Args: + container_id: Container ID + tail: Number of lines to retrieve + + Returns: + (stdout, stderr) as strings + "" + if container_id not in self.validator.session_containers: + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + logs = container.logs(tail=tail, stdout=True, stderr=True) + + # Podman logs returns a generator of frames, need to consume it + if hasattr(logs, '__iter__') and not isinstance(logs, (str, bytes)): + # It's a generator/iterator, consume it + logs_bytes = b''.join(logs) + logs_str = logs_bytes.decode('utf-8', errors='replace') + elif isinstance(logs, bytes): + logs_str = logs.decode('utf-8', errors='replace') + else: + logs_str = str(logs) + + # For simplicity, return all logs in stdout (Podman combines them) + return logs_str, " + + except Exception as e: + self.audit_logger.log( + event_type="EXECUTION_REQUEST", + severity="ERROR", + message=f"Failed to get container logs: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def wait_for_container( + self", + container_id: str", + timeout: int = 300 + ) -> int: + "" + Wait for container to exit. + + Args: + container_id: Container ID + timeout: Timeout in seconds + + Returns: + Exit code + + Raises: + TimeoutError: If container doesn't exit within timeout + "" + if container_id not in self.validator.session_containers: + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + result = container.wait(timeout=timeout) + + # Extract exit code from result + if isinstance(result, dict): + exit_code = result.get("StatusCode", 0) + else: + exit_code = result + + return exit_code + + except Exception as e: + self.audit_logger.log( + event_type="EXECUTION_REQUEST", + severity="ERROR", + message=f"Failed to wait for container: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def cleanup_old_containers( + self", + max_age: timedelta = timedelta(hours=24) + ) -> int: + "" + Cleanup containers older than max_age. + + Args: + max_age: Maximum age for containers + + Returns: + Number of containers removed + "" + try: + # Get all containers with mcp-forge.session label + containers = self.podman.client.containers.list( + all=True", + filters={"label": ["mcp-forge.session"]} + ) + + removed_count = 0 + now = datetime.now() + + for container in containers: + # Get creation time + created_str = container.attrs.get("Created", ") + if not created_str: + continue + + # Parse ISO format timestamp + try: + # Remove fractional seconds and timezone for parsing + created_str = created_str.split('.')[0] + created = datetime.fromisoformat(created_str.replace('Z', '')) + except (ValueError, AttributeError): + continue + + age = now - created + + if age > max_age: + try: + container.remove(force=True) + removed_count += 1 + + self.audit_logger.log( + event_type="CONTAINER_REMOVE", + severity="INFO", + message=f"Cleaned up old container: {container.id}", + details={ + "container_id": container.id", + "age_hours": age.total_seconds() / 3600 + } + ) + except Exception as e: + self.audit_logger.log( + event_type="CONTAINER_REMOVE", + severity="WARNING", + message=f"Failed to remove old container: {e}", + details={"container_id": container.id, "error": str(e)} + ) + + return removed_count + + except Exception as e: + self.audit_logger.log( + event_type="CONTAINER_REMOVE", + severity="ERROR", + message=f"Cleanup failed: {e}", + details={"error": str(e)} + ) + raise diff --git a/src/pod_executor/jupyter/__init__.py b/src/pod_executor/jupyter/__init__.py new file mode 100644 index 0000000..e8d339a --- /dev/null +++ b/src/pod_executor/jupyter/__init__.py @@ -0,0 +1,14 @@ +"""Jupyter stateful code executor.""" + +from pod_executor.jupyter.backend import JupyterBackend +from pod_executor.jupyter.kernel import JupyterKernelManager, KernelError +from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError + +__all__ = [ + "JupyterBackend", + "JupyterKernelManager", + "KernelError", + "SessionManager", + "SessionState", + "SessionError", +] diff --git a/src/pod_executor/jupyter/backend.py b/src/pod_executor/jupyter/backend.py new file mode 100644 index 0000000..dd27681 --- /dev/null +++ b/src/pod_executor/jupyter/backend.py @@ -0,0 +1,229 @@ +"""Jupyter backend for stateful code execution.""" + +from typing import Optional, Dict, List +import hashlib + +from pod_executor.containers.manager import SecureContainerManager +from pod_executor.security.resource_limits import ResourceLimits, parse_memory_string +from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger +from pod_executor.simple.executor import ExecutionResult +from pod_executor.jupyter.kernel import JupyterKernelManager +from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError + + +class JupyterBackend: + """Stateful code execution backend using Jupyter kernels.""" + + def __init__( + self, + container_manager: SecureContainerManager, + image: str, + default_timeout: int = 300, + default_memory: str = "512m", + default_cpu_quota: int = 100000, + max_timeout: int = 3600, + max_memory: str = "2g", + max_cpu_quota: int = 200000, + idle_timeout: int = 3600, + max_sessions: int = 10, + resource_limits: Optional[ResourceLimits] = None, + audit_logger: Optional[AuditLoggerProtocol] = None + ): + """ + Initialize Jupyter backend. + + Args: + container_manager: Container lifecycle manager + image: Docker/Podman image with ipykernel installed + default_timeout: Default execution timeout in seconds + default_memory: Default memory limit (e.g., "512m", "1g") + default_cpu_quota: Default CPU quota (100000 = 1 CPU) + max_timeout: Maximum allowed timeout + max_memory: Maximum allowed memory + max_cpu_quota: Maximum allowed CPU quota + idle_timeout: Session idle timeout in seconds + max_sessions: Maximum concurrent sessions + resource_limits: Default resource limits for kernels + audit_logger: Optional audit logger (uses NullAuditLogger if None) + """ + self.container_manager = container_manager + self.image = image + self.default_timeout = default_timeout + self.default_memory = default_memory + self.default_cpu_quota = default_cpu_quota + self.max_timeout = max_timeout + self.max_memory = max_memory + self.max_cpu_quota = max_cpu_quota + self.idle_timeout = idle_timeout + self.max_sessions = max_sessions + self.audit_logger = audit_logger or NullAuditLogger() + + # Initialize kernel manager + kernel_manager = JupyterKernelManager( + container_manager=container_manager, + image=image, + resource_limits=resource_limits + ) + + # Initialize session manager + self.session_manager = SessionManager( + kernel_manager=kernel_manager, + idle_timeout=idle_timeout, + max_sessions=max_sessions, + audit_logger=self.audit_logger + ) + + def execute( + self, + code: str, + session_id: str, + timeout: Optional[int] = None, + memory: Optional[str] = None, + cpu_quota: Optional[int] = None, + volumes: Optional[Dict[str, dict]] = None, + injection_code: Optional[str] = None, + bridge_socket_path: Optional[str] = None + ) -> ExecutionResult: + """ + Execute code in stateful session. + + Creates session if it doesn't exist, reuses existing session otherwise. + Session maintains namespace state across multiple executions. + + Args: + code: Python code to execute + session_id: Unique session identifier + timeout: Max execution time in seconds (uses default if None) + memory: Memory limit string (uses default if None) + cpu_quota: CPU quota (uses default if None) + volumes: Volume mounts dict + injection_code: Optional code executed once at session start + bridge_socket_path: Optional path to socket for mounting + + Returns: + ExecutionResult with execution output and metadata + + Raises: + ValueError: If limits exceed configured maximums + SessionError: If session operation fails + """ + # Use defaults if not specified + timeout = timeout if timeout is not None else self.default_timeout + memory = memory if memory is not None else self.default_memory + cpu_quota = cpu_quota if cpu_quota is not None else self.default_cpu_quota + + # Validate limits against maximums + self._validate_limits(timeout, memory, cpu_quota) + + # Log execution (hash code, don't log actual content) + code_hash = hashlib.sha256(code.encode()).hexdigest() + self.audit_logger.log( + event_type="execution.request", + severity="info", + message="Stateful code execution requested", + session_id=session_id, + details={ + "code_hash": code_hash, + "timeout": timeout, + "memory": memory, + "cpu_quota": cpu_quota + } + ) + + # Check if session exists, create if needed + try: + self.session_manager.get_session(session_id) + except SessionError: + # Session doesn't exist, create it + resource_limits = ResourceLimits( + memory=memory, + cpu_quota=cpu_quota, + storage="1g", + timeout=timeout + ) + + self.session_manager.create_session( + session_id=session_id, + resource_limits=resource_limits, + volumes=volumes, + injection_code=injection_code, + bridge_socket_path=bridge_socket_path + ) + + # Execute in session + result = self.session_manager.execute_in_session( + session_id=session_id, + code=code, + timeout=timeout + ) + + return result + + def document_state( + self, + session_id: str, + variables: Dict[str, str], + note: str = "", + clear: bool = False + ) -> dict: + """ + Document important variables in session. + + Args: + session_id: Session to document + variables: Dictionary of variable_name -> description + note: Optional note about session state + clear: If True, replace all documented variables; if False, merge + + Returns: + Dictionary with updated state info + """ + self.session_manager.document_state( + session_id=session_id, + variables=variables, + note=note, + clear=clear + ) + + state = self.session_manager.get_session_state(session_id) + return state.to_dict() + + def get_session_state(self, session_id: str) -> SessionState: + """Get documented state for session.""" + return self.session_manager.get_session_state(session_id) + + def destroy_session(self, session_id: str) -> None: + """Destroy session and cleanup kernel.""" + self.session_manager.destroy_session(session_id) + + def list_sessions(self) -> List[dict]: + """List all active sessions with metadata.""" + return self.session_manager.list_sessions() + + def cleanup_idle_sessions(self) -> int: + """Cleanup sessions idle beyond configured timeout.""" + return self.session_manager.cleanup_idle_sessions() + + def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None: + """ + Validate resource limits against configured maximums. + + Raises: + ValueError: If any limit exceeds maximum + """ + if timeout > self.max_timeout: + raise ValueError( + f"Timeout {timeout} exceeds maximum {self.max_timeout}" + ) + + memory_bytes = parse_memory_string(memory) + max_memory_bytes = parse_memory_string(self.max_memory) + if memory_bytes > max_memory_bytes: + raise ValueError( + f"Memory {memory} exceeds maximum {self.max_memory}" + ) + + if cpu_quota > self.max_cpu_quota: + raise ValueError( + f"CPU quota {cpu_quota} exceeds maximum {self.max_cpu_quota}" + ) diff --git a/src/pod_executor/jupyter/kernel.py b/src/pod_executor/jupyter/kernel.py new file mode 100644 index 0000000..55a1001 --- /dev/null +++ b/src/pod_executor/jupyter/kernel.py @@ -0,0 +1,631 @@ +""" +Real Jupyter kernel management for stateful execution. + +This module implements proper Jupyter kernel management: +- jupyter-client runs on host (MCP-Forge server) +- ipykernel runs inside Podman containers +- Communication via ZMQ protocol +- 1:1 mapping: one container per session, one kernel per container +""" + +from typing import Dict, Optional, List, Any +from dataclasses import dataclass +from datetime import datetime, timedelta +import uuid +import json +import tempfile +import time +import socket +from pathlib import Path + +from jupyter_client.blocking.client import BlockingKernelClient +import zmq + +from pod_executor.containers.manager import SecureContainerManager, ContainerConfig +from pod_executor.security.resource_limits import ResourceLimits +from pod_executor.simple.executor import ExecutionResult + + +class KernelError(Exception): + """Raised when kernel operations fail.""" + pass + + +@dataclass +class KernelInfo: + """Information about a running kernel.""" + kernel_id: str + container_id: str + session_id: str + connection_file: Path + connection_info: Dict[str, Any] # ZMQ ports and keys + started_at: datetime + last_activity: datetime + client: Optional[BlockingKernelClient] = None + + def to_dict(self) -> dict: + """Convert to dictionary for JSON serialization.""" + return { + "kernel_id": self.kernel_id, + "container_id": self.container_id, + "session_id": self.session_id, + "started_at": self.started_at.isoformat(), + "last_activity": self.last_activity.isoformat(), + "connection_info": { + "shell_port": self.connection_info.get("shell_port"), + "iopub_port": self.connection_info.get("iopub_port"), + "stdin_port": self.connection_info.get("stdin_port"), + "control_port": self.connection_info.get("control_port"), + "hb_port": self.connection_info.get("hb_port"), + } + } + + +class JupyterKernelManager: + """ + Manages IPython kernels in containers via jupyter-client. + + Architecture: + - This class runs on host (MCP-Forge server process) + - Creates one container per session with ipykernel running inside + - Connects to kernel via ZMQ protocol (jupyter-client) + - Communicates using Jupyter message protocol + + Each session gets: + - Dedicated container + - Dedicated kernel process + - Isolated Python namespace + - Independent resource limits + """ + + def __init__( + self, + container_manager: SecureContainerManager, + image: str, + resource_limits: Optional[ResourceLimits] = None + ): + """ + Initialize kernel manager. + + Args: + container_manager: Container lifecycle manager + image: Docker/Podman image with ipykernel installed + resource_limits: Default resource limits for kernels + """ + self.container_manager = container_manager + self.image = image + self.resource_limits = resource_limits + self.kernels: Dict[str, KernelInfo] = {} + + def start_kernel( + self, + session_id: str, + volumes: Optional[Dict[str, dict]] = None, + injection_code: Optional[str] = None, + bridge_socket_path: Optional[str] = None + ) -> str: + """ + Start IPython kernel in dedicated container. + + Process: + 1. Generate ZMQ connection info (ports, keys) + 2. Create connection file + 3. Create container with ipykernel command + 4. Mount bridge socket if provided (for MCP tools) + 5. Start container + 6. Wait for kernel to be ready + 7. Connect jupyter-client to kernel via ZMQ + 8. Execute injection code (MCP tools setup) if provided + 9. Verify kernel is responsive + + Args: + session_id: Session ID this kernel belongs to + volumes: Optional volume mounts + injection_code: Optional MCP tool injection code to execute at startup + bridge_socket_path: Optional path to MCP bridge socket for mounting + + Returns: + kernel_id: Unique identifier for this kernel + + Raises: + KernelError: If kernel startup fails + """ + kernel_id = f"kernel-{uuid.uuid4().hex[:16]}" + + # Generate connection info + connection_info = self._generate_connection_info() + + # Create connection file + connection_file = self._create_connection_file(kernel_id, connection_info) + + try: + # Set up volumes (user volumes + bridge socket + connection file) + container_volumes = volumes.copy() if volumes else {} + if bridge_socket_path: + container_volumes[bridge_socket_path] = { + "bind": bridge_socket_path, + "mode": "rw" + } + + # Mount connection file into container + container_connection_path = f"/tmp/kernel-{kernel_id}.json" + container_volumes[str(connection_file)] = { + "bind": container_connection_path, + "mode": "ro" + } + + # Create container with ipykernel using host networking + config = ContainerConfig( + image=self.image, + command=[ + "python", "-m", "ipykernel_launcher", + "-f", container_connection_path + ], + resource_limits=self.resource_limits, + volumes=container_volumes, + network_mode="host", # Use host network for ZMQ communication + port_bindings={ + connection_info["shell_port"]: connection_info["shell_port"], + connection_info["iopub_port"]: connection_info["iopub_port"], + connection_info["stdin_port"]: connection_info["stdin_port"], + connection_info["control_port"]: connection_info["control_port"], + connection_info["hb_port"]: connection_info["hb_port"], + } + ) + + container_id = self.container_manager.create_container( + config, + session_id=session_id, + name=f"jupyter-{kernel_id}" + ) + + # Start container + self.container_manager.start_container(container_id) + + # Wait for kernel to be ready with polling + if not self._wait_for_kernel_ready(connection_info, timeout=30): + raise KernelError(f"Kernel {kernel_id} failed to start within timeout") + + # Connect client + client = self._connect_client(connection_info) + + # Verify kernel is responsive + if not self._verify_kernel(client): + raise KernelError(f"Kernel {kernel_id} not responsive") + + # Execute injection code if provided (MCP tools setup) + if injection_code: + self._execute_injection_code(client, injection_code, kernel_id) + + # Register kernel + now = datetime.utcnow() + kernel_info = KernelInfo( + kernel_id=kernel_id, + container_id=container_id, + session_id=session_id, + connection_file=connection_file, + connection_info=connection_info, + started_at=now, + last_activity=now, + client=client + ) + self.kernels[kernel_id] = kernel_info + + return kernel_id + + except Exception as e: + # Cleanup on failure + connection_file.unlink(missing_ok=True) + raise KernelError(f"Failed to start kernel: {e}") from e + + def execute_code( + self, + kernel_id: str, + code: str, + timeout: int = 300 + ) -> ExecutionResult: + """ + Execute code in kernel via ZMQ. + + Uses jupyter-client to: + 1. Send execute_request message + 2. Receive stream (stdout/stderr) messages + 3. Receive execute_result/display_data messages + 4. Collect and parse all output + + Args: + kernel_id: Kernel to execute in + code: Python code to execute + timeout: Maximum execution time in seconds + + Returns: + ExecutionResult with stdout, stderr, result + + Raises: + KernelError: If kernel not found or execution fails + """ + kernel_info = self._get_kernel(kernel_id) + client = kernel_info.client + + if not client: + raise KernelError(f"Kernel {kernel_id} has no connected client") + + start_time = time.time() + + try: + # Execute code + _msg_id = client.execute(code, silent=False, store_history=True) + + # Collect output + stdout_parts = [] + stderr_parts = [] + result = None + has_error = False + + # Wait for execution to complete + while True: + try: + msg = client.get_iopub_msg(timeout=timeout) + msg_type = msg['header']['msg_type'] + content = msg['content'] + + if msg_type == 'stream': + if content['name'] == 'stdout': + stdout_parts.append(content['text']) + elif content['name'] == 'stderr': + stderr_parts.append(content['text']) + + elif msg_type == 'execute_result': + result = content.get('data', {}).get('text/plain', '') + + elif msg_type == 'error': + has_error = True + stderr_parts.append('\n'.join(content['traceback'])) + + elif msg_type == 'status': + if content['execution_state'] == 'idle': + break + + except zmq.error.Again: + break + + execution_time = time.time() - start_time + + # Update last activity + kernel_info.last_activity = datetime.utcnow() + + stderr_text = ''.join(stderr_parts) + + return ExecutionResult( + success=(not has_error), + stdout=''.join(stdout_parts), + stderr=stderr_text, + result=result, + execution_time=execution_time, + exit_code=1 if has_error else 0, + error=stderr_text if has_error else None + ) + + except Exception as e: + execution_time = time.time() - start_time + return ExecutionResult( + success=False, + stdout='', + stderr=str(e), + result=None, + execution_time=execution_time, + exit_code=1, + error=str(e) + ) + + def shutdown_kernel(self, kernel_id: str) -> None: + """ + Shutdown kernel and cleanup container. + + 1. Send shutdown_request via ZMQ + 2. Wait for kernel shutdown + 3. Stop and remove container + 4. Cleanup connection file + + Args: + kernel_id: Kernel to shutdown + """ + kernel_info = self._get_kernel(kernel_id) + + try: + # Shutdown kernel + if kernel_info.client: + kernel_info.client.shutdown() + kernel_info.client.stop_channels() + + # Stop and remove container + self.container_manager.stop_container(kernel_info.container_id) + self.container_manager.remove_container(kernel_info.container_id) + + # Cleanup connection file + kernel_info.connection_file.unlink(missing_ok=True) + + finally: + # Remove from registry + del self.kernels[kernel_id] + + def inspect_namespace(self, kernel_id: str) -> List[str]: + """ + Get list of variables in kernel namespace. + + Executes introspection code: + [var for var in dir() if not var.startswith('_')] + + Args: + kernel_id: Kernel to inspect + + Returns: + List of variable names + """ + code = "[var for var in dir() if not var.startswith('_')]" + result = self.execute_code(kernel_id, code, timeout=5) + + if result.success and result.result: + # Parse result (it's a string representation of a list) + try: + return eval(result.result) # nosec - controlled code + except Exception: + return [] + return [] + + def get_variable_info( + self, + kernel_id: str, + variable_name: str + ) -> Dict[str, Any]: + """ + Get detailed information about a variable. + + Executes introspection code to get: + - type(var).__name__ + - sys.getsizeof(var) if available + - var.shape if hasattr(var, 'shape') + - repr(var)[:100] + + Args: + kernel_id: Kernel to inspect + variable_name: Name of variable to inspect + + Returns: + Dict with type, size, shape, repr + """ + code = f""" +import sys +_var = {variable_name} +_info = {{ + 'type': type(_var).__name__, + 'repr': repr(_var)[:100], +}} +try: + _info['size_bytes'] = sys.getsizeof(_var) +except: + pass +if hasattr(_var, 'shape'): + _info['shape'] = _var.shape +_info +""" + result = self.execute_code(kernel_id, code, timeout=5) + + if result.success and result.result: + try: + return eval(result.result) # nosec - controlled code + except Exception: + return {} + return {} + + def restart_kernel(self, kernel_id: str) -> None: + """ + Restart kernel (namespace reset, container kept). + + Strategy: shutdown current kernel and start new one in same container. + Note: In a full implementation, we'd use KernelManager.restart_kernel(). + + Args: + kernel_id: Kernel to restart + """ + kernel_info = self._get_kernel(kernel_id) + + # For now, just record activity - full restart implementation requires + # KernelManager integration (not just BlockingKernelClient) + # TODO: Implement proper kernel restart via KernelManager + kernel_info.last_activity = datetime.utcnow() + + def cleanup_idle_kernels( + self, + idle_timeout: timedelta + ) -> int: + """ + Cleanup kernels idle longer than timeout. + + Args: + idle_timeout: Maximum idle time before cleanup + + Returns: + Number of kernels cleaned up + """ + now = datetime.utcnow() + cleaned_up = 0 + + for kernel_id in list(self.kernels.keys()): + kernel_info = self.kernels[kernel_id] + idle_time = now - kernel_info.last_activity + + if idle_time > idle_timeout: + try: + self.shutdown_kernel(kernel_id) + cleaned_up += 1 + except Exception: + pass # Continue cleanup even if one fails + + return cleaned_up + + def _get_kernel(self, kernel_id: str) -> KernelInfo: + """Get kernel info or raise error.""" + if kernel_id not in self.kernels: + raise KernelError(f"Kernel {kernel_id} not found") + return self.kernels[kernel_id] + + def _generate_connection_info(self) -> Dict[str, Any]: + """Generate ZMQ connection information with allocated ports.""" + import secrets + + # Allocate 5 ports for ZMQ channels + ports = self._allocate_ports(5) + + return { + "shell_port": ports[0], + "iopub_port": ports[1], + "stdin_port": ports[2], + "control_port": ports[3], + "hb_port": ports[4], + "ip": "127.0.0.1", + "key": secrets.token_hex(32), + "transport": "tcp", + "signature_scheme": "hmac-sha256", + "kernel_name": "python3" + } + + def _allocate_ports(self, count: int) -> List[int]: + """ + Allocate available ports for ZMQ. + + Args: + count: Number of ports to allocate + + Returns: + List of allocated port numbers + """ + ports = [] + for _ in range(count): + # Let OS assign available port + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(('127.0.0.1', 0)) # Bind to any available port + port = sock.getsockname()[1] + sock.close() + ports.append(port) + return ports + + def _create_connection_file( + self, + kernel_id: str, + connection_info: Dict[str, Any] + ) -> Path: + """Create connection file for kernel.""" + # Create temp file + fd, path = tempfile.mkstemp(suffix=f"-kernel-{kernel_id}.json") + + # Write connection info + with open(fd, 'w') as f: + json.dump(connection_info, f) + + return Path(path) + + def _connect_client(self, connection_info: Dict[str, Any]) -> BlockingKernelClient: + """Connect jupyter-client to kernel.""" + client = BlockingKernelClient() + client.load_connection_info(connection_info) + client.start_channels() + return client + + def _verify_kernel(self, client: BlockingKernelClient, timeout: int = 10) -> bool: + """Verify kernel is responsive.""" + try: + client.wait_for_ready(timeout=timeout) + return True + except Exception: + return False + + def _wait_for_kernel_ready( + self, + connection_info: Dict[str, Any], + timeout: int = 30, + poll_interval: float = 0.5 + ) -> bool: + """ + Wait for kernel to be ready by polling ports. + + Args: + connection_info: Kernel connection information + timeout: Maximum time to wait in seconds + poll_interval: Time between polls in seconds + + Returns: + True if kernel is ready, False if timeout + """ + start_time = time.time() + shell_port = connection_info["shell_port"] + + while time.time() - start_time < timeout: + try: + # Try to connect to shell port + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(1) + result = sock.connect_ex(('127.0.0.1', shell_port)) + sock.close() + + if result == 0: + # Port is open, kernel is ready + return True + except Exception: + pass + + time.sleep(poll_interval) + + return False + + def _execute_injection_code( + self, + client: BlockingKernelClient, + injection_code: str, + kernel_id: str + ) -> None: + """ + Execute MCP tool injection code on kernel startup. + + This runs once when the kernel starts to set up MCP tools. + Unlike regular code execution, we don't capture output. + + Args: + client: Connected kernel client + injection_code: Python code to inject (MCP tools setup) + kernel_id: Kernel ID for error messages + + Raises: + KernelError: If injection code fails to execute + """ + try: + # Execute injection code silently + _msg_id = client.execute(injection_code, silent=True, store_history=False) + + # Wait for execution to complete + timeout = 10 # Injection should be fast + while True: + try: + msg = client.get_iopub_msg(timeout=timeout) + msg_type = msg['header']['msg_type'] + + if msg_type == 'error': + content = msg['content'] + error_msg = '\n'.join(content.get('traceback', [str(content)])) + raise KernelError( + f"MCP injection failed in kernel {kernel_id}: {error_msg}" + ) + + elif msg_type == 'status': + if msg['content']['execution_state'] == 'idle': + break # Injection complete + + except zmq.error.Again: + break # Timeout, assume success + + except KernelError: + raise + except Exception as e: + raise KernelError( + f"Failed to execute MCP injection code in kernel {kernel_id}: {e}" + ) from e diff --git a/src/pod_executor/jupyter/sessions.py b/src/pod_executor/jupyter/sessions.py new file mode 100644 index 0000000..34ccf8c --- /dev/null +++ b/src/pod_executor/jupyter/sessions.py @@ -0,0 +1,437 @@ +"""Session management for stateful execution.""" + +from typing import Dict, Optional, List, Any +from dataclasses import dataclass, field +from datetime import datetime, timedelta + +from pod_executor.jupyter.kernel import JupyterKernelManager +from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger +from pod_executor.security.resource_limits import ResourceLimits +from pod_executor.simple.executor import ExecutionResult + + +class SessionError(Exception): + """Raised when session operations fail.""" + pass + + +@dataclass +class SessionState: + """Documented state for a session.""" + session_id: str + documented_variables: Dict[str, str] = field(default_factory=dict) + note: str = "" + last_updated: datetime = field(default_factory=datetime.utcnow) + all_variables: List[str] = field(default_factory=list) + introspection: Dict[str, dict] = field(default_factory=dict) + + def to_dict(self) -> dict: + """Convert to dictionary for JSON serialization.""" + return { + "session_id": self.session_id, + "documented_variables": self.documented_variables, + "note": self.note, + "last_updated": self.last_updated.isoformat(), + "all_variables": self.all_variables, + "introspection": self.introspection + } + + +class Session: + """Stateful execution session.""" + + def __init__( + self, + session_id: str, + kernel_id: str, + created_at: datetime, + resource_limits: ResourceLimits + ): + """ + Initialize session. + + Args: + session_id: Unique session identifier + kernel_id: ID of associated kernel + created_at: Session creation timestamp + resource_limits: Resource limits for this session + """ + self.session_id = session_id + self.kernel_id = kernel_id + self.created_at = created_at + self.last_activity = created_at + self.resource_limits = resource_limits + self.state = SessionState(session_id=session_id) + self.documented_variables: Dict[str, str] = {} + self.documentation_note: Optional[str] = None + + def update_activity(self) -> None: + """Update last activity timestamp.""" + self.last_activity = datetime.utcnow() + + def is_idle(self, timeout: timedelta) -> bool: + """ + Check if session is idle beyond timeout. + + Args: + timeout: Maximum idle time + + Returns: + True if session has been idle longer than timeout + """ + now = datetime.utcnow() + idle_time = now - self.last_activity + return idle_time > timeout + + def to_dict(self) -> dict: + """Convert to dictionary for serialization.""" + return { + "session_id": self.session_id, + "kernel_id": self.kernel_id, + "created_at": self.created_at.isoformat(), + "last_activity": self.last_activity.isoformat(), + "state": self.state.to_dict() + } + + +class SessionManager: + """Manages stateful execution sessions.""" + + def __init__( + self, + kernel_manager: JupyterKernelManager, + idle_timeout: int = 3600, + max_sessions: int = 10, + audit_logger: Optional[AuditLoggerProtocol] = None + ): + """ + Initialize session manager. + + Args: + kernel_manager: Kernel lifecycle manager + idle_timeout: Session idle timeout in seconds (default: 3600) + max_sessions: Maximum concurrent sessions (default: 10) + audit_logger: Optional audit logger (uses NullAuditLogger if None) + """ + self.kernel_manager = kernel_manager + self.idle_timeout = idle_timeout + self.max_sessions = max_sessions + self.audit_logger = audit_logger or NullAuditLogger() + self.sessions: Dict[str, Session] = {} + + def create_session( + self, + session_id: str, + resource_limits: ResourceLimits, + volumes: Optional[Dict[str, dict]] = None, + injection_code: Optional[str] = None, + bridge_socket_path: Optional[str] = None + ) -> Session: + """ + Create new stateful session. + + Args: + session_id: Unique identifier for session + resource_limits: Resource limits for session + volumes: Optional volume mounts + injection_code: Optional MCP tool injection code to execute at startup + bridge_socket_path: Optional path to MCP bridge socket for mounting + + Returns: + Created Session object + + Raises: + SessionError: If session_id already exists + SessionError: If max concurrent sessions exceeded + """ + if session_id in self.sessions: + raise SessionError(f"Session {session_id} already exists") + + # Check max concurrent limit + self._enforce_max_concurrent() + + # Start kernel with MCP injection if provided + kernel_id = self.kernel_manager.start_kernel( + session_id, + volumes=volumes, + injection_code=injection_code, + bridge_socket_path=bridge_socket_path + ) + + # Create session + now = datetime.utcnow() + session = Session( + session_id=session_id, + kernel_id=kernel_id, + created_at=now, + resource_limits=resource_limits + ) + + self.sessions[session_id] = session + + # Log session creation + self.audit_logger.log( + event_type="SESSION_CREATE, + severity="INFO, + message=f"Session created: {session_id}", + session_id=session_id, + details={ + "kernel_id": kernel_id, + "memory": resource_limits.memory_bytes, + "cpu_quota": resource_limits.cpu_quota + } + ) + + return session + + def session_exists(self, session_id: str) -> bool: + """ + Check if session exists. + + Args: + session_id: Session identifier + + Returns: + True if session exists, False otherwise + """ + return session_id in self.sessions + + def get_session(self, session_id: str) -> Session: + """ + Get session by ID. + + Args: + session_id: Session identifier + + Returns: + Session object + + Raises: + SessionError: If session doesn't exist + """ + if session_id not in self.sessions: + raise SessionError(f"Session {session_id} not found") + + return self.sessions[session_id] + + async def document_variables( + self, + session_id: str, + variables: Dict[str, str], + note: Optional[str] = None, + clear: bool = False + ) -> Dict: + """ + Document important variables in a session. + + Args: + session_id: Session identifier + variables: Dict mapping variable names to descriptions + note: Optional general note about session state + clear: Whether to clear existing documentation first + + Returns: + Result dict with success status and documented count + """ + session = self.get_session(session_id) + + if clear: + session.documented_variables = {} + + # Store variable documentation in session + if not hasattr(session, 'documented_variables'): + session.documented_variables = {} + + session.documented_variables.update(variables) + + if note: + session.documentation_note = note + + return { + "success": True, + "documented_count": len(variables), + "total_documented": len(session.documented_variables) + } + + def execute_in_session( + self, + session_id: str, + code: str, + timeout: int = 300 + ) -> ExecutionResult: + """ + Execute code in session kernel. + + Args: + session_id: Session to execute in + code: Python code to execute + timeout: Maximum execution time + + Returns: + ExecutionResult with output + + Raises: + SessionError: If session doesn't exist + """ + session = self.get_session(session_id) + + # Update activity + session.update_activity() + + # Execute in kernel + result = self.kernel_manager.execute_code( + session.kernel_id, + code, + timeout=timeout + ) + + return result + + def document_state( + self, + session_id: str, + variables: Dict[str, str], + note: str = "", + clear: bool = False + ) -> None: + """ + Document important variables in session. + + Updates session.state with variable descriptions and runs + introspection to capture current namespace state. + + Args: + session_id: Session to document + variables: Dictionary of variable_name -> description + note: Optional note about session state + clear: If True, replace all documented variables; if False, merge + + Raises: + SessionError: If session doesn't exist + """ + session = self.get_session(session_id) + + # Update documented variables + if clear: + session.state.documented_variables = variables.copy() + else: + session.state.documented_variables.update(variables) + + # Update note if provided + if note: + session.state.note = note + + # Run introspection to get current namespace state + session.state.all_variables = self.kernel_manager.inspect_namespace(session.kernel_id) + + # Get variable info for documented variables + session.state.introspection = {} + for var_name in variables.keys(): + if var_name in session.state.all_variables: + try: + info = self.kernel_manager.get_variable_info(session.kernel_id, var_name) + session.state.introspection[var_name] = info + except Exception: + # Variable might not exist yet + pass + + # Update timestamp + session.state.last_updated = datetime.utcnow() + session.update_activity() + + def get_session_state(self, session_id: str) -> SessionState: + """ + Get documented state for session. + + Args: + session_id: Session identifier + + Returns: + SessionState object + + Raises: + SessionError: If session doesn't exist + """ + session = self.get_session(session_id) + return session.state + + def destroy_session(self, session_id: str) -> None: + """ + Destroy session and cleanup kernel. + + Args: + session_id: Session to destroy + + Raises: + SessionError: If session doesn't exist + """ + session = self.get_session(session_id) + + # Shutdown kernel + try: + self.kernel_manager.shutdown_kernel(session.kernel_id) + except Exception as e: + # Log but continue with cleanup + self.audit_logger.log( + event_type="SESSION_DESTROY, + severity="WARNING, + message=f"Error shutting down kernel for session {session_id}", + session_id=session_id, + error=str(e) + ) + + # Remove session + del self.sessions[session_id] + + # Log destruction + self.audit_logger.log( + event_type="SESSION_DESTROY, + severity="INFO, + message=f"Session destroyed: {session_id}", + session_id=session_id + ) + + def cleanup_idle_sessions(self) -> int: + """ + Cleanup sessions idle beyond configured timeout. + + Returns: + Number of sessions cleaned up + """ + timeout = timedelta(seconds=self.idle_timeout) + sessions_to_remove = [] + + for session_id, session in self.sessions.items(): + if session.is_idle(timeout): + sessions_to_remove.append(session_id) + + # Destroy idle sessions + for session_id in sessions_to_remove: + try: + self.destroy_session(session_id) + except Exception: + # Best effort cleanup + pass + + return len(sessions_to_remove) + + def list_sessions(self) -> List[dict]: + """ + List all active sessions with metadata. + + Returns: + List of session dictionaries + """ + return [session.to_dict() for session in self.sessions.values()] + + def _enforce_max_concurrent(self) -> None: + """ + Enforce max concurrent sessions limit. + + Raises: + SessionError: If at max concurrent sessions + """ + if len(self.sessions) >= self.max_sessions: + raise SessionError( + f"Maximum concurrent sessions ({self.max_sessions}) reached" + ) diff --git a/src/pod_executor/security/__init__.py b/src/pod_executor/security/__init__.py new file mode 100644 index 0000000..2aee95f --- /dev/null +++ b/src/pod_executor/security/__init__.py @@ -0,0 +1,33 @@ +"""Security components for pod_executor.""" + +from pod_executor.security.resource_limits import ( + ResourceLimits, + parse_memory_string, + parse_cpu_quota, + parse_storage_string, +) +from pod_executor.security.audit import ( + AuditLoggerProtocol, + NullAuditLogger, + SimpleFileAuditLogger, +) +from pod_executor.security.validation import ( + SecurityError, + OperationValidatorProtocol, + NoOpValidator, + BasicValidator, +) + +__all__ = [ + "ResourceLimits", + "parse_memory_string", + "parse_cpu_quota", + "parse_storage_string", + "AuditLoggerProtocol", + "NullAuditLogger", + "SimpleFileAuditLogger", + "SecurityError", + "OperationValidatorProtocol", + "NoOpValidator", + "BasicValidator", +] diff --git a/src/pod_executor/security/audit.py b/src/pod_executor/security/audit.py new file mode 100644 index 0000000..5fca46f --- /dev/null +++ b/src/pod_executor/security/audit.py @@ -0,0 +1,132 @@ +""" +Audit logger protocol for pod_executor. + +Provides a protocol (interface) for audit logging that can be implemented +by consuming applications. A default no-op implementation is provided. +""" + +from typing import Protocol, Any, Optional, Dict +from pathlib import Path + + +class AuditLoggerProtocol(Protocol): + """Protocol for audit logging (optional dependency).""" + + def log( + self, + event_type: str, + severity: str, + message: str, + session_id: Optional[str] = None, + user_id: Optional[str] = None, + details: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + **kwargs: Any + ) -> None: + """ + Log an audit event. + + Args: + event_type: Type of event (e.g., "container.create", "execution.request") + severity: Severity level ("info", "warning", "error", "critical") + message: Human-readable message describing the event + session_id: Optional session ID associated with event + user_id: Optional user ID associated with event + details: Optional dictionary of additional details + error: Optional error message if event represents an error + **kwargs: Additional keyword arguments for extensibility + """ + ... + + +class NullAuditLogger: + """No-op audit logger for standalone usage without audit requirements.""" + + def log( + self, + event_type: str = "", + severity: str = "info", + message: str = "", + session_id: Optional[str] = None, + user_id: Optional[str] = None, + details: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + **kwargs: Any + ) -> None: + """Do nothing - audit logging disabled.""" + pass + + +class SimpleFileAuditLogger: + """ + Simple file-based audit logger for basic use cases. + + Logs events to a JSON Lines file (one JSON object per line). + Thread-safe via file locking. + """ + + def __init__(self, log_path: Path): + """ + Initialize file audit logger. + + Args: + log_path: Path to audit log file + """ + self.log_path = Path(log_path) + self._ensure_log_file() + + def _ensure_log_file(self) -> None: + """Ensure log file and directory exist.""" + self.log_path.parent.mkdir(parents=True, exist_ok=True) + if not self.log_path.exists(): + self.log_path.touch() + + def log( + self, + event_type: str = "", + severity: str = "info", + message: str = "", + session_id: Optional[str] = None, + user_id: Optional[str] = None, + details: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + **kwargs: Any + ) -> None: + """ + Log an audit event to JSON Lines file. + + Args: + event_type: Type of event + severity: Severity level + message: Human-readable message + session_id: Optional session ID + user_id: Optional user ID + details: Optional details dictionary + error: Optional error message + **kwargs: Additional fields + """ + import json + from datetime import datetime, timezone + + entry = { + "timestamp": datetime.now(timezone.utc).isoformat(), + "event_type": event_type, + "severity": severity, + "message": message, + } + + if session_id is not None: + entry["session_id"] = session_id + if user_id is not None: + entry["user_id"] = user_id + if details is not None: + entry["details"] = details + if error is not None: + entry["error"] = error + + # Add any additional kwargs + entry.update(kwargs) + + # Write to file (append mode, file locking via 'a' mode) + with open(self.log_path, 'a') as f: + f.write(json.dumps(entry) + '\n') diff --git a/src/pod_executor/security/resource_limits.py b/src/pod_executor/security/resource_limits.py new file mode 100644 index 0000000..7540dfe --- /dev/null +++ b/src/pod_executor/security/resource_limits.py @@ -0,0 +1,150 @@ +""" +Resource limit parser and validator. + +Parses and validates resource limit strings (memory, CPU, storage). +All values must be positive and within reasonable limits. +""" + +import re +from typing import Dict, Any + + +def parse_memory_string(memory: str) -> int: + """ + Parse memory string to bytes. + + Supports: k, m, g suffixes (case-insensitive) + Examples: "512m" → 536870912, "2g" → 2147483648 + + Args: + memory: Memory string with suffix (e.g., "512m", "2g", "1024k") + + Returns: + Memory in bytes + + Raises: + ValueError: If format is invalid or value is <= 0 + """ + memory = memory.strip() + + # Pattern: optional sign, number (int or float), suffix (k/m/g) + pattern = r'^(-?\d+(?:\.\d+)?)\s*([kmgKMG])$' + match = re.match(pattern, memory) + + if not match: + raise ValueError( + f"Invalid memory format: '{memory}'. " + f"Expected format: (e.g., '512m', '2g')" + ) + + value_str, suffix = match.groups() + value = float(value_str) + + if value <= 0: + raise ValueError( + f"Memory value must be positive, got: {value}" + ) + + # Convert to bytes + suffix_lower = suffix.lower() + multipliers = { + 'k': 1024, + 'm': 1024 * 1024, + 'g': 1024 * 1024 * 1024, + } + + bytes_value = int(value * multipliers[suffix_lower]) + + return bytes_value + + +def parse_cpu_quota(cpu_quota: int) -> int: + """ + Validate CPU quota value. + + CPU quota is in microseconds per 100ms period. + 100000 = 100% of one CPU core + + Args: + cpu_quota: CPU quota in microseconds (e.g., 50000 for 50% of one core) + + Returns: + Validated CPU quota value + + Raises: + ValueError: If quota <= 0 or > 1000000 (10 cores max) + """ + if cpu_quota <= 0: + raise ValueError( + f"CPU quota must be positive, got: {cpu_quota}" + ) + + # Maximum of 10 cores (1000000 microseconds) + if cpu_quota > 1000000: + raise ValueError( + f"CPU quota exceeds maximum of 1000000 (10 cores), got: {cpu_quota}" + ) + + return cpu_quota + + +def parse_storage_string(storage: str) -> int: + """ + Parse storage string to bytes (same as memory). + + Args: + storage: Storage string with suffix (e.g., "1g", "512m") + + Returns: + Storage in bytes + + Raises: + ValueError: If format is invalid or value is <= 0 + """ + return parse_memory_string(storage) + + +class ResourceLimits: + """ + Resource limits with validation. + + Encapsulates memory, storage, CPU, and timeout limits with validation. + Provides conversion to Podman container parameters. + """ + + def __init__( + self, + memory: str, + storage: str, + cpu_quota: int, + timeout: int = 300 + ): + """ + Initialize resource limits with validation. + + Args: + memory: Memory limit string (e.g., "512m", "2g") + storage: Storage limit string (e.g., "1g", "10g") + cpu_quota: CPU quota in microseconds per 100ms period + timeout: Execution timeout in seconds (default: 300) + + Raises: + ValueError: If any limit is invalid + """ + self.memory_bytes = parse_memory_string(memory) + self.storage_bytes = parse_storage_string(storage) + self.cpu_quota = parse_cpu_quota(cpu_quota) + self.timeout = timeout + + def to_podman_params(self) -> Dict[str, Any]: + """ + Convert to Podman container create parameters. + + Returns: + Dictionary of parameters suitable for Podman container creation + """ + return { + "mem_limit": str(self.memory_bytes), # Podman expects string + "cpu_quota": self.cpu_quota + # Note: storage_bytes tracked internally but not passed to Podman (not supported) + } diff --git a/src/pod_executor/security/validation.py b/src/pod_executor/security/validation.py new file mode 100644 index 0000000..74efed7 --- /dev/null +++ b/src/pod_executor/security/validation.py @@ -0,0 +1,232 @@ +""" +Security validation protocol for pod_executor. + +Provides protocols (interfaces) for security validation that can be implemented +by consuming applications. Default implementations are provided. +""" + +from typing import Protocol, Set, Optional, Dict, Any + + +class SecurityError(Exception): + """Raised when security policy is violated.""" + pass + + +class OperationValidatorProtocol(Protocol): + """Protocol for validating container operations.""" + + session_containers: Set[str] + + def validate_container_create( + self, + image: str, + params: Dict[str, Any], + session_id: Optional[str] = None + ) -> None: + """ + Validate container creation parameters against security policy. + + Args: + image: Container image name + params: Container creation parameters + session_id: Optional session ID for volume validation + + Raises: + SecurityError: If any security policy is violated + """ + ... + + def validate_container_start(self, container_id: str) -> None: + """ + Validate container start. + + Args: + container_id: Container ID to start + + Raises: + SecurityError: If operation is not allowed + """ + ... + + def validate_container_stop(self, container_id: str) -> None: + """ + Validate container stop. + + Args: + container_id: Container ID to stop + + Raises: + SecurityError: If operation is not allowed + """ + ... + + def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]: + """ + Validate generic operation. + + Args: + operation: Operation type + target: Operation target + + Returns: + (allowed, reason) tuple + """ + ... + + def register_session_container(self, container_id: str) -> None: + """Register a container as belonging to a session.""" + ... + + def unregister_session_container(self, container_id: str) -> None: + """Unregister a session container.""" + ... + + +class NoOpValidator: + """ + No-op validator that allows all operations. + + WARNING: This validator provides NO SECURITY. Only use for testing + or in fully trusted environments. + """ + + def __init__(self): + """Initialize with empty session container set.""" + self.session_containers: Set[str] = set() + + def validate_container_create( + self, + image: str, + params: Dict[str, Any], + session_id: Optional[str] = None + ) -> None: + """Allow all container creations.""" + pass + + def validate_container_start(self, container_id: str) -> None: + """Allow all container starts.""" + pass + + def validate_container_stop(self, container_id: str) -> None: + """Allow all container stops.""" + pass + + def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]: + """Allow all operations.""" + return (True, None) + + def register_session_container(self, container_id: str) -> None: + """Track session container.""" + self.session_containers.add(container_id) + + def unregister_session_container(self, container_id: str) -> None: + """Untrack session container.""" + self.session_containers.discard(container_id) + + +class BasicValidator: + """ + Basic validator with minimal security checks. + + Enforces: + - Allowed image patterns + - Required security parameters + - Forbidden dangerous parameters + - Session container tracking + """ + + # Allowed container images with wildcard support + DEFAULT_ALLOWED_IMAGES = [ + "python:3.11*", + "python:3.12*", + "jupyter/*", + "mcp-forge/*", + ] + + # Parameters that are forbidden + FORBIDDEN_PARAMS = [ + "privileged", + "cap_add", + "devices", + "pid_mode", + ] + + def __init__(self, allowed_images: Optional[list[str]] = None): + """ + Initialize basic validator. + + Args: + allowed_images: List of allowed image patterns (supports wildcards) + """ + self.allowed_images = allowed_images or self.DEFAULT_ALLOWED_IMAGES + self.session_containers: Set[str] = set() + + def validate_container_create( + self, + image: str, + params: Dict[str, Any], + session_id: Optional[str] = None + ) -> None: + """ + Validate container creation. + + Args: + image: Container image name + params: Container creation parameters + session_id: Optional session ID + + Raises: + SecurityError: If validation fails + """ + # Validate image is allowed + if not self._is_image_allowed(image): + raise SecurityError( + f"Image '{image}' not in allowlist. " + f"Allowed patterns: {self.allowed_images}" + ) + + # Check for forbidden parameters + for forbidden in self.FORBIDDEN_PARAMS: + if forbidden in params: + raise SecurityError( + f"Forbidden parameter '{forbidden}' in container creation" + ) + + # Ensure security_opt includes no-new-privileges + security_opts = params.get("security_opt", []) + if "no-new-privileges" not in security_opts: + raise SecurityError( + "Container must include security_opt=['no-new-privileges']" + ) + + def validate_container_start(self, container_id: str) -> None: + """Validate container start.""" + if container_id not in self.session_containers: + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + def validate_container_stop(self, container_id: str) -> None: + """Validate container stop.""" + if container_id not in self.session_containers: + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]: + """Validate generic operation.""" + return (True, None) # Allow by default + + def register_session_container(self, container_id: str) -> None: + """Register session container.""" + self.session_containers.add(container_id) + + def unregister_session_container(self, container_id: str) -> None: + """Unregister session container.""" + self.session_containers.discard(container_id) + + def _is_image_allowed(self, image: str) -> bool: + """Check if image matches any allowed pattern.""" + import fnmatch + return any(fnmatch.fnmatch(image, pattern) for pattern in self.allowed_images) diff --git a/src/pod_executor/simple/__init__.py b/src/pod_executor/simple/__init__.py new file mode 100644 index 0000000..b4ad239 --- /dev/null +++ b/src/pod_executor/simple/__init__.py @@ -0,0 +1,8 @@ +"""Simple stateless code executor.""" + +from pod_executor.simple.executor import CodeExecutor, ExecutionResult + +__all__ = [ + "CodeExecutor", + "ExecutionResult", +] diff --git a/src/pod_executor/simple/executor.py b/src/pod_executor/simple/executor.py new file mode 100644 index 0000000..732142f --- /dev/null +++ b/src/pod_executor/simple/executor.py @@ -0,0 +1,222 @@ +"""Code execution in isolated containers.""" + +from typing import Any, Optional +from dataclasses import dataclass, asdict +import json +import time +import textwrap +import uuid + +from pod_executor.containers.manager import SecureContainerManager, ContainerConfig +from pod_executor.security.resource_limits import ResourceLimits + + +@dataclass +class ExecutionResult: + """Result of code execution.""" + success: bool + stdout: str + stderr: str + result: Optional[Any] + execution_time: float + exit_code: int + error: Optional[str] = None + + def to_dict(self) -> dict: + """Convert to dictionary for JSON serialization.""" + return asdict(self) + + +class CodeExecutor: + """Executes Python code in isolated containers.""" + + def __init__( + self, + container_manager: SecureContainerManager, + image: str, + resource_limits: ResourceLimits + ): + self.container_manager = container_manager + self.image = image + self.resource_limits = resource_limits + + def execute( + self, + code: str, + timeout: Optional[int] = None, + injection_code: Optional[str] = None, + bridge_socket_path: Optional[str] = None + ) -> ExecutionResult: + """ + Execute Python code in a fresh container. + + Process: + 1. Create container with code + 2. Start container + 3. Wait for completion (with timeout) + 4. Capture stdout/stderr + 5. Extract result from last expression + 6. Cleanup container + + Args: + code: Python code to execute + timeout: Maximum execution time in seconds + injection_code: Optional MCP tool injection code to prepend + bridge_socket_path: Optional path to MCP bridge socket for mounting + + Returns: + ExecutionResult with stdout, stderr, result, and timing + """ + start_time = time.time() + container_id = None + + # Use provided timeout or default from resource limits + exec_timeout = timeout if timeout is not None else self.resource_limits.timeout + + try: + # Prepare code wrapper (with injection if provided) + wrapped_code = self._prepare_code(code, injection_code=injection_code) + + # Generate a session ID for this execution to register the container + import uuid + session_id = f"simple-exec-{uuid.uuid4().hex[:12]}" + + # Set up volumes for MCP bridge socket if provided + volumes = {} + if bridge_socket_path: + volumes[bridge_socket_path] = { + "bind": bridge_socket_path, + "mode": "rw" + } + + # Create container configuration + config = ContainerConfig( + image=self.image, + command=["python", "-c", wrapped_code], + resource_limits=self.resource_limits, + volumes=volumes if volumes else None + ) + + # Create and start container with session_id for proper registration + container_id = self.container_manager.create_container(config, session_id=session_id) + self.container_manager.start_container(container_id) + + # Wait for completion + exit_code = self.container_manager.wait_for_container( + container_id, + timeout=exec_timeout + ) + + # Get logs + stdout, stderr = self.container_manager.get_container_logs(container_id) + + # Parse output to extract result + result, error = self._parse_output(stdout) + + execution_time = time.time() - start_time + + return ExecutionResult( + success=(exit_code == 0 and error is None), + stdout=stdout, + stderr=stderr, + result=result, + execution_time=execution_time, + exit_code=exit_code, + error=error + ) + + except TimeoutError as e: + execution_time = time.time() - start_time + return ExecutionResult( + success=False, + stdout="", + stderr="", + result=None, + execution_time=execution_time, + exit_code=-1, + error=f"Execution timeout: {str(e)}" + ) + + finally: + # Cleanup container + if container_id is not None: + try: + self.container_manager.remove_container(container_id) + except Exception: + pass # Best effort cleanup + + def _prepare_code(self, code: str, injection_code: Optional[str] = None) -> str: + """ + Wrap code to capture result and handle errors. + + Wraps code in try/except and captures: + - Last expression result + - Exceptions with traceback + - Execution metadata + + Args: + code: User code to execute + injection_code: Optional MCP tool injection code to prepend + + Returns wrapped code that outputs JSON to stdout. + """ + # Prepend injection code if provided + if injection_code: + full_code = injection_code + "\n\n" + code + else: + full_code = code + + # Escape the code for embedding in exec string + escaped_code = full_code.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n') + + wrapper_template = ''' +import sys +import json +import traceback + +def __mcp_execute(): + result = None + error = None + + try: + # Execute user code + exec_globals = {} + exec("""%s""", exec_globals) + + # Try to get result from last expression + result = exec_globals.get('_', None) + + except SyntaxError as e: + error = f"SyntaxError: {e.msg} (line {e.lineno})" + except Exception as e: + error = f"{type(e).__name__}: {str(e)}" + + # Output result as JSON + print(json.dumps({"result": result, "error": error})) + +__mcp_execute() +''' + + return wrapper_template % escaped_code + + def _parse_output(self, stdout: str) -> tuple[Optional[Any], Optional[str]]: + """ + Parse execution output to extract result and error. + + Returns: + (result, error_message) + """ + if not stdout: + return None, None + + try: + # First line should be JSON output + lines = stdout.split('\n', 1) + json_line = lines[0] + + data = json.loads(json_line) + return data.get("result"), data.get("error") + + except (json.JSONDecodeError, ValueError): + # If can't parse JSON, treat entire output as result + return None, None