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

ms_enclave/sandbox/manager/http_manager.py

@@ -0,0 +1,405 @@
+"""HTTP-based sandbox manager for remote sandbox services."""
+
+from typing import Any, Dict, List, Optional, Union
+
+import aiohttp
+
+from ms_enclave.utils import get_logger
+
+from ..model import SandboxConfig, SandboxInfo, SandboxStatus, SandboxType, ToolExecutionRequest, ToolResult
+from .base import SandboxManager
+
+logger = get_logger()
+
+
+class HttpSandboxManager(SandboxManager):
+    """HTTP-based sandbox manager for remote services.
+    """
+
+    def __init__(self, base_url: str, timeout: int = 30, api_key: Optional[str] = None):
+        """Initialize HTTP sandbox manager.
+
+        Args:
+            base_url: Base URL of the sandbox service
+            timeout: Request timeout in seconds
+            api_key: Optional API key to include as ``X-API-Key`` header for all requests
+        """
+        super().__init__()
+        self.base_url = base_url.rstrip('/')
+        self.timeout = aiohttp.ClientTimeout(total=timeout)
+        self._default_headers: Optional[Dict[str, str]] = {'X-API-Key': api_key} if api_key else None
+        self._session: Optional[aiohttp.ClientSession] = None
+
+    async def start(self) -> None:
+        """Start the HTTP sandbox manager."""
+        if self._running:
+            return
+
+        self._connector = aiohttp.TCPConnector()
+        self._session = aiohttp.ClientSession(
+            connector=self._connector, timeout=self.timeout, headers=self._default_headers
+        )
+        self._running = True
+        logger.info(f'HTTP sandbox manager started, connected to {self.base_url}')
+
+    async def stop(self) -> None:
+        """Stop the HTTP sandbox manager."""
+        if not self._running:
+            return
+
+        self._running = False
+
+        # Clean up all sandboxes created by this manager
+        if self._sandboxes:
+            logger.info(f'Cleaning up {len(self._sandboxes)} sandboxes created by this manager')
+            await self.cleanup_all_sandboxes()
+
+        if self._session:
+            await self._session.close()
+            self._session = None
+        logger.info('HTTP sandbox manager stopped')
+
+    async def create_sandbox(
+        self,
+        sandbox_type: SandboxType,
+        config: Optional[Union[SandboxConfig, Dict]] = None,
+        sandbox_id: Optional[str] = None
+    ) -> str:
+        """Create a new sandbox via HTTP API.
+
+        Args:
+            sandbox_type: Type of sandbox to create
+            config: Sandbox configuration
+            sandbox_id: Optional sandbox ID
+
+        Returns:
+            Sandbox ID
+
+        Raises:
+            ValueError: If sandbox type is not supported
+            RuntimeError: If sandbox creation fails
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        # Match server's endpoint format: POST /sandbox/create
+        params = {'sandbox_type': sandbox_type.value if isinstance(sandbox_type, SandboxType) else sandbox_type}
+        if isinstance(config, SandboxConfig):
+            payload = config.model_dump(exclude_none=True)
+        elif isinstance(config, dict):
+            payload = config
+        else:
+            payload = {}
+
+        try:
+            async with self._session.post(f'{self.base_url}/sandbox/create', params=params, json=payload) as response:
+                if response.status == 200:
+                    data = await response.json()
+                    sandbox_id = data['sandbox_id']
+
+                    # Get sandbox info and store in _sandboxes
+                    sandbox_info = await self.get_sandbox_info(sandbox_id)
+                    if sandbox_info:
+                        self._sandboxes[sandbox_id] = sandbox_info
+
+                    logger.info(f'Created sandbox {sandbox_id} via HTTP API')
+                    return sandbox_id
+                else:
+                    error_data = await response.json()
+                    raise RuntimeError(f"HTTP {response.status}: {error_data.get('detail', 'Unknown error')}")
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error creating sandbox: {e}')
+            raise RuntimeError(f'Failed to create sandbox: {e}')
+
+    async def get_sandbox_info(self, sandbox_id: str) -> Optional[SandboxInfo]:
+        """Get sandbox information via HTTP API.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            Sandbox information or None if not found
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        try:
+            # Match server's endpoint format: GET /sandbox/{sandbox_id}
+            async with self._session.get(f'{self.base_url}/sandbox/{sandbox_id}') as response:
+                if response.status == 200:
+                    data = await response.json()
+                    return SandboxInfo.model_validate(data)
+                elif response.status == 404:
+                    return None
+                else:
+                    error_data = await response.json()
+                    logger.error(f'Error getting sandbox info: HTTP {response.status}: {error_data}')
+                    return None
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error getting sandbox info: {e}')
+            return None
+
+    async def list_sandboxes(self, status_filter: Optional[SandboxStatus] = None) -> List[SandboxInfo]:
+        """List all sandboxes via HTTP API.
+
+        Args:
+            status_filter: Optional status filter
+
+        Returns:
+            List of sandbox information
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        params = {}
+        if status_filter:
+            params['status'] = status_filter.value
+
+        try:
+            # Match server's endpoint format: GET /sandboxes
+            async with self._session.get(f'{self.base_url}/sandboxes', params=params) as response:
+                if response.status == 200:
+                    data = await response.json()
+                    return [SandboxInfo.model_validate(item) for item in data]
+                else:
+                    error_data = await response.json()
+                    logger.error(f'Error listing sandboxes: HTTP {response.status}: {error_data}')
+                    return []
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error listing sandboxes: {e}')
+            return []
+
+    async def stop_sandbox(self, sandbox_id: str) -> bool:
+        """Stop a sandbox via HTTP API.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            True if stopped successfully, False if not found
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        try:
+            # Match server's endpoint format: POST /sandbox/{sandbox_id}/stop
+            async with self._session.post(f'{self.base_url}/sandbox/{sandbox_id}/stop') as response:
+                if response.status == 200:
+                    logger.info(f'Stopped sandbox {sandbox_id} via HTTP API')
+                    return True
+                elif response.status == 404:
+                    logger.warning(f'Sandbox {sandbox_id} not found for stopping')
+                    return False
+                else:
+                    error_data = await response.json()
+                    logger.error(f'Error stopping sandbox: HTTP {response.status}: {error_data}')
+                    return False
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error stopping sandbox: {e}')
+            return False
+
+    async def delete_sandbox(self, sandbox_id: str) -> bool:
+        """Delete a sandbox via HTTP API.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            True if deleted successfully, False if not found
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        try:
+            # Match server's endpoint format: DELETE /sandbox/{sandbox_id}
+            async with self._session.delete(f'{self.base_url}/sandbox/{sandbox_id}') as response:
+                if response.status == 200:
+                    # Remove from tracking when successfully deleted
+                    self._sandboxes.pop(sandbox_id, None)
+                    logger.info(f'Deleted sandbox {sandbox_id} via HTTP API')
+                    return True
+                elif response.status == 404:
+                    # Also remove from tracking if not found (already deleted)
+                    self._sandboxes.pop(sandbox_id, None)
+                    logger.warning(f'Sandbox {sandbox_id} not found for deletion')
+                    return False
+                else:
+                    error_data = await response.json()
+                    logger.error(f'Error deleting sandbox: HTTP {response.status}: {error_data}')
+                    return False
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error deleting sandbox: {e}')
+            return False
+
+    async def execute_tool(self, sandbox_id: str, tool_name: str, parameters: Dict[str, Any]) -> ToolResult:
+        """Execute tool in sandbox via HTTP API.
+
+        Args:
+            sandbox_id: Sandbox ID
+            tool_name: Tool name to execute
+            parameters: Tool parameters
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            ValueError: If sandbox or tool not found
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        # Create proper request object to match server expectations
+        request = ToolExecutionRequest(sandbox_id=sandbox_id, tool_name=tool_name, parameters=parameters)
+        payload = request.model_dump()
+
+        try:
+            # Match server's endpoint format: POST /sandbox/tool/execute
+            async with self._session.post(f'{self.base_url}/sandbox/tool/execute', json=payload) as response:
+                if response.status == 200:
+                    data = await response.json()
+                    return ToolResult.model_validate(data)
+                elif response.status == 404:
+                    error_data = await response.json()
+                    raise ValueError(error_data.get('detail', f'Sandbox {sandbox_id} not found'))
+                elif response.status == 500:
+                    error_data = await response.json()
+                    raise RuntimeError(error_data.get('detail', 'Internal server error'))
+                else:
+                    error_data = await response.json()
+                    raise RuntimeError(f"HTTP {response.status}: {error_data.get('detail', 'Unknown error')}")
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error executing tool: {e}')
+            raise RuntimeError(f'Failed to execute tool: {e}')
+
+    async def get_sandbox_tools(self, sandbox_id: str) -> Dict[str, Any]:
+        """Get available tools for a sandbox via HTTP API.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            List of available tool types
+
+        Raises:
+            ValueError: If sandbox not found
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        try:
+            # Match server's endpoint format: GET /sandbox/{sandbox_id}/tools
+            async with self._session.get(f'{self.base_url}/sandbox/{sandbox_id}/tools') as response:
+                if response.status == 200:
+                    data = await response.json()
+                    return data
+                elif response.status == 404:
+                    error_data = await response.json()
+                    raise ValueError(error_data.get('detail', f'Sandbox {sandbox_id} not found'))
+                else:
+                    error_data = await response.json()
+                    raise RuntimeError(f"HTTP {response.status}: {error_data.get('detail', 'Unknown error')}")
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error getting sandbox tools: {e}')
+            raise RuntimeError(f'Failed to get sandbox tools: {e}')
+
+    async def cleanup_all_sandboxes(self) -> None:
+        """Clean up all sandboxes created by this manager via HTTP API."""
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        try:
+            # Clean up only the sandboxes created by this manager
+            sandbox_ids = list(self._sandboxes.keys())
+            deleted_count = 0
+
+            for sandbox_id in sandbox_ids:
+                try:
+                    if await self.delete_sandbox(sandbox_id):
+                        deleted_count += 1
+                except Exception as e:
+                    logger.error(f'Error deleting sandbox {sandbox_id}: {e}')
+
+            logger.info(f'Cleaned up {deleted_count} sandboxes created by this manager')
+
+        except Exception as e:
+            logger.error(f'HTTP client error cleaning up sandboxes: {e}')
+
+    async def get_stats(self) -> Dict[str, Any]:
+        """Get server statistics via HTTP API.
+
+        Returns:
+            Server statistics dictionary
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        try:
+            # Get server stats
+            async with self._session.get(f'{self.base_url}/stats') as response:
+                if response.status == 200:
+                    server_stats = await response.json()
+                else:
+                    error_data = await response.json()
+                    logger.error(f'Error getting server stats: HTTP {response.status}: {error_data}')
+                    server_stats = {}
+
+            # Add local tracking stats
+            from collections import Counter
+            status_counter = Counter()
+            type_counter = Counter()
+
+            for sandbox_info in self._sandboxes.values():
+                status_counter[sandbox_info.status.value] += 1
+                type_counter[sandbox_info.sandbox_type.value] += 1
+
+            local_stats = {
+                'manager_type': 'http',
+                'base_url': self.base_url,
+                'tracked_sandboxes': len(self._sandboxes),
+                'tracked_status_counts': dict(status_counter),
+                'tracked_sandbox_types': dict(type_counter),
+                'running': self._running,
+            }
+
+            # Combine server and local stats
+            return {**server_stats, **local_stats}
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error getting server stats: {e}')
+            return {
+                'manager_type': 'http',
+                'base_url': self.base_url,
+                'tracked_sandboxes': len(self._sandboxes),
+                'running': self._running,
+                'error': str(e)
+            }
+
+    async def health_check(self) -> Dict[str, Any]:
+        """Perform health check via HTTP API.
+
+        Returns:
+            Health check result
+        """
+        if not self._session:
+            raise RuntimeError('Manager not started')
+
+        try:
+            # Match server's endpoint format: GET /health
+            async with self._session.get(f'{self.base_url}/health') as response:
+                if response.status == 200:
+                    return await response.json()
+                else:
+                    error_data = await response.json()
+                    logger.error(f'Health check failed: HTTP {response.status}: {error_data}')
+                    return {'healthy': False, 'error': f'HTTP {response.status}'}
+
+        except aiohttp.ClientError as e:
+            logger.error(f'HTTP client error during health check: {e}')
+            return {'healthy': False, 'error': str(e)}

