skilllite 0.1.0__py3-none-any.whl

skilllite/mcp/server.py

@@ -0,0 +1,734 @@
+"""
+MCP Server implementation for SkillLite.
+
+This module provides a secure code execution sandbox server implementing
+the Model Context Protocol (MCP).
+
+Security Model:
+    The MCP server implements a two-phase execution model for security:
+    
+    1. **scan_code**: First, scan the code for security issues. This returns
+       a detailed report of any potential risks found in the code.
+    
+    2. **execute_code**: Then, execute the code. If security issues were found,
+       the caller must explicitly set `confirmed=true` to proceed.
+    
+    This allows LLM clients to present security warnings to users and get
+    explicit confirmation before executing potentially dangerous code.
+
+Environment Variables:
+    SKILLBOX_SANDBOX_LEVEL: Default sandbox level (1/2/3, default: 3)
+    SKILLBOX_PATH: Path to skillbox binary
+    MCP_SANDBOX_TIMEOUT: Execution timeout in seconds (default: 30)
+"""
+
+import asyncio
+import hashlib
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+
+def _load_dotenv():
+    """Load environment variables from .env file if it exists."""
+    # Try to find .env file in current directory or parent directories
+    current_dir = Path.cwd()
+    
+    for search_dir in [current_dir] + list(current_dir.parents)[:3]:
+        env_file = search_dir / ".env"
+        if env_file.exists():
+            try:
+                with open(env_file, "r") as f:
+                    for line in f:
+                        line = line.strip()
+                        # Skip comments and empty lines
+                        if not line or line.startswith("#"):
+                            continue
+                        # Parse KEY=VALUE
+                        if "=" in line:
+                            key, _, value = line.partition("=")
+                            key = key.strip()
+                            value = value.strip()
+                            # Remove quotes if present
+                            if value and value[0] in ('"', "'") and value[-1] == value[0]:
+                                value = value[1:-1]
+                            # Only set if not already in environment
+                            if key and key not in os.environ:
+                                os.environ[key] = value
+                return True
+            except Exception:
+                pass
+    return False
+
+
+# Load .env file on module import
+_load_dotenv()
+
+try:
+    from mcp.server import Server
+    from mcp.server.stdio import stdio_server
+    from mcp.types import (
+        TextContent,
+        Tool,
+        CallToolResult,
+    )
+    MCP_AVAILABLE = True
+except ImportError:
+    MCP_AVAILABLE = False
+    # Define stub types for when MCP is not available
+    Server = None  # type: ignore
+    stdio_server = None  # type: ignore
+    TextContent = None  # type: ignore
+    Tool = None  # type: ignore
+    CallToolResult = None  # type: ignore
+
+
+class SecurityScanResult:
+    """Result of a security scan."""
+    
+    def __init__(
+        self,
+        is_safe: bool,
+        issues: List[Dict[str, Any]],
+        scan_id: str,
+        code_hash: str,
+        high_severity_count: int = 0,
+        medium_severity_count: int = 0,
+        low_severity_count: int = 0,
+    ):
+        self.is_safe = is_safe
+        self.issues = issues
+        self.scan_id = scan_id
+        self.code_hash = code_hash
+        self.high_severity_count = high_severity_count
+        self.medium_severity_count = medium_severity_count
+        self.low_severity_count = low_severity_count
+        self.timestamp = time.time()
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "is_safe": self.is_safe,
+            "issues": self.issues,
+            "scan_id": self.scan_id,
+            "code_hash": self.code_hash,
+            "high_severity_count": self.high_severity_count,
+            "medium_severity_count": self.medium_severity_count,
+            "low_severity_count": self.low_severity_count,
+            "requires_confirmation": self.high_severity_count > 0,
+        }
+    
+    def format_report(self) -> str:
+        """Format a human-readable security report."""
+        if not self.issues:
+            return "✅ Security scan passed. No issues found."
+        
+        lines = [
+            f"📋 Security Scan Report (ID: {self.scan_id[:8]})",
+            f"   Found {len(self.issues)} item(s) for review:",
+            "",
+        ]
+        
+        severity_icons = {
+            "Critical": "🔴",
+            "High": "🟠",
+            "Medium": "🟡",
+            "Low": "🟢",
+        }
+        
+        for idx, issue in enumerate(self.issues, 1):
+            severity = issue.get("severity", "Medium")
+            icon = severity_icons.get(severity, "⚪")
+            lines.append(f"  {icon} #{idx} [{severity}] {issue.get('issue_type', 'Unknown')}")
+            lines.append(f"     ├─ Rule: {issue.get('rule_id', 'N/A')}")
+            lines.append(f"     ├─ Line {issue.get('line_number', '?')}: {issue.get('description', '')}")
+            lines.append(f"     └─ Code: {issue.get('code_snippet', '')[:60]}...")
+            lines.append("")
+        
+        if self.high_severity_count > 0:
+            lines.append("⚠️  High severity issues found. Confirmation required to execute.")
+            lines.append(f"   To proceed, call execute_code with confirmed=true and scan_id=\"{self.scan_id}\"")
+        else:
+            lines.append("ℹ️  Only low/medium severity issues found. Safe to execute.")
+        
+        return "\n".join(lines)
+
+
+class SandboxExecutor:
+    """Secure code execution sandbox using Rust skillbox."""
+    
+    # Cache scan results for confirmation flow (scan_id -> result)
+    _scan_cache: Dict[str, SecurityScanResult] = {}
+    # Cache expiry time in seconds
+    SCAN_CACHE_TTL = 300  # 5 minutes
+    
+    def __init__(self):
+        from ..sandbox.skillbox import find_binary
+        
+        self.skillbox_path = os.getenv("SKILLBOX_PATH") or find_binary() or "./skillbox/target/release/skillbox"
+        self.timeout = int(os.getenv("MCP_SANDBOX_TIMEOUT", "30"))
+        self.runtime_available = os.path.exists(self.skillbox_path) and os.access(self.skillbox_path, os.X_OK)
+        
+        # Read default sandbox level from environment variable
+        # SKILLBOX_SANDBOX_LEVEL: 1=no sandbox, 2=sandbox only, 3=sandbox+scan (default)
+        default_level = os.getenv("SKILLBOX_SANDBOX_LEVEL", "3")
+        try:
+            self.default_sandbox_level = int(default_level)
+            if self.default_sandbox_level not in [1, 2, 3]:
+                self.default_sandbox_level = 3
+        except ValueError:
+            self.default_sandbox_level = 3
+    
+    def _generate_code_hash(self, language: str, code: str) -> str:
+        """Generate a hash of the code for verification."""
+        content = f"{language}:{code}"
+        return hashlib.sha256(content.encode()).hexdigest()
+    
+    def _generate_scan_id(self, code_hash: str) -> str:
+        """Generate a unique scan ID."""
+        timestamp = str(time.time())
+        return hashlib.sha256(f"{code_hash}:{timestamp}".encode()).hexdigest()[:16]
+    
+    def _cleanup_expired_scans(self):
+        """Remove expired scan results from cache."""
+        current_time = time.time()
+        expired_ids = [
+            scan_id for scan_id, result in self._scan_cache.items()
+            if current_time - result.timestamp > self.SCAN_CACHE_TTL
+        ]
+        for scan_id in expired_ids:
+            del self._scan_cache[scan_id]
+    
+    def _create_temp_skill(self, language: str, code: str) -> Tuple[str, str]:
+        """Create a temporary skill directory with the code file."""
+        skill_dir = tempfile.mkdtemp(prefix="mcp_skill_")
+        
+        # Create scripts subdirectory (required by skillbox)
+        scripts_dir = os.path.join(skill_dir, "scripts")
+        os.makedirs(scripts_dir, exist_ok=True)
+        
+        ext = self.get_file_extension(language)
+        entry_point = f"scripts/main.{ext}"
+        
+        skill_md_content = f"""---
+name: mcp-execution
+entry_point: {entry_point}
+language: {language}
+description: MCP code execution skill
+network:
+  enabled: true
+---
+
+This skill executes code from MCP.
+"""
+        with open(os.path.join(skill_dir, "SKILL.md"), "w") as f:
+            f.write(skill_md_content)
+        
+        code_file = os.path.join(scripts_dir, f"main.{ext}")
+        with open(code_file, "w") as f:
+            f.write(code)
+        os.chmod(code_file, 0o755)
+        
+        return skill_dir, code_file
+    
+    def scan_code(self, language: str, code: str) -> SecurityScanResult:
+        """Scan code for security issues without executing it."""
+        if not self.runtime_available:
+            return SecurityScanResult(
+                is_safe=False,
+                issues=[{"severity": "Critical", "issue_type": "RuntimeError", 
+                        "description": f"skillbox not found at {self.skillbox_path}",
+                        "rule_id": "system", "line_number": 0, "code_snippet": ""}],
+                scan_id="error",
+                code_hash="",
+                high_severity_count=1,
+            )
+        
+        self._cleanup_expired_scans()
+        code_hash = self._generate_code_hash(language, code)
+        scan_id = self._generate_scan_id(code_hash)
+        
+        try:
+            skill_dir, code_file = self._create_temp_skill(language, code)
+            
+            try:
+                result = subprocess.run(
+                    [self.skillbox_path, "security-scan", code_file],
+                    capture_output=True,
+                    text=True,
+                    timeout=30
+                )
+                
+                issues = self._parse_scan_output(result.stdout + result.stderr)
+                high_count = sum(1 for i in issues if i.get("severity") in ["Critical", "High"])
+                medium_count = sum(1 for i in issues if i.get("severity") == "Medium")
+                low_count = sum(1 for i in issues if i.get("severity") == "Low")
+                
+                scan_result = SecurityScanResult(
+                    is_safe=high_count == 0,
+                    issues=issues,
+                    scan_id=scan_id,
+                    code_hash=code_hash,
+                    high_severity_count=high_count,
+                    medium_severity_count=medium_count,
+                    low_severity_count=low_count,
+                )
+                
+                self._scan_cache[scan_id] = scan_result
+                return scan_result
+                
+            finally:
+                shutil.rmtree(skill_dir, ignore_errors=True)
+                
+        except subprocess.TimeoutExpired:
+            return SecurityScanResult(
+                is_safe=False,
+                issues=[{"severity": "Critical", "issue_type": "Timeout",
+                        "description": "Security scan timed out",
+                        "rule_id": "system", "line_number": 0, "code_snippet": ""}],
+                scan_id=scan_id,
+                code_hash=code_hash,
+                high_severity_count=1,
+            )
+        except Exception as e:
+            return SecurityScanResult(
+                is_safe=False,
+                issues=[{"severity": "Critical", "issue_type": "ScanError",
+                        "description": str(e),
+                        "rule_id": "system", "line_number": 0, "code_snippet": ""}],
+                scan_id=scan_id,
+                code_hash=code_hash,
+                high_severity_count=1,
+            )
+    
+    def _parse_scan_output(self, output: str) -> List[Dict[str, Any]]:
+        """Parse security scan output into structured issues."""
+        issues = []
+        lines = output.split("\n")
+        
+        current_issue = None
+        for line in lines:
+            line = line.strip()
+            
+            if line.startswith(("🔴", "🟠", "🟡", "🟢", "🚨")):
+                if current_issue:
+                    issues.append(current_issue)
+                
+                severity = "Medium"
+                if "🔴" in line or "🚨" in line:
+                    severity = "Critical" if "🚨" in line else "High"
+                elif "🟠" in line:
+                    severity = "High"
+                elif "🟡" in line:
+                    severity = "Medium"
+                elif "🟢" in line:
+                    severity = "Low"
+                
+                issue_type = "Unknown"
+                if "[" in line and "]" in line:
+                    start = line.index("[") + 1
+                    end = line.index("]")
+                    issue_type = line[end+1:].strip() if end + 1 < len(line) else "Unknown"
+                
+                current_issue = {
+                    "severity": severity,
+                    "issue_type": issue_type,
+                    "rule_id": "",
+                    "line_number": 0,
+                    "description": "",
+                    "code_snippet": "",
+                }
+            elif current_issue:
+                if "Rule:" in line:
+                    current_issue["rule_id"] = line.split("Rule:")[-1].strip()
+                elif "Line" in line and ":" in line:
+                    parts = line.split(":", 1)
+                    if len(parts) > 1:
+                        try:
+                            line_part = parts[0].replace("├─", "").replace("Line", "").strip()
+                            current_issue["line_number"] = int(line_part)
+                            current_issue["description"] = parts[1].strip()
+                        except ValueError:
+                            current_issue["description"] = line
+                elif "Code:" in line:
+                    current_issue["code_snippet"] = line.split("Code:")[-1].strip()
+        
+        if current_issue:
+            issues.append(current_issue)
+        
+        return issues
+    
+    def verify_scan(self, scan_id: str, code_hash: str) -> Optional[SecurityScanResult]:
+        """Verify a scan result exists and matches the code hash."""
+        self._cleanup_expired_scans()
+        
+        if scan_id not in self._scan_cache:
+            return None
+        
+        result = self._scan_cache[scan_id]
+        if result.code_hash != code_hash:
+            return None
+        
+        return result
+    
+    def execute(
+        self,
+        language: str,
+        code: str,
+        confirmed: bool = False,
+        scan_id: Optional[str] = None,
+        sandbox_level: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """Execute code in a secure sandbox using Rust skillbox.
+        
+        Args:
+            language: Programming language (python, javascript, bash)
+            code: Code to execute
+            confirmed: Whether user has confirmed execution despite security warnings
+            scan_id: Scan ID from previous scan_code call (required when confirmed=True)
+            sandbox_level: Override sandbox level (default: from SKILLBOX_SANDBOX_LEVEL env or 3)
+        """
+        # Use default sandbox level from environment if not specified
+        if sandbox_level is None:
+            sandbox_level = self.default_sandbox_level
+        if not self.runtime_available:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": f"skillbox not found at {self.skillbox_path}. Please build it with: cd skillbox && cargo build --release",
+                "exit_code": 1
+            }
+        
+        code_hash = self._generate_code_hash(language, code)
+        
+        if sandbox_level >= 3 and not confirmed:
+            scan_result = self.scan_code(language, code)
+            if scan_result.high_severity_count > 0:
+                return {
+                    "success": False,
+                    "stdout": "",
+                    "stderr": (
+                        f"🔐 Security Review Required\n\n"
+                        f"{scan_result.format_report()}\n\n"
+                        f"To execute this code, call execute_code again with:\n"
+                        f"  - confirmed: true\n"
+                        f"  - scan_id: \"{scan_result.scan_id}\"\n"
+                    ),
+                    "exit_code": 2,
+                    "requires_confirmation": True,
+                    "scan_id": scan_result.scan_id,
+                    "security_issues": scan_result.to_dict(),
+                }
+        
+        if confirmed and scan_id:
+            cached_result = self.verify_scan(scan_id, code_hash)
+            if not cached_result:
+                return {
+                    "success": False,
+                    "stdout": "",
+                    "stderr": (
+                        "❌ Invalid or expired scan_id. The code may have been modified.\n"
+                        "Please run scan_code again to get a new scan_id."
+                    ),
+                    "exit_code": 3,
+                }
+        
+        try:
+            skill_dir, _ = self._create_temp_skill(language, code)
+            
+            try:
+                env = os.environ.copy()
+                env["SKILLBOX_AUTO_APPROVE"] = "true"
+                
+                cmd = [self.skillbox_path, "run"]
+                if sandbox_level in [1, 2, 3]:
+                    cmd.extend(["--sandbox-level", str(sandbox_level)])
+                cmd.extend([skill_dir, "{}"])
+                
+                result = subprocess.run(
+                    cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=self.timeout,
+                    env=env,
+                )
+                
+                return {
+                    "success": result.returncode == 0,
+                    "stdout": result.stdout,
+                    "stderr": result.stderr,
+                    "exit_code": result.returncode
+                }
+            finally:
+                shutil.rmtree(skill_dir, ignore_errors=True)
+                
+        except subprocess.TimeoutExpired:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": f"Execution timed out after {self.timeout} seconds",
+                "exit_code": 124
+            }
+        except Exception as e:
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": str(e),
+                "exit_code": 1
+            }
+    
+    def get_file_extension(self, language: str) -> str:
+        """Get file extension for the given language."""
+        extensions = {
+            "python": "py",
+            "javascript": "js",
+            "bash": "sh"
+        }
+        return extensions.get(language, "py")
+
+
+class MCPServer:
+    """MCP server for SkillLite sandbox execution.
+    
+    This server provides two tools for secure code execution:
+    
+    1. **scan_code**: Scan code for security issues before execution.
+       Returns a detailed report and a scan_id for confirmation.
+    
+    2. **execute_code**: Execute code in a sandbox. If high-severity
+       security issues are found, requires explicit confirmation.
+    
+    Example workflow:
+        1. Call scan_code to check for security issues
+        2. Review the security report with the user
+        3. If user approves, call execute_code with confirmed=true and scan_id
+    """
+    
+    def __init__(self):
+        if not MCP_AVAILABLE:
+            raise ImportError(
+                "MCP library not available. Please install it with: "
+                "pip install skilllite[mcp]"
+            )
+        self.server = Server("skilllite-mcp-server")
+        self.executor = SandboxExecutor()
+        self._setup_handlers()
+    
+    def _setup_handlers(self):
+        """Setup MCP server handlers."""
+        
+        @self.server.list_tools()
+        async def list_tools() -> List[Tool]:
+            return [
+                Tool(
+                    name="scan_code",
+                    description=(
+                        "Scan code for security issues before execution. "
+                        "Returns a security report with any potential risks found. "
+                        "Use this before execute_code to review security implications."
+                    ),
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "language": {
+                                "type": "string",
+                                "enum": ["python", "javascript", "bash"],
+                                "description": "Programming language of the code"
+                            },
+                            "code": {
+                                "type": "string",
+                                "description": "Code to scan for security issues"
+                            }
+                        },
+                        "required": ["language", "code"]
+                    }
+                ),
+                Tool(
+                    name="execute_code",
+                    description=(
+                        "Execute code in a secure sandbox environment. "
+                        "If security issues are found, you must set confirmed=true "
+                        "and provide the scan_id from a previous scan_code call."
+                    ),
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "language": {
+                                "type": "string",
+                                "enum": ["python", "javascript", "bash"],
+                                "description": "Programming language to execute"
+                            },
+                            "code": {
+                                "type": "string",
+                                "description": "Code to execute"
+                            },
+                            "confirmed": {
+                                "type": "boolean",
+                                "description": (
+                                    "Set to true to confirm execution despite security warnings. "
+                                    "Required when high-severity issues are found."
+                                ),
+                                "default": False
+                            },
+                            "scan_id": {
+                                "type": "string",
+                                "description": (
+                                    "The scan_id from a previous scan_code call. "
+                                    "Required when confirmed=true to verify the code hasn't changed."
+                                )
+                            },
+                            "sandbox_level": {
+                                "type": "integer",
+                                "enum": [1, 2, 3],
+                                "description": (
+                                    "Sandbox security level: "
+                                    "1=no sandbox, 2=sandbox only, 3=sandbox+security scan (default)"
+                                ),
+                                "default": 3
+                            }
+                        },
+                        "required": ["language", "code"]
+                    }
+                )
+            ]
+        
+        @self.server.call_tool()
+        async def call_tool(
+            name: str,
+            arguments: Dict[str, Any]
+        ) -> "CallToolResult":
+            if name == "scan_code":
+                return await self._handle_scan_code(arguments)
+            elif name == "execute_code":
+                return await self._handle_execute_code(arguments)
+            else:
+                return CallToolResult(
+                    isError=True,
+                    content=[
+                        TextContent(
+                            type="text",
+                            text=f"Unknown tool: {name}"
+                        )
+                    ]
+                )
+    
+    async def _handle_scan_code(self, arguments: Dict[str, Any]) -> "CallToolResult":
+        """Handle scan_code tool call."""
+        language = arguments.get("language")
+        code = arguments.get("code")
+        
+        if not language or not code:
+            return CallToolResult(
+                isError=True,
+                content=[
+                    TextContent(
+                        type="text",
+                        text="Missing required arguments: language and code"
+                    )
+                ]
+            )
+        
+        scan_result = self.executor.scan_code(language, code)
+        
+        report = scan_result.format_report()
+        
+        result_json = json.dumps(scan_result.to_dict(), indent=2)
+        
+        return CallToolResult(
+            content=[
+                TextContent(
+                    type="text",
+                    text=f"{report}\n\n---\nScan Details (JSON):\n{result_json}"
+                )
+            ]
+        )
+    
+    async def _handle_execute_code(self, arguments: Dict[str, Any]) -> "CallToolResult":
+        """Handle execute_code tool call."""
+        language = arguments.get("language")
+        code = arguments.get("code")
+        confirmed = arguments.get("confirmed", False)
+        scan_id = arguments.get("scan_id")
+        # sandbox_level: None means use default from environment variable
+        sandbox_level = arguments.get("sandbox_level")
+        
+        if not language or not code:
+            return CallToolResult(
+                isError=True,
+                content=[
+                    TextContent(
+                        type="text",
+                        text="Missing required arguments: language and code"
+                    )
+                ]
+            )
+        
+        result = self.executor.execute(
+            language=language,
+            code=code,
+            confirmed=confirmed,
+            scan_id=scan_id,
+            sandbox_level=sandbox_level,
+        )
+        
+        if result.get("requires_confirmation"):
+            return CallToolResult(
+                content=[
+                    TextContent(
+                        type="text",
+                        text=result["stderr"]
+                    )
+                ]
+            )
+        
+        output_lines = []
+        if result["stdout"]:
+            output_lines.append(f"Output:\n{result['stdout']}")
+        if result["stderr"]:
+            output_lines.append(f"Errors:\n{result['stderr']}")
+        
+        output_text = "\n".join(output_lines) if output_lines else "Execution completed successfully (no output)"
+        
+        if result["success"]:
+            return CallToolResult(
+                content=[
+                    TextContent(
+                        type="text",
+                        text=output_text
+                    )
+                ]
+            )
+        else:
+            return CallToolResult(
+                isError=True,
+                content=[
+                    TextContent(
+                        type="text",
+                        text=f"Execution failed:\n{output_text}"
+                    )
+                ]
+            )
+    
+    async def run(self):
+        """Run the MCP server."""
+        async with stdio_server() as (read_stream, write_stream):
+            await self.server.run(read_stream, write_stream, self.server.create_initialization_options())
+
+
+async def main():
+    """Main entry point for the MCP sandbox server."""
+    if not MCP_AVAILABLE:
+        print("Error: MCP library not available")
+        print("Please install it with: pip install skilllite[mcp]")
+        return
+    
+    server = MCPServer()
+    await server.run()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())