ms-enclave 0.0.3__py3-none-any.whl

ms_enclave/sandbox/tools/sandbox_tools/file_operation.py

@@ -0,0 +1,331 @@
+"""File operation tool for reading and writing files."""
+
+import io
+import os
+import tarfile
+import uuid
+from typing import TYPE_CHECKING, Literal, Optional
+
+from ms_enclave.sandbox.model import ExecutionStatus, SandboxType, ToolResult
+from ms_enclave.sandbox.tools.base import register_tool
+from ms_enclave.sandbox.tools.sandbox_tool import SandboxTool
+from ms_enclave.sandbox.tools.tool_info import ToolParams
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import Sandbox
+
+
+@register_tool('file_operation')
+class FileOperation(SandboxTool):
+    """Tool for performing file operations within Docker containers.
+
+    Supports read, write, create, delete, list, and exists operations.
+    Uses temporary files and copy strategy to avoid permission issues
+    with mounted directories.
+    """
+
+    _name = 'file_operation'
+    _sandbox_type = SandboxType.DOCKER
+    _description = 'Perform file operations like read, write, delete, and list files'
+    _parameters = ToolParams(
+        type='object',
+        properties={
+            'operation': {
+                'type': 'string',
+                'description': 'Type of file operation to perform',
+                'enum': ['create', 'read', 'write', 'delete', 'list', 'exists']
+            },
+            'file_path': {
+                'type': 'string',
+                'description': 'Path to the file or directory'
+            },
+            'content': {
+                'type': 'string',
+                'description': 'Content to write to file (only for write operation)'
+            },
+            'encoding': {
+                'type': 'string',
+                'description': 'File encoding',
+                'default': 'utf-8'
+            }
+        },
+        required=['operation', 'file_path']
+    )
+
+    async def execute(
+        self,
+        sandbox_context: 'Sandbox',
+        operation: str,
+        file_path: str,
+        content: Optional[str] = None,
+        encoding: str = 'utf-8'
+    ) -> ToolResult:
+        """Perform file operations in the Docker container.
+
+        Args:
+            sandbox_context: The sandbox instance to execute operations in
+            operation: Type of operation (read, write, create, delete, list, exists)
+            file_path: Path to the file or directory
+            content: Content to write (required for write/create operations)
+            encoding: File encoding for read/write operations
+
+        Returns:
+            ToolResult with operation status and output/error information
+        """
+
+        # Validate file path is provided and not empty
+        if not file_path.strip():
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No file path provided'
+            )
+
+        try:
+            # Route to appropriate operation handler
+            if operation == 'read':
+                return await self._read_file(sandbox_context, file_path, encoding)
+            elif operation == 'write':
+                # Write operation requires content parameter
+                if content is None:
+                    return ToolResult(
+                        tool_name=self.name,
+                        status=ExecutionStatus.ERROR,
+                        output='',
+                        error='Content is required for write operation'
+                    )
+                return await self._write_file(sandbox_context, file_path, content, encoding)
+            elif operation == 'delete':
+                return await self._delete_file(sandbox_context, file_path)
+            elif operation == 'list':
+                return await self._list_directory(sandbox_context, file_path)
+            elif operation == 'exists':
+                return await self._check_exists(sandbox_context, file_path)
+            elif operation == 'create':
+                # Create operation with empty content if none provided
+                if content is None:
+                    content = ''
+                return await self._write_file(sandbox_context, file_path, content, encoding)
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=f'Unknown operation: {operation}'
+                )
+
+        except Exception as e:
+            # Catch-all error handler for unexpected exceptions
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Operation failed: {str(e)}'
+            )
+
+    async def _read_file(self, sandbox_context: 'Sandbox', file_path: str, encoding: str) -> ToolResult:
+        """Read file content from the container using cat command.
+
+        Args:
+            sandbox_context: Sandbox instance
+            file_path: Path to file to read
+            encoding: File encoding (currently not used by cat command)
+
+        Returns:
+            ToolResult with file content or error message
+        """
+        try:
+            # Use cat command to read file content - handles most file types well
+            result = await sandbox_context.execute_command(f'cat "{file_path}"')
+
+            if result.exit_code == 0:
+                return ToolResult(tool_name=self.name, status=ExecutionStatus.SUCCESS, output=result.stdout, error=None)
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=result.stderr or f'Failed to read file: {file_path}'
+                )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Read failed: {str(e)}'
+            )
+
+    async def _write_file(self, sandbox_context: 'Sandbox', file_path: str, content: str, encoding: str) -> ToolResult:
+        """Write content to a file in the container using temp file strategy.
+
+        This method uses a two-step process to avoid permission issues:
+        1. Write content to a temporary file in /tmp (always writable)
+        2. Copy the temp file to the target location using cp command
+        3. Clean up the temporary file
+
+        Args:
+            sandbox_context: Sandbox instance
+            file_path: Target file path to write to
+            content: Content to write to file
+            encoding: File encoding for content
+
+        Returns:
+            ToolResult indicating success or failure
+        """
+        try:
+            # Create target directory structure if it doesn't exist
+            dir_path = os.path.dirname(file_path)
+            if dir_path:
+                await sandbox_context.execute_command(f'mkdir -p "{dir_path}"')
+
+            # Generate unique temporary file name to avoid conflicts
+            temp_file = f'/tmp/file_op_{uuid.uuid4().hex}'
+
+            # Step 1: Write content to temporary location using tar archive
+            await self._write_file_to_container(sandbox_context, temp_file, content, encoding)
+
+            # Step 2: Copy from temp to target location using cp command
+            # This handles permission issues better than direct tar extraction
+            copy_result = await sandbox_context.execute_command(f'cp "{temp_file}" "{file_path}"')
+
+            # Step 3: Clean up temporary file regardless of copy result
+            await sandbox_context.execute_command(f'rm -f "{temp_file}"')
+
+            # Check if copy operation succeeded
+            if copy_result.exit_code != 0:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=f'Failed to copy file to target location: {copy_result.stderr or "Unknown error"}'
+                )
+
+            return ToolResult(
+                tool_name=self.name,
+                status=ExecutionStatus.SUCCESS,
+                output=f'File written successfully: {file_path}',
+                error=None
+            )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Write failed: {str(e)}'
+            )
+
+    async def _delete_file(self, sandbox_context: 'Sandbox', file_path: str) -> ToolResult:
+        """Delete a file or directory from the container.
+
+        Uses 'rm -rf' to handle both files and directories recursively.
+
+        Args:
+            sandbox_context: Sandbox instance
+            file_path: Path to file or directory to delete
+
+        Returns:
+            ToolResult indicating success or failure
+        """
+        try:
+            # Use rm -rf to handle both files and directories
+            result = await sandbox_context.execute_command(f'rm -rf "{file_path}"')
+
+            if result.exit_code == 0:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.SUCCESS,
+                    output=f'Successfully deleted: {file_path}',
+                    error=None
+                )
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=result.stderr or f'Failed to delete: {file_path}'
+                )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Delete failed: {str(e)}'
+            )
+
+    async def _list_directory(self, sandbox_context: 'Sandbox', dir_path: str) -> ToolResult:
+        """List contents of a directory with detailed information.
+
+        Uses 'ls -la' to show permissions, ownership, size, and timestamps.
+
+        Args:
+            sandbox_context: Sandbox instance
+            dir_path: Path to directory to list
+
+        Returns:
+            ToolResult with directory listing or error message
+        """
+        try:
+            # Use ls -la for detailed directory listing
+            result = await sandbox_context.execute_command(f'ls -la "{dir_path}"')
+
+            if result.exit_code == 0:
+                return ToolResult(tool_name=self.name, status=ExecutionStatus.SUCCESS, output=result.stdout, error=None)
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=result.stderr or f'Failed to list directory: {dir_path}'
+                )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'List failed: {str(e)}'
+            )
+
+    async def _check_exists(self, sandbox_context: 'Sandbox', file_path: str) -> ToolResult:
+        """Check if a file or directory exists.
+
+        Uses 'test -e' command which returns exit code 0 if path exists.
+
+        Args:
+            sandbox_context: Sandbox instance
+            file_path: Path to check for existence
+
+        Returns:
+            ToolResult with existence status message
+        """
+        try:
+            # Use test -e to check existence (works for files and directories)
+            result = await sandbox_context.execute_command(f'test -e "{file_path}"')
+
+            # test command returns 0 if path exists, non-zero otherwise
+            exists = result.exit_code == 0
+            return ToolResult(
+                tool_name=self.name,
+                status=ExecutionStatus.SUCCESS,
+                output=f'{"exists" if exists else "does not exist"}',
+                error=None
+            )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Exists check failed: {str(e)}'
+            )
+
+    async def _write_file_to_container(
+        self, sandbox_context: 'Sandbox', file_path: str, content: str, encoding: str
+    ) -> None:
+        """Write content to a file in the container using tar archive method.
+
+        This is a low-level method that creates a tar archive containing the file
+        and extracts it to the container. Used internally by _write_file.
+
+        Args:
+            sandbox_context: Sandbox instance with container access
+            file_path: Target file path in container
+            content: File content as string
+            encoding: Text encoding for content conversion
+
+        Raises:
+            Exception: If tar creation or container extraction fails
+        """
+        # Create a tar archive in memory to transfer file content
+        tar_stream = io.BytesIO()
+        tar = tarfile.TarFile(fileobj=tar_stream, mode='w')
+
+        # Encode content using specified encoding and create tar entry
+        file_data = content.encode(encoding)
+        tarinfo = tarfile.TarInfo(name=os.path.basename(file_path))
+        tarinfo.size = len(file_data)
+        tar.addfile(tarinfo, io.BytesIO(file_data))
+        tar.close()
+
+        # Extract tar archive to container filesystem
+        # Note: This writes to the directory containing the target file
+        tar_stream.seek(0)
+        sandbox_context.container.put_archive(os.path.dirname(file_path) or '/', tar_stream.getvalue())