ms_enclave/sandbox/manager/local_manager.py

@@ -0,0 +1,295 @@
+"""Sandbox environment manager."""
+
+import asyncio
+import time
+from collections import Counter
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional, Union
+
+from ms_enclave.utils import get_logger
+
+from ..boxes import Sandbox, SandboxFactory
+from ..model import SandboxConfig, SandboxInfo, SandboxStatus, SandboxType, ToolResult
+from .base import SandboxManager
+
+logger = get_logger()
+
+
+class LocalSandboxManager(SandboxManager):
+    """Manager for sandbox environments."""
+
+    def __init__(self, cleanup_interval: int = 300):  # 5 minutes
+        """Initialize sandbox manager.
+
+        Args:
+            cleanup_interval: Interval between cleanup runs in seconds
+        """
+        super().__init__()
+        self._cleanup_interval = cleanup_interval
+        self._cleanup_task: Optional[asyncio.Task] = None
+
+    async def start(self) -> None:
+        """Start the sandbox manager."""
+        if self._running:
+            return
+
+        self._running = True
+        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
+        logger.info('Local sandbox manager started')
+
+    async def stop(self) -> None:
+        """Stop the sandbox manager."""
+        if not self._running:
+            return
+
+        self._running = False
+
+        # Cancel cleanup task
+        if self._cleanup_task:
+            self._cleanup_task.cancel()
+            try:
+                await self._cleanup_task
+            except asyncio.CancelledError:
+                pass
+
+        # Stop and cleanup all sandboxes
+        await self.cleanup_all_sandboxes()
+        logger.info('Local sandbox manager stopped')
+
+    async def create_sandbox(
+        self,
+        sandbox_type: SandboxType,
+        config: Optional[Union[SandboxConfig, Dict]] = None,
+        sandbox_id: Optional[str] = None
+    ) -> str:
+        """Create a new sandbox.
+
+        Args:
+            sandbox_type: Type of sandbox to create
+            config: Sandbox configuration
+            sandbox_id: Optional sandbox ID
+
+        Returns:
+            Sandbox ID
+
+        Raises:
+            ValueError: If sandbox type is not supported
+            RuntimeError: If sandbox creation fails
+        """
+        try:
+            # Create sandbox instance
+            sandbox = SandboxFactory.create_sandbox(sandbox_type, config, sandbox_id)
+
+            # Start the sandbox
+            await sandbox.start()
+
+            # Store sandbox
+            self._sandboxes[sandbox.id] = sandbox
+
+            logger.info(f'Created and started sandbox {sandbox.id} of type {sandbox_type}')
+            return sandbox.id
+
+        except Exception as e:
+            logger.error(f'Failed to create sandbox of type {sandbox_type}: {e}')
+            raise RuntimeError(f'Failed to create sandbox: {e}')
+
+    async def get_sandbox(self, sandbox_id: str) -> Optional[Sandbox]:
+        """Get sandbox by ID.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            Sandbox instance or None if not found
+        """
+        return self._sandboxes.get(sandbox_id)
+
+    async def get_sandbox_info(self, sandbox_id: str) -> Optional[SandboxInfo]:
+        """Get sandbox information.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            Sandbox information or None if not found
+        """
+        sandbox = self._sandboxes.get(sandbox_id)
+        if sandbox:
+            return sandbox.get_info()
+        return None
+
+    async def list_sandboxes(self, status_filter: Optional[SandboxStatus] = None) -> List[SandboxInfo]:
+        """List all sandboxes.
+
+        Args:
+            status_filter: Optional status filter
+
+        Returns:
+            List of sandbox information
+        """
+        result = []
+        for sandbox in self._sandboxes.values():
+            info = sandbox.get_info()
+            if status_filter is None or info.status == status_filter:
+                result.append(info)
+        return result
+
+    async def stop_sandbox(self, sandbox_id: str) -> bool:
+        """Stop a sandbox.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            True if stopped successfully, False if not found
+        """
+        sandbox = self._sandboxes.get(sandbox_id)
+        if not sandbox:
+            logger.warning(f'Sandbox {sandbox_id} not found for stopping')
+            return False
+
+        try:
+            await sandbox.stop()
+            logger.info(f'Stopped sandbox {sandbox_id}')
+            return True
+        except Exception as e:
+            logger.error(f'Error stopping sandbox {sandbox_id}: {e}')
+            return False
+
+    async def delete_sandbox(self, sandbox_id: str) -> bool:
+        """Delete a sandbox.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            True if deleted successfully, False if not found
+        """
+        sandbox = self._sandboxes.get(sandbox_id)
+        if not sandbox:
+            logger.warning(f'Sandbox {sandbox_id} not found for deletion')
+            return False
+
+        try:
+            await sandbox.stop()
+            del self._sandboxes[sandbox_id]
+            logger.info(f'Deleted sandbox {sandbox_id}')
+            return True
+        except Exception as e:
+            logger.error(f'Error deleting sandbox {sandbox_id}: {e}')
+            return False
+
+    async def execute_tool(self, sandbox_id: str, tool_name: str, parameters: Dict[str, Any]) -> ToolResult:
+        """Execute tool in sandbox.
+
+        Args:
+            sandbox_id: Sandbox ID
+            tool_name: Tool name to execute
+            parameters: Tool parameters
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            ValueError: If sandbox or tool not found
+        """
+        sandbox = self._sandboxes.get(sandbox_id)
+        if not sandbox:
+            raise ValueError(f'Sandbox {sandbox_id} not found')
+
+        if sandbox.status != SandboxStatus.RUNNING:
+            raise ValueError(f'Sandbox {sandbox_id} is not running (status: {sandbox.status})')
+
+        result = await sandbox.execute_tool(tool_name, parameters)
+        return result
+
+    async def get_sandbox_tools(self, sandbox_id: str) -> Dict[str, Any]:
+        """Get available tools for a sandbox.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            Dictionary of available tool types
+
+        Raises:
+            ValueError: If sandbox not found
+        """
+        sandbox = self._sandboxes.get(sandbox_id)
+        if not sandbox:
+            raise ValueError(f'Sandbox {sandbox_id} not found')
+
+        return sandbox.get_available_tools()
+
+    async def cleanup_all_sandboxes(self) -> None:
+        """Clean up all sandboxes."""
+        sandbox_ids = list(self._sandboxes.keys())
+        logger.info(f'Cleaning up {len(sandbox_ids)} sandboxes')
+
+        for sandbox_id in sandbox_ids:
+            try:
+                await self.delete_sandbox(sandbox_id)
+            except Exception as e:
+                logger.error(f'Error cleaning up sandbox {sandbox_id}: {e}')
+
+    async def get_stats(self) -> Dict[str, Any]:
+        """Get manager statistics.
+
+        Returns:
+            Statistics dictionary
+        """
+        status_counter = Counter()
+        type_counter = Counter()
+
+        for sandbox in self._sandboxes.values():
+            status_counter[sandbox.status.value] += 1
+            type_counter[sandbox.sandbox_type.value] += 1
+
+        stats = {
+            'manager_type': 'local',
+            'total_sandboxes': len(self._sandboxes),
+            'status_counts': dict(status_counter),
+            'sandbox_types': dict(type_counter),
+            'running': self._running,
+            'cleanup_interval': self._cleanup_interval,
+        }
+
+        return stats
+
+    async def _cleanup_loop(self) -> None:
+        """Background cleanup loop."""
+        while self._running:
+            try:
+                await self._cleanup_expired_sandboxes()
+                await asyncio.sleep(self._cleanup_interval)
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f'Error in cleanup loop: {e}')
+                await asyncio.sleep(self._cleanup_interval)
+
+    async def _cleanup_expired_sandboxes(self) -> None:
+        """Clean up expired sandboxes."""
+        current_time = datetime.now()
+        expired_sandboxes = []
+
+        for sandbox_id, sandbox in self._sandboxes.items():
+            # Check if sandbox is in error state or stopped for too long
+            if sandbox.status in [SandboxStatus.ERROR, SandboxStatus.STOPPED]:
+                # Clean up after 1 hour
+                if current_time - sandbox.updated_at > timedelta(hours=1):
+                    expired_sandboxes.append(sandbox_id)
+            # Check for very old sandboxes (48 hours)
+            elif current_time - sandbox.created_at > timedelta(hours=48):
+                expired_sandboxes.append(sandbox_id)
+
+        # Clean up expired sandboxes
+        if expired_sandboxes:
+            logger.info(f'Found {len(expired_sandboxes)} expired sandboxes to clean up')
+
+        for sandbox_id in expired_sandboxes:
+            try:
+                logger.info(f'Cleaning up expired sandbox: {sandbox_id}')
+                await self.delete_sandbox(sandbox_id)
+            except Exception as e:
+                logger.error(f'Error cleaning up expired sandbox {sandbox_id}: {e}')