th-analyzer/commando_mcp_server.py

#!/usr/bin/env python3
"""Commando VM — MCP Tool Server
Exposes all 112+ Commando VM tools as MCP tools via HTTP/SSE.
Run: python3 commando_mcp_server.py
"""

import os
import sys
import json
import subprocess
import tempfile
from pathlib import Path

# ─── MCP Protocol Implementation (minimal) ────────────────

TOOLS_DIR = Path("C:/tools")
COMMANDS = {
    "mimikatz": str(TOOLS_DIR / "Mimikatz" / "mimikatz.exe"),
    "hashcat": str(TOOLS_DIR / "hashcat" / "hashcat.exe"),
    "bloodhound": str(TOOLS_DIR / "BloodHound"),
    "crackmapexec": str(TOOLS_DIR / "CrackMapExecWin"),
    "impacket": str(TOOLS_DIR / "impacket"),
    "sysinternals": str(TOOLS_DIR / "Sysinternals"),
}

def discover_tools():
    """Discover all executable tools in C:\\tools."""
    tools = []
    for dirpath, dirnames, filenames in os.walk(str(TOOLS_DIR)):
        for f in filenames:
            if f.endswith(('.exe', '.ps1', '.py', '.bat')):
                rel = os.path.relpath(os.path.join(dirpath, f), str(TOOLS_DIR))
                tools.append(rel.split('\\')[0])  # category
    return sorted(set(tools))

def run_tool(tool_name, args, timeout=60):
    """Run a tool from Commando VM tools directory and return output."""
    # Find the tool
    for dirpath, dirnames, filenames in os.walk(str(TOOLS_DIR)):
        for f in filenames:
            if f.lower().startswith(tool_name.lower()) or tool_name.lower() in f.lower():
                exe = os.path.join(dirpath, f)
                try:
                    cmd = [exe] + args.split() if args else [exe]
                    result = subprocess.run(
                        cmd, capture_output=True, text=True, timeout=timeout
                    )
                    return {
                        "tool": f,
                        "stdout": result.stdout[:5000],
                        "stderr": result.stderr[:1000],
                        "exit_code": result.returncode
                    }
                except subprocess.TimeoutExpired:
                    return {"tool": f, "error": f"Timed out after {timeout}s"}
                except Exception as e:
                    return {"tool": f, "error": str(e)}
    return {"error": f"Tool '{tool_name}' not found in Commando VM"}

def run_powershell(script, timeout=60):
    """Run a PowerShell command and return output."""
    try:
        result = subprocess.run(
            ["powershell", "-Command", script],
            capture_output=True, text=True, timeout=timeout
        )
        return {
            "stdout": result.stdout[:5000],
            "stderr": result.stderr[:1000],
            "exit_code": result.returncode
        }
    except subprocess.TimeoutExpired:
        return {"error": f"Timed out after {timeout}s"}
    except Exception as e:
        return {"error": str(e)}

# ─── MCP Tool Definitions ────────────────────────────────

def handle_request(request):
    """Handle MCP protocol request."""
    req = json.loads(request) if isinstance(request, str) else request
    method = req.get("method", "")
    req_id = req.get("id", 1)
    
    if method == "initialize":
        return {"jsonrpc": "2.0", "id": req_id, "result": {
            "protocolVersion": "0.1.0",
            "capabilities": {"tools": {}},
            "serverInfo": {"name": "commando-mcp", "version": "1.0.0"}
        }}
    
    elif method == "tools/list":
        tools = discover_tools()
        categories = sorted(set(tools))
        
        tool_list = [
            {
                "name": f"commando_{cat.lower().replace('-','_').replace(' ','_')}",
                "description": f"Run {cat} tools from Commando VM",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "tool": {
                            "type": "string",
                            "description": f"Tool name from {cat} category"
                        },
                        "args": {
                            "type": "string",
                            "description": "Command-line arguments"
                        },
                        "timeout": {
                            "type": "number",
                            "description": "Timeout in seconds (default: 60)"
                        }
                    },
                    "required": ["tool"]
                }
            }
            for cat in categories[:50]  # Max 50 tools to avoid overwhelm
        ]
        
        # Add generic tool runner
        tool_list.append({
            "name": "commando_run",
            "description": "Run any tool from Commando VM's C:\\tools directory",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "tool": {"type": "string", "description": "Tool name (e.g., mimikatz, hashcat, nmap)"},
                    "args": {"type": "string", "description": "Command-line arguments"},
                    "timeout": {"type": "number", "description": "Timeout in seconds (default: 60)"}
                },
                "required": ["tool"]
            }
        })
        
        tool_list.append({
            "name": "commando_powershell",
            "description": "Run a PowerShell command on Commando VM",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "PowerShell command to execute"},
                    "timeout": {"type": "number", "description": "Timeout in seconds (default: 60)"}
                },
                "required": ["command"]
            }
        })
        
        tool_list.append({
            "name": "commando_list_tools",
            "description": "List all available tools in Commando VM",
            "inputSchema": {
                "type": "object",
                "properties": {}
            }
        })
        
        return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": tool_list}}
    
    elif method == "tools/call":
        params = req.get("params", {})
        name = params.get("name", "")
        args = params.get("arguments", {})
        
        if name == "commando_list_tools":
            tools = discover_tools()
            categories = sorted(set(tools))
            return {"jsonrpc": "2.0", "id": req_id, "result": {
                "content": [{
                    "type": "text",
                    "text": f"Commando VM has {len(tools)} tool categories:\n" + "\n".join(f"  - {t}" for t in categories)
                }]
            }}
        
        elif name == "commando_run":
            result = run_tool(args.get("tool", ""), args.get("args", ""), args.get("timeout", 60))
            output = ""
            if "stdout" in result:
                output = result["stdout"]
            elif "error" in result:
                output = f"Error: {result['error']}"
            return {"jsonrpc": "2.0", "id": req_id, "result": {
                "content": [{"type": "text", "text": output}]
            }}
        
        elif name == "commando_powershell":
            result = run_powershell(args.get("command", ""), args.get("timeout", 60))
            output = result.get("stdout", result.get("error", "No output"))
            return {"jsonrpc": "2.0", "id": req_id, "result": {
                "content": [{"type": "text", "text": output}]
            }}
        
        else:
            # Try running by tool name from description
            cat_name = name.replace("commando_", "").replace("_", "-")
            result = run_tool(cat_name, args.get("args", ""), args.get("timeout", 60))
            return {"jsonrpc": "2.0", "id": req_id, "result": {
                "content": [{"type": "text", "text": result.get("stdout", result.get("error", "No output"))}]
            }}
    
    return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": "Method not found"}}

