ms-enclave 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

ms_enclave/sandbox/model/__init__.py

@@ -0,0 +1,21 @@
+"""Data models for sandbox system."""
+
+from .base import ExecutionStatus, SandboxStatus, SandboxType, ToolType
+from .config import (
+    DockerNotebookConfig,
+    DockerSandboxConfig,
+    FileOperationConfig,
+    PythonExecutorConfig,
+    SandboxConfig,
+    ShellExecutorConfig,
+    ToolConfig,
+)
+from .requests import (
+    ExecuteCodeRequest,
+    ExecuteCommandRequest,
+    FileOperationRequest,
+    ReadFileRequest,
+    ToolExecutionRequest,
+    WriteFileRequest,
+)
+from .responses import CommandResult, HealthCheckResult, SandboxInfo, ToolResult

ms_enclave/sandbox/model/base.py

@@ -0,0 +1,36 @@
+"""Base data models."""
+
+from enum import Enum
+
+
+class SandboxStatus(str, Enum):
+    """Sandbox status enumeration."""
+
+    INITIALIZING = 'initializing'
+    RUNNING = 'running'
+    STOPPING = 'stopping'
+    STOPPED = 'stopped'
+    ERROR = 'error'
+
+
+class SandboxType(str, Enum):
+    """Sandbox type enumeration."""
+    DOCKER = 'docker'
+    DOCKER_NOTEBOOK = 'docker_notebook'
+    DUMMY = 'dummy'
+
+
+class ToolType(str, Enum):
+    """Tool type enumeration."""
+    SANDBOX = 'sandbox'
+    FUNCTION = 'function'
+    EXTERNAL = 'external'
+
+
+class ExecutionStatus(str, Enum):
+    """Execution status enumeration."""
+
+    SUCCESS = 'success'
+    ERROR = 'error'
+    TIMEOUT = 'timeout'
+    CANCELLED = 'cancelled'

ms_enclave/sandbox/model/config.py

@@ -0,0 +1,97 @@
+"""Configuration data models."""
+
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class SandboxConfig(BaseModel):
+    """Base sandbox configuration."""
+
+    timeout: int = Field(default=30, description='Default timeout in seconds')
+    tools_config: Dict[str, Dict[
+        str, Any]] = Field(default_factory=dict, description='Configuration for tools within the sandbox')
+    working_dir: str = Field(default='/sandbox', description='Default working directory')
+    env_vars: Dict[str, str] = Field(default_factory=dict, description='Environment variables')
+    resource_limits: Dict[str, Any] = Field(default_factory=dict, description='Resource limits')
+
+
+class DockerSandboxConfig(SandboxConfig):
+    """Docker-specific sandbox configuration."""
+
+    image: str = Field('python:3.11-slim', description='Docker image name')
+    command: Optional[Union[str, List[str]]] = Field(None, description='Container command')
+    volumes: Dict[str, Dict[str, str]] = Field(
+        default_factory=dict,
+        description="Volume mounts. Format: { host_path: {'bind': container_path, 'mode': 'rw|ro'} }"
+    )
+    ports: Dict[str, str] = Field(default_factory=dict, description='Port mappings')
+    network: Optional[str] = Field('bridge', description='Network name')
+    memory_limit: str = Field(default='1g', description='Memory limit')
+    cpu_limit: float = Field(default=1.0, description='CPU limit')
+    network_enabled: bool = Field(default=True, description='Enable network access')
+    privileged: bool = Field(default=False, description='Run in privileged mode')
+    remove_on_exit: bool = Field(default=True, description='Remove container on exit')
+
+    @field_validator('memory_limit')
+    def validate_memory_limit(cls, v):
+        """Validate memory limit format."""
+        if not isinstance(v, str):
+            raise ValueError('Memory limit must be a string')
+        # Basic validation for memory format (e.g., '512m', '1g', '2G')
+        import re
+        if not re.match(r'^\d+[kmgKMG]?$', v):
+            raise ValueError('Invalid memory limit format')
+        return v
+
+    @field_validator('cpu_limit')
+    def validate_cpu_limit(cls, v):
+        """Validate CPU limit."""
+        if v <= 0:
+            raise ValueError('CPU limit must be positive')
+        return v
+
+
+class DockerNotebookConfig(DockerSandboxConfig):
+    """Docker Notebook-specific sandbox configuration."""
+
+    image: str = Field('jupyter-kernel-gateway', description='Docker image name for Jupyter Notebook')
+    host: str = Field('127.0.0.1', description='Host for Jupyter Notebook')
+    port: int = Field(8888, description='Port for Jupyter Notebook')
+    token: Optional[str] = Field(None, description='Token for Jupyter Notebook access')
+
+
+class ToolConfig(BaseModel):
+    """Tool configuration."""
+
+    enabled: bool = Field(default=True, description='Whether tool is enabled')
+    timeout: int = Field(default=30, description='Tool execution timeout')
+    parameters: Dict[str, Any] = Field(default_factory=dict, description='Tool parameters')
+    restrictions: Dict[str, Any] = Field(default_factory=dict, description='Tool restrictions')
+
+
+class PythonExecutorConfig(ToolConfig):
+    """Python executor configuration."""
+
+    python_path: str = Field(default='python3', description='Python executable path')
+    allowed_modules: Optional[List[str]] = Field(None, description='Allowed modules (None = all)')
+    blocked_modules: List[str] = Field(default_factory=list, description='Blocked modules')
+    max_output_size: int = Field(default=1024 * 1024, description='Maximum output size in bytes')
+
+
+class ShellExecutorConfig(ToolConfig):
+    """Shell executor configuration."""
+
+    shell_path: str = Field(default='/bin/bash', description='Shell executable path')
+    allowed_commands: Optional[List[str]] = Field(None, description='Allowed commands (None = all)')
+    blocked_commands: List[str] = Field(default_factory=list, description='Blocked commands')
+    max_output_size: int = Field(default=1024 * 1024, description='Maximum output size in bytes')
+
+
+class FileOperationConfig(ToolConfig):
+    """File operation configuration."""
+
+    allowed_paths: Optional[List[str]] = Field(None, description='Allowed paths (None = all)')
+    blocked_paths: List[str] = Field(default_factory=list, description='Blocked paths')
+    max_file_size: int = Field(default=10 * 1024 * 1024, description='Maximum file size in bytes')
+    allowed_extensions: Optional[List[str]] = Field(None, description='Allowed file extensions')