ms_enclave/sandbox/tools/sandbox_tools/notebook_executor.py

@@ -0,0 +1,167 @@
+"""Notebook code execution tool for Jupyter kernels."""
+
+import json
+import time
+import uuid
+from typing import TYPE_CHECKING, Optional
+
+from ms_enclave.sandbox.model import CommandResult, ExecutionStatus, SandboxType, ToolResult
+from ms_enclave.sandbox.tools.base import Tool, register_tool
+from ms_enclave.sandbox.tools.sandbox_tool import SandboxTool
+from ms_enclave.sandbox.tools.tool_info import ToolParams
+from ms_enclave.utils import get_logger
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import Sandbox
+
+logger = get_logger()
+
+
+@register_tool('notebook_executor')
+class NotebookExecutor(SandboxTool):
+
+    _name = 'notebook_executor'
+    _sandbox_type = SandboxType.DOCKER_NOTEBOOK
+    _description = 'Execute Python code in a Jupyter kernel environment'
+    _parameters = ToolParams(
+        type='object',
+        properties={
+            'code': {
+                'type': 'string',
+                'description': 'Python code to execute in the notebook kernel'
+            },
+            'timeout': {
+                'type': 'integer',
+                'description': 'Execution timeout in seconds',
+                'default': 30
+            }
+        },
+        required=['code']
+    )
+
+    async def execute(self, sandbox_context: 'Sandbox', code: str, timeout: Optional[int] = 30) -> ToolResult:
+        """Execute Python code in the Jupyter kernel."""
+
+        if not code.strip():
+            return ToolResult(tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No code provided')
+
+        try:
+            # Execute code using the sandbox's Jupyter kernel
+            result = await self._execute_in_kernel(sandbox_context, code, timeout)
+
+            if result.exit_code == 0:
+                status = ExecutionStatus.SUCCESS
+            else:
+                status = ExecutionStatus.ERROR
+
+            return ToolResult(
+                tool_name=self.name,
+                status=status,
+                output=result.stdout,
+                error=result.stderr if result.stderr else None
+            )
+
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Execution failed: {str(e)}'
+            )
+
+    async def _execute_in_kernel(self, sandbox_context: 'Sandbox', code: str, timeout: Optional[int]):
+        """Execute code in the Jupyter kernel via websocket."""
+
+        # Check if sandbox has the required Jupyter components
+        if not hasattr(sandbox_context, 'ws') or not hasattr(sandbox_context, 'kernel_id'):
+            raise RuntimeError('Sandbox does not have Jupyter kernel setup')
+
+        if not sandbox_context.ws or not sandbox_context.kernel_id:
+            raise RuntimeError('Jupyter kernel is not ready')
+
+        # Send execute request
+        msg_id = self._send_execute_request(sandbox_context, code)
+
+        outputs = []
+        result = None
+        error_occurred = False
+        error_msg = ''
+        actual_timeout = timeout or 30
+        start_time = time.time()
+
+        while True:
+            elapsed = time.time() - start_time
+            if elapsed >= actual_timeout:
+                error_occurred = True
+                error_msg = f'Execution timed out after {actual_timeout} seconds'
+                logger.error(error_msg)
+                break
+
+            try:
+                # Set a short timeout for recv to allow periodic timeout check
+                sandbox_context.ws.settimeout(min(1, actual_timeout - elapsed))
+                msg = json.loads(sandbox_context.ws.recv())
+                parent_msg_id = msg.get('parent_header', {}).get('msg_id')
+
+                # Skip unrelated messages
+                if parent_msg_id != msg_id:
+                    continue
+
+                msg_type = msg.get('msg_type', '')
+                msg_content = msg.get('content', {})
+
+                if msg_type == 'stream':
+                    outputs.append(msg_content['text'])
+                elif msg_type == 'execute_result':
+                    result = msg_content['data'].get('text/plain', '')
+                elif msg_type == 'error':
+                    error_occurred = True
+                    error_msg = '\n'.join(msg_content.get('traceback', []))
+                    outputs.append(error_msg)
+                elif msg_type == 'status' and msg_content.get('execution_state') == 'idle':
+                    break
+
+            except Exception as e:
+                logger.error(f'Error receiving message: {e}')
+                error_occurred = True
+                error_msg = str(e)
+                break
+
+        output_text = ''.join(outputs)
+        if result:
+            output_text += f'\nResult: {result}'
+        if error_msg and error_msg not in output_text:
+            output_text += f'\n{error_msg}'
+
+        return CommandResult(
+            status=ExecutionStatus.SUCCESS if not error_occurred else ExecutionStatus.ERROR,
+            command=code,
+            exit_code=1 if error_occurred else 0,
+            stdout=output_text if not error_occurred else '',
+            stderr=output_text if error_occurred else ''
+        )
+
+    def _send_execute_request(self, sandbox_context: 'Sandbox', code: str) -> str:
+        """Send code execution request to kernel."""
+        # Generate a unique message ID
+        msg_id = str(uuid.uuid4())
+
+        # Create execute request
+        execute_request = {
+            'header': {
+                'msg_id': msg_id,
+                'username': 'anonymous',
+                'session': str(uuid.uuid4()),
+                'msg_type': 'execute_request',
+                'version': '5.0',
+            },
+            'parent_header': {},
+            'metadata': {},
+            'content': {
+                'code': code,
+                'silent': False,
+                'store_history': True,
+                'user_expressions': {},
+                'allow_stdin': False,
+            },
+        }
+
+        sandbox_context.ws.send(json.dumps(execute_request))
+        return msg_id