# ─── STDIO Transport ─────────────────────────────────────

def main():
    """Run MCP server over stdio (for Hermes integration)."""
    import sys
    for line in sys.stdin:
        if not line.strip():
            continue
        try:
            response = handle_request(line.strip())
            print(json.dumps(response), flush=True)
        except Exception as e:
            print(json.dumps({"jsonrpc": "2.0", "error": {"code": -32700, "message": str(e)}}), flush=True)

if __name__ == "__main__":
    import sys
    
    if "--http" in sys.argv:
        # HTTP/SSE mode — for Hermes MCP integration
        from http.server import HTTPServer, BaseHTTPRequestHandler
        import urllib.parse
        import queue
        
        PORT = int(sys.argv[sys.argv.index("--http") + 1]) if "--http" in sys.argv and len(sys.argv) > sys.argv.index("--http") + 1 else 8092
        
        # Store SSE response queues per session
        sse_clients = {}
        next_session = 1
        
        class MCPHandler(BaseHTTPRequestHandler):
            def do_GET(self):
                parsed = urllib.parse.urlparse(self.path)
                
                if parsed.path == "/health":
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Access-Control-Allow-Origin", "*")
                    self.end_headers()
                    self.wfile.write(json.dumps({"status": "ok", "server": "commando-mcp", "tools": len(discover_tools())}).encode())
                    return
                
                if parsed.path == "/mcp" or parsed.path == "/sse":
                    # SSE connection
                    accept = self.headers.get("Accept", "")
                    if "text/event-stream" not in accept and parsed.path == "/mcp":
                        # Return initialize for non-SSE clients
                        self.send_response(200)
                        self.send_header("Content-Type", "application/json")
                        self.send_header("Access-Control-Allow-Origin", "*")
                        self.send_header("MCP-Version", "0.1.0")
                        self.end_headers()
                        resp = handle_request({"jsonrpc": "2.0", "id": "init", "method": "initialize"})
                        self.wfile.write(json.dumps(resp).encode())
                        return
                    
                    # SSE mode
                    global next_session
                    session_id = f"sess_{next_session}"
                    next_session += 1
                    msg_queue = queue.Queue()
                    sse_clients[session_id] = msg_queue
                    
                    endpoint = f"/mcp?session={session_id}"
                    
                    self.send_response(200)
                    self.send_header("Content-Type", "text/event-stream")
                    self.send_header("Cache-Control", "no-cache")
                    self.send_header("Connection", "keep-alive")
                    self.send_header("Access-Control-Allow-Origin", "*")
                    self.end_headers()
                    
                    # Send session event
                    self.wfile.write(f"event: endpoint\ndata: {endpoint}\n\n".encode())
                    self.wfile.write(f"event: session_id\ndata: {session_id}\n\n".encode())
                    self.wfile.flush()
                    
                    # Keep connection open and send responses
                    while True:
                        try:
                            msg = msg_queue.get(timeout=30)
                            self.wfile.write(f"data: {json.dumps(msg)}\n\n".encode())
                            self.wfile.flush()
                        except queue.Empty:
                            self.wfile.write(": keepalive\n\n".encode())
                            self.wfile.flush()
                
                else:
                    self.send_response(404)
                    self.end_headers()
            
            def do_POST(self):
                global next_session
                content_length = int(self.headers.get("Content-Length", 0))
                body = self.rfile.read(content_length) if content_length > 0 else b"{}"
                parsed = urllib.parse.urlparse(self.path)
                
                # Parse session from query string
                params = urllib.parse.parse_qs(parsed.query)
                session_id = params.get("session", [None])[0]
                
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                
                try:
                    req = json.loads(body.decode())
                    resp = handle_request(req)
                    
                    if session_id and session_id in sse_clients:
                        # Send response via SSE
                        sse_clients[session_id].put(resp)
                        self.wfile.write(json.dumps({"jsonrpc": "2.0", "id": req.get("id"), "result": {"_delivered": "via_sse"}}).encode())
                    else:
                        # Direct response
                        self.wfile.write(json.dumps(resp).encode())
                except Exception as e:
                    self.wfile.write(json.dumps({"jsonrpc": "2.0", "error": {"code": -32700, "message": str(e)}}).encode())
            
            def log_message(self, format, *args):
                pass  # Suppress logs
        
        print(f"Commando MCP server running on port {PORT}")
        print(f"  Endpoint: http://0.0.0.0:{PORT}/mcp")
        print(f"  Health:   http://0.0.0.0:{PORT}/health")
        server = HTTPServer(("0.0.0.0", PORT), MCPHandler)
        server.serve_forever()
    else:
        # STDIO mode
        main()