ms_enclave/sandbox/model/requests.py

@@ -0,0 +1,57 @@
+"""Request data models."""
+
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class ExecuteCodeRequest(BaseModel):
+    """Request model for code execution."""
+
+    code: str = Field(..., description='Code to execute')
+    language: str = Field(default='python', description='Programming language')
+    timeout: Optional[int] = Field(None, description='Execution timeout in seconds')
+    working_dir: Optional[str] = Field(None, description='Working directory')
+    env: Dict[str, str] = Field(default_factory=dict, description='Environment variables')
+    capture_output: bool = Field(True, description='Whether to capture stdout/stderr')
+
+
+class ExecuteCommandRequest(BaseModel):
+    """Request model for shell command execution."""
+
+    command: Union[str, List[str]] = Field(..., description='Command to execute')
+    timeout: Optional[int] = Field(None, description='Execution timeout in seconds')
+    working_dir: Optional[str] = Field(None, description='Working directory')
+    env: Dict[str, str] = Field(default_factory=dict, description='Environment variables')
+    shell: bool = Field(True, description='Execute in shell')
+
+
+class FileOperationRequest(BaseModel):
+    """Base request model for file operations."""
+
+    path: str = Field(..., description='File path')
+
+
+class ReadFileRequest(FileOperationRequest):
+    """Request model for reading files."""
+
+    encoding: Optional[str] = Field('utf-8', description='File encoding')
+    binary: bool = Field(False, description='Read as binary')
+
+
+class WriteFileRequest(FileOperationRequest):
+    """Request model for writing files."""
+
+    content: Union[str, bytes] = Field(..., description='File content')
+    encoding: Optional[str] = Field('utf-8', description='File encoding')
+    binary: bool = Field(False, description='Write as binary')
+    create_dirs: bool = Field(True, description='Create parent directories if needed')
+
+
+class ToolExecutionRequest(BaseModel):
+    """Request model for tool execution."""
+
+    sandbox_id: str = Field(..., description='Sandbox identifier')
+    tool_name: str = Field(..., description='Name of tool to execute')
+    parameters: Dict[str, Any] = Field(default_factory=dict, description='Tool parameters')
+    timeout: Optional[int] = Field(None, description='Execution timeout in seconds')

ms_enclave/sandbox/model/responses.py