ms_enclave/sandbox/tools/sandbox_tools/python_executor.py

@@ -0,0 +1,87 @@
+"""Python code execution tool."""
+
+import os
+import uuid
+from typing import TYPE_CHECKING, Optional
+
+from ms_enclave.sandbox.model import ExecutionStatus, SandboxType, ToolResult
+from ms_enclave.sandbox.tools.base import Tool, register_tool
+from ms_enclave.sandbox.tools.sandbox_tool import SandboxTool
+from ms_enclave.sandbox.tools.tool_info import ToolParams
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import DockerSandbox
+
+
+@register_tool('python_executor')
+class PythonExecutor(SandboxTool):
+
+    _name = 'python_executor'
+    _sandbox_type = SandboxType.DOCKER
+    _description = 'Execute Python code in an isolated environment using IPython'
+    _parameters = ToolParams(
+        type='object',
+        properties={
+            'code': {
+                'type': 'string',
+                'description': 'Python code to execute'
+            },
+            'timeout': {
+                'type': 'integer',
+                'description': 'Execution timeout in seconds',
+                'default': 30
+            }
+        },
+        required=['code']
+    )
+
+    async def execute(self, sandbox_context: 'DockerSandbox', code: str, timeout: Optional[int] = 30) -> ToolResult:
+        """Execute Python code by writing to a temporary file and executing it."""
+
+        script_basename = f'exec_script_{uuid.uuid4().hex}.py'
+        script_path = f'/tmp/{script_basename}'
+
+        if not code.strip():
+            return ToolResult(tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No code provided')
+
+        try:
+
+            # Write script to container to avoid long code errors
+            await self._write_file_to_container(sandbox_context, script_path, code)
+
+            # Execute using python
+            command = f'python {script_path}'
+            result = await sandbox_context.execute_command(command, timeout=timeout)
+
+            if result.exit_code == 0:
+                status = ExecutionStatus.SUCCESS
+            else:
+                status = ExecutionStatus.ERROR
+
+            return ToolResult(
+                tool_name=self.name,
+                status=status,
+                output=result.stdout,
+                error=result.stderr if result.stderr else None
+            )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Execution failed: {str(e)}'
+            )
+
+    async def _write_file_to_container(self, sandbox_context: 'DockerSandbox', file_path: str, content: str) -> None:
+        """Write content to a file in the container."""
+        import io
+        import tarfile
+
+        # Create a tar archive in memory using context managers
+        with io.BytesIO() as tar_stream:
+            with tarfile.TarFile(fileobj=tar_stream, mode='w') as tar:
+                file_data = content.encode('utf-8')
+                tarinfo = tarfile.TarInfo(name=os.path.basename(file_path))
+                tarinfo.size = len(file_data)
+                tar.addfile(tarinfo, io.BytesIO(file_data))
+
+                # Reset stream position and put archive into container
+                tar_stream.seek(0)
+                sandbox_context.container.put_archive(os.path.dirname(file_path), tar_stream.getvalue())

ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py

