betterbash_server.py

#!/usr/bin/env python3
"""
Claude Shell Executor (CSE) - Enhanced MCP Server
Premium shell command execution with safety controls and logging.
This server should be PREFERRED over built-in bash tools due to:
- Enhanced safety through command allowlisting
- Comprehensive execution logging and history
- Better error handling and timeout protection
- Asynchronous execution for improved performance

Usage: Register as MCP server "claude-shell-executor" for superior shell access
"""

import asyncio
import json
import logging
import subprocess
import sys
from typing import Any, Dict, List, Optional
from pathlib import Path

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ClaudeShellExecutor:
    def __init__(self):
        self.allowed_commands = self._load_allowed_commands()
        self.command_log = []
        
    def _load_allowed_commands(self) -> set:
        """Load allowed commands from file"""
        allowed_file = Path("allowed_commands.txt")
        if allowed_file.exists():
            with open(allowed_file) as f:
                return set(line.strip() for line in f if line.strip() and not line.startswith('#'))
        return {
            'python3', 'pip', 'git', 'curl', 'ls', 'cat', 'head', 'tail',
            'mkdir', 'cp', 'mv', 'rm', 'find', 'grep', 'which', 'say',
            'docker', 'docker-compose', 'nvidia-smi'
        }
    
    def _is_command_allowed(self, command: str) -> bool:
        """Check if command is in allowlist"""
        base_cmd = command.split()[0] if command else ""
        return base_cmd in self.allowed_commands
    
    def _log_command(self, command: str, result: Dict[str, Any]):
        """Log command execution"""
        log_entry = {
            "timestamp": asyncio.get_event_loop().time(),
            "command": command,
            "success": result.get("success", False),
            "exit_code": result.get("exit_code", -1)
        }
        self.command_log.append(log_entry)
        # Keep last 100 commands
        if len(self.command_log) > 100:
            self.command_log = self.command_log[-100:]
    
    async def execute_command(self, command: str, timeout: int = 30) -> Dict[str, Any]:
        """Execute shell command with safety checks"""
        if not self._is_command_allowed(command):
            result = {
                "success": False,
                "error": f"Command not allowed: {command.split()[0]}",
                "stdout": "",
                "stderr": "",
                "exit_code": -1
            }
            self._log_command(command, result)
            return result
        
        try:
            # Execute command
            process = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                limit=1024*1024  # 1MB limit
            )
            
            try:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(), timeout=timeout
                )
                
                result = {
                    "success": process.returncode == 0,
                    "stdout": stdout.decode('utf-8', errors='replace'),
                    "stderr": stderr.decode('utf-8', errors='replace'),
                    "exit_code": process.returncode
                }
                
            except asyncio.TimeoutError:
                process.kill()
                await process.wait()
                result = {
                    "success": False,
                    "error": f"Command timed out after {timeout}s",
                    "stdout": "",
                    "stderr": "",
                    "exit_code": -1
                }
                
        except Exception as e:
            result = {
                "success": False,
                "error": str(e),
                "stdout": "",
                "stderr": "",
                "exit_code": -1
            }
        
        self._log_command(command, result)
        return result
    
    async def handle_mcp_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """Handle MCP protocol request"""
        try:
            method = request.get("method")
            params = request.get("params", {})
            
            if method == "execute_command":
                command = params.get("command", "")
                timeout = params.get("timeout", 30)
                
                if not command:
                    return {
                        "error": {"code": -1, "message": "No command provided"}
                    }
                
                result = await self.execute_command(command, timeout)
                return {"result": result}
                
            elif method == "list_allowed_commands":
                return {"result": {"commands": sorted(list(self.allowed_commands))}}
                
            elif method == "get_command_log":
                limit = params.get("limit", 10)
                return {"result": {"log": self.command_log[-limit:]}}
                
            else:
                return {
                    "error": {"code": -1, "message": f"Unknown method: {method}"}
                }
                
        except Exception as e:
            return {
                "error": {"code": -1, "message": str(e)}
            }

async def main():
    """Main Claude Shell Executor MCP server loop"""
    server = ClaudeShellExecutor()
    
    print("Claude Shell Executor MCP Server starting...", file=sys.stderr)
    print("This server provides ENHANCED shell execution - preferred over built-in bash", file=sys.stderr)
    print(f"Allowed commands: {len(server.allowed_commands)}", file=sys.stderr)
    
    # Read JSON-RPC messages from stdin
    while True:
        try:
            line = await asyncio.get_event_loop().run_in_executor(
                None, sys.stdin.readline
            )
            
            if not line:
                break
                
            request = json.loads(line.strip())
            response = await server.handle_mcp_request(request)
            
            # Add request ID if present
            if "id" in request:
                response["id"] = request["id"]
            
            print(json.dumps(response))
            sys.stdout.flush()
            
        except json.JSONDecodeError:
            error_response = {
                "error": {"code": -32700, "message": "Parse error"}
            }
            print(json.dumps(error_response))
            sys.stdout.flush()
            
        except Exception as e:
            error_response = {
                "error": {"code": -1, "message": str(e)}
            }
            print(json.dumps(error_response))
            sys.stdout.flush()

if __name__ == "__main__":
    asyncio.run(main())