@@ -0,0 +1,57 @@
+"""Response data models."""
+
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from .base import ExecutionStatus, SandboxStatus
+
+
+class SandboxInfo(BaseModel):
+    """Information about a sandbox instance."""
+
+    id: str = Field(..., description='Sandbox identifier')
+    status: SandboxStatus = Field(..., description='Current status')
+    type: str = Field(..., description="Sandbox type (e.g., 'docker')")
+    config: Dict[str, Any] = Field(default_factory=dict, description='Sandbox configuration')
+    created_at: datetime = Field(default_factory=datetime.now, description='Creation timestamp')
+    updated_at: datetime = Field(default_factory=datetime.now, description='Last update timestamp')
+    metadata: Dict[str, Any] = Field(default_factory=dict, description='Additional metadata')
+    available_tools: Dict[str, Any] = Field(default_factory=dict, description='Available tools')
+
+
+class ExecutionResult(BaseModel):
+    """Base class for execution results."""
+
+    status: ExecutionStatus = Field(..., description='Execution status')
+    execution_time: Optional[float] = Field(None, description='Execution time in seconds')
+    timestamp: datetime = Field(default_factory=datetime.now, description='Execution timestamp')
+
+
+class ToolResult(ExecutionResult):
+    """Result of tool execution."""
+
+    tool_name: str = Field(..., description='Name of tool executed')
+    output: Any = Field(None, description='Tool execution result')
+    metadata: Dict[str, Any] = Field(default_factory=dict, description='Additional metadata')
+    error: Optional[str] = Field(None, description='Error message if failed')
+
+
+class CommandResult(ExecutionResult):
+    """Command execution result."""
+
+    command: Union[str, List[str]] = Field(..., description='Executed command')
+    exit_code: int = Field(..., description='Exit code of the command')
+    stdout: Optional[str] = Field(None, description='Standard output')
+    stderr: Optional[str] = Field(None, description='Standard error output')
+
+
+class HealthCheckResult(BaseModel):
+    """Health check result."""
+
+    healthy: bool = Field(..., description='Health status')
+    version: str = Field(..., description='System version')
+    uptime: float = Field(..., description='System uptime in seconds')
+    active_sandboxes: int = Field(..., description='Number of active sandboxes')
+    system_info: Dict[str, Any] = Field(default_factory=dict, description='System information')

ms_enclave/sandbox/server/server.py