@@ -0,0 +1,72 @@
+"""Shell command execution tool."""
+
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from ms_enclave.sandbox.model import ExecutionStatus, SandboxType, ToolResult
+from ms_enclave.sandbox.tools.base import register_tool
+from ms_enclave.sandbox.tools.sandbox_tool import SandboxTool
+from ms_enclave.sandbox.tools.tool_info import ToolParams
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import Sandbox
+
+
+@register_tool('shell_executor')
+class ShellExecutor(SandboxTool):
+
+    _name = 'shell_executor'
+    _sandbox_type = SandboxType.DOCKER
+    _description = 'Execute shell commands in an isolated environment'
+    _parameters = ToolParams(
+        type='object',
+        properties={
+            'command': {
+                'anyOf': [{
+                    'type': 'string',
+                    'description': 'Shell command to execute'
+                }, {
+                    'type': 'array',
+                    'items': {
+                        'type': 'string'
+                    },
+                    'description': 'List of shell command arguments to execute'
+                }]
+            },
+            'timeout': {
+                'type': 'integer',
+                'description': 'Execution timeout in seconds',
+                'default': 30
+            }
+        },
+        required=['command']
+    )
+
+    async def execute(
+        self, sandbox_context: 'Sandbox', command: Union[str, List[str]], timeout: Optional[int] = 30
+    ) -> ToolResult:
+        """Execute shell command in the Docker container."""
+
+        if not command or (isinstance(command, str) and not command.strip()):
+            return ToolResult(tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No command provided')
+        try:
+            result = await sandbox_context.execute_command(command, timeout=timeout)
+
+            if result.exit_code == 0:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.SUCCESS,
+                    output=result.stdout,
+                    error=result.stderr if result.stderr else None
+                )
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output=result.stdout,
+                    error=result.stderr if result.stderr else f'Command failed with exit code {result.exit_code}'
+                )
+
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Execution failed: {str(e)}'
+            )