@@ -0,0 +1,195 @@
+"""FastAPI server for sandbox system with optional API key authentication."""
+
+import time
+from contextlib import asynccontextmanager
+from typing import Any, Dict, List, Optional
+
+from fastapi import BackgroundTasks, Depends, FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+from ..manager import LocalSandboxManager
+from ..model import (
+    DockerSandboxConfig,
+    HealthCheckResult,
+    SandboxConfig,
+    SandboxInfo,
+    SandboxStatus,
+    SandboxType,
+    ToolExecutionRequest,
+    ToolResult,
+)
+
+
+class SandboxServer:
+    """FastAPI-based sandbox server.
+    """
+
+    def __init__(self, cleanup_interval: int = 300, api_key: Optional[str] = None):
+        """Initialize sandbox server.
+
+        Args:
+            cleanup_interval: Cleanup interval in seconds
+            api_key: Optional API key to protect endpoints. If None, auth is disabled.
+        """
+        self.manager = LocalSandboxManager(cleanup_interval)
+        self.api_key: Optional[str] = api_key
+        self.app = FastAPI(
+            title='Sandbox API',
+            description='Agent sandbox execution environment',
+            version='1.0.0',
+            lifespan=self.lifespan
+        )
+        self._setup_middleware()
+        self._setup_auth_middleware()
+        self._setup_routes()
+        self.start_time = time.time()
+
+    @asynccontextmanager
+    async def lifespan(self, app: FastAPI):
+        """Application lifespan management."""
+        # Startup
+        await self.manager.start()
+        yield
+        # Shutdown
+        await self.manager.stop()
+
+    def _setup_middleware(self):
+        """Setup middleware."""
+        self.app.add_middleware(
+            CORSMiddleware,
+            allow_origins=['*'],
+            allow_credentials=True,
+            allow_methods=['*'],
+            allow_headers=['*'],
+        )
+
+    def _setup_auth_middleware(self) -> None:
+        """Setup optional API key authentication middleware.
+
+        When ``self.api_key`` is None, this middleware is a no-op.
+        Otherwise, it enforces that every request includes the correct API key
+        either via ``X-API-Key`` header or ``api_key`` query parameter.
+        """
+
+        @self.app.middleware('http')
+        async def auth_middleware(request: Request, call_next):  # type: ignore[unused-ignore]
+            # Fast path when auth is disabled
+            if not self.api_key:
+                return await call_next(request)
+
+            provided = request.headers.get('x-api-key') or request.query_params.get('api_key')
+            if provided == self.api_key:
+                return await call_next(request)
+
+            return JSONResponse(status_code=401, content={'detail': 'Unauthorized'})
+
+    def _setup_routes(self):
+        """Setup API routes."""
+
+        # Health check
+        @self.app.get('/health', response_model=HealthCheckResult)
+        async def health_check():
+            """Health check endpoint."""
+            stats = await self.manager.get_stats()
+            return HealthCheckResult(
+                healthy=True,
+                version='1.0.0',
+                uptime=time.time() - self.start_time,
+                active_sandboxes=stats['total_sandboxes'],
+                system_info=stats
+            )
+
+        # Sandbox management
+        @self.app.post('/sandbox/create')
+        async def create_sandbox(sandbox_type: SandboxType, config: Optional[Dict] = None):
+            """Create a new sandbox."""
+            try:
+                sandbox_id = await self.manager.create_sandbox(sandbox_type, config)
+
+                return {'sandbox_id': sandbox_id}
+
+            except Exception as e:
+                raise HTTPException(status_code=400, detail=str(e))
+
+        @self.app.get('/sandboxes')
+        async def list_sandboxes(status: Optional[SandboxStatus] = None):
+            """List all sandboxes."""
+            sandboxes = await self.manager.list_sandboxes(status)
+            return sandboxes
+
+        @self.app.get('/sandbox/{sandbox_id}', response_model=SandboxInfo)
+        async def get_sandbox_info(sandbox_id: str):
+            """Get sandbox information."""
+            info = await self.manager.get_sandbox_info(sandbox_id)
+            if not info:
+                raise HTTPException(status_code=404, detail='Sandbox not found')
+            return info
+
+        @self.app.post('/sandbox/{sandbox_id}/stop')
+        async def stop_sandbox(sandbox_id: str):
+            """Stop a sandbox."""
+            success = await self.manager.stop_sandbox(sandbox_id)
+            if not success:
+                raise HTTPException(status_code=404, detail='Sandbox not found')
+            return {'message': 'Sandbox stopped successfully'}
+
+        @self.app.delete('/sandbox/{sandbox_id}')
+        async def delete_sandbox(sandbox_id: str):
+            """Delete a sandbox."""
+            success = await self.manager.delete_sandbox(sandbox_id)
+            if not success:
+                raise HTTPException(status_code=404, detail='Sandbox not found')
+            return {'message': 'Sandbox deleted successfully'}
+
+        # Tool execution
+        @self.app.post('/sandbox/tool/execute', response_model=ToolResult)
+        async def execute_tool(request: ToolExecutionRequest):
+            """Execute tool in sandbox."""
+            try:
+                result = await self.manager.execute_tool(request.sandbox_id, request.tool_name, request.parameters)
+                return result
+            except ValueError as e:
+                raise HTTPException(status_code=404, detail=str(e))
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+        @self.app.get('/sandbox/{sandbox_id}/tools')
+        async def get_sandbox_tools(sandbox_id: str):
+            """Get available tools for a sandbox."""
+            try:
+                tools = await self.manager.get_sandbox_tools(sandbox_id)
+                return tools
+            except ValueError as e:
+                raise HTTPException(status_code=404, detail=str(e))
+
+        # System info
+        @self.app.get('/stats')
+        async def get_stats():
+            """Get system statistics."""
+            return await self.manager.get_stats()
+
+    def run(self, host: str = '0.0.0.0', port: int = 8000, **kwargs):
+        """Run the server.
+
+        Args:
+            host: Host to bind to
+            port: Port to bind to
+            **kwargs: Additional uvicorn arguments
+        """
+        import uvicorn
+        uvicorn.run(self.app, host=host, port=port, **kwargs)
+
+
+# Create a default server instance
+def create_server(cleanup_interval: int = 300, api_key: Optional[str] = None) -> SandboxServer:
+    """Create a sandbox server instance.
+
+    Args:
+        cleanup_interval: Cleanup interval in seconds
+        api_key: Optional API key to protect endpoints. If None, auth is disabled.
+
+    Returns:
+        Sandbox server instance
+    """
+    return SandboxServer(cleanup_interval, api_key)

ms_enclave/sandbox/tools/__init__.py

@@ -0,0 +1,4 @@
+"""Tool interfaces and implementations."""
+
+from .base import Tool, ToolFactory, register_tool
+from .sandbox_tools import FileOperation, PythonExecutor, ShellExecutor

ms_enclave/sandbox/tools/base.py

@@ -0,0 +1,95 @@
+"""Base tool interface and factory."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type
+
+from ..model import SandboxType, ToolResult
+from .tool_info import ToolParams
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import Sandbox
+
+
+class Tool(ABC):
+    """Abstract base class for all tools."""
+
+    def __init__(
+        self,
+        *,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        parameters: Optional[ToolParams] = None,
+        enabled: bool = True,
+        timeout: Optional[int] = None,
+        **kwargs,
+    ):
+        self._name = name or self.__class__.__name__
+        self._description = description
+        self._parameters = parameters
+        self.enabled = enabled
+        self.timeout = timeout
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def description(self) -> Optional[str]:
+        return self._description
+
+    @property
+    def parameters(self) -> Optional[ToolParams]:
+        return self._parameters
+
+    @property
+    def schema(self) -> Dict:
+        return {
+            'type': 'function',
+            'function': {
+                'name': self.name,
+                'description': self._description,
+                'parameters': self._parameters.model_dump(exclude_none=True) if self._parameters else {},
+            },
+        }
+
+    @property
+    @abstractmethod
+    def required_sandbox_type(self) -> Optional[SandboxType]:
+        """Get the required sandbox type for this tool."""
+        pass
+
+    @abstractmethod
+    async def execute(self, sandbox_context: 'Sandbox', **kwargs) -> ToolResult:
+        """Execute the tool with given sandbox context and parameters."""
+        pass
+
+
+class ToolFactory:
+    """Factory for creating tool instances."""
+
+    _tools: Dict[str, Type[Tool]] = {}
+
+    @classmethod
+    def register_tool(cls, tool_name: str, tool_class: Type[Tool]):
+        cls._tools[tool_name] = tool_class
+
+    @classmethod
+    def create_tool(cls, tool_name: str, **kwargs) -> Tool:
+        if tool_name not in cls._tools:
+            raise ValueError(f'Tool name {tool_name} is not registered')
+
+        tool_class = cls._tools[tool_name]
+        return tool_class(**kwargs)
+
+    @classmethod
+    def get_available_tools(cls) -> list[str]:
+        return list(cls._tools.keys())
+
+
+def register_tool(tool_name: str):
+
+    def decorator(tool_class: Type[Tool]):
+        ToolFactory.register_tool(tool_name, tool_class)
+        return tool_class
+
+    return decorator

ms_enclave/sandbox/tools/sandbox_tool.py

@@ -0,0 +1,46 @@
+from typing import Any, Dict, Optional
+
+from ..model import SandboxType
+from .base import Tool
+from .tool_info import ToolParams
+
+
+class SandboxTool(Tool):
+    """Built-in tool class"""
+
+    _name: Optional[str] = None
+    _description: Optional[str] = None
+    _parameters: Optional[ToolParams] = None
+    _sandbox_type: Optional[SandboxType] = None
+
+    def __init__(
+        self,
+        *,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        parameters: Optional[ToolParams] = None,
+        sandbox: Optional[Any] = None,
+        sandbox_type: Optional[SandboxType] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the tool.
+
+        Note: Once the sandbox is set, it does not change.
+        """
+        super().__init__(
+            name=name,
+            description=description,
+            parameters=parameters,
+            **kwargs,
+        )
+        self._sandbox = sandbox
+        self._name = name or self.__class__._name
+        self._description = description or self.__class__._description
+        self._parameters = parameters or self.__class__._parameters
+        self._sandbox_type = sandbox_type or self.__class__._sandbox_type
+
+    @property
+    def required_sandbox_type(self) -> Optional[SandboxType]:
+        """Get the required sandbox type for this tool."""
+        return self._sandbox_type

ms_enclave/sandbox/tools/sandbox_tools/__init__.py

@@ -0,0 +1,4 @@
+from .file_operation import FileOperation
+from .notebook_executor import NotebookExecutor
+from .python_executor import PythonExecutor
+from .shell_executor import ShellExecutor