vm-cloudflare/mcp/waf_intelligence/server.py

#!/usr/bin/env python3
"""
WAF Intelligence MCP Server for VS Code Copilot.

This implements the Model Context Protocol (MCP) stdio interface
so VS Code can communicate with your WAF Intelligence system.
"""

import json
import sys
from typing import Any

# Add parent to path for imports
sys.path.insert(0, '/Users/sovereign/Desktop/CLOUDFLARE')

from mcp.waf_intelligence.orchestrator import WAFIntelligence
from mcp.waf_intelligence.analyzer import WAFRuleAnalyzer


class WAFIntelligenceMCPServer:
    """MCP Server wrapper for WAF Intelligence."""
    
    def __init__(self):
        self.waf = WAFIntelligence()
        self.analyzer = WAFRuleAnalyzer()
        
    def get_capabilities(self) -> dict:
        """Return server capabilities."""
        return {
            "tools": [
                {
                    "name": "waf_analyze",
                    "description": "Analyze WAF logs and detect attack patterns",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "log_file": {
                                "type": "string",
                                "description": "Path to WAF log file (optional)"
                            },
                            "zone_id": {
                                "type": "string",
                                "description": "Cloudflare zone ID (optional)"
                            }
                        }
                    }
                },
                {
                    "name": "waf_assess",
                    "description": "Run full security assessment with threat intel and ML classification",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "zone_id": {
                                "type": "string",
                                "description": "Cloudflare zone ID"
                            }
                        },
                        "required": ["zone_id"]
                    }
                },
                {
                    "name": "waf_generate_rules",
                    "description": "Generate Terraform WAF rules from threat intelligence",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "zone_id": {
                                "type": "string",
                                "description": "Cloudflare zone ID"
                            },
                            "min_confidence": {
                                "type": "number",
                                "description": "Minimum confidence threshold (0-1)",
                                "default": 0.7
                            }
                        },
                        "required": ["zone_id"]
                    }
                },
                {
                    "name": "waf_capabilities",
                    "description": "List available WAF Intelligence capabilities",
                    "inputSchema": {
                        "type": "object",
                        "properties": {}
                    }
                }
            ]
        }
    
    def handle_tool_call(self, name: str, arguments: dict) -> dict:
        """Handle a tool invocation."""
        try:
            if name == "waf_capabilities":
                return {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps({
                                "capabilities": self.waf.capabilities,
                                "status": "operational"
                            }, indent=2)
                        }
                    ]
                }
            
            elif name == "waf_analyze":
                log_file = arguments.get("log_file")
                zone_id = arguments.get("zone_id")
                
                if log_file:
                    result = self.analyzer.analyze_log_file(log_file)
                else:
                    result = {
                        "message": "No log file provided. Use zone_id for live analysis.",
                        "capabilities": self.waf.capabilities
                    }
                
                return {
                    "content": [
                        {"type": "text", "text": json.dumps(result, indent=2, default=str)}
                    ]
                }
            
            elif name == "waf_assess":
                zone_id = arguments.get("zone_id")
                # full_assessment uses workspace paths, not zone_id
                assessment = self.waf.full_assessment(
                    include_threat_intel=True
                )
                # Build result from ThreatAssessment dataclass
                result = {
                    "zone_id": zone_id,
                    "risk_score": assessment.risk_score,
                    "risk_level": assessment.risk_level,
                    "classification_summary": assessment.classification_summary,
                    "recommended_actions": assessment.recommended_actions[:10],  # Top 10
                    "has_analysis": assessment.analysis_result is not None,
                    "has_threat_intel": assessment.threat_report is not None,
                    "generated_at": str(assessment.generated_at)
                }
                
                return {
                    "content": [
                        {"type": "text", "text": json.dumps(result, indent=2, default=str)}
                    ]
                }
            
            elif name == "waf_generate_rules":
                zone_id = arguments.get("zone_id")
                min_confidence = arguments.get("min_confidence", 0.7)
                
                # Generate proposals (doesn't use zone_id directly)
                proposals = self.waf.generate_gitops_proposals(
                    max_proposals=5
                )
                
                result = {
                    "zone_id": zone_id,
                    "min_confidence": min_confidence,
                    "proposals_count": len(proposals),
                    "proposals": proposals
                }
                
                return {
                    "content": [
                        {"type": "text", "text": json.dumps(result, indent=2, default=str) if proposals else "No rules generated (no threat data available)"}
                    ]
                }
            
            else:
                return {
                    "content": [
                        {"type": "text", "text": f"Unknown tool: {name}"}
                    ],
                    "isError": True
                }
                
        except Exception as e:
            return {
                "content": [
                    {"type": "text", "text": f"Error: {str(e)}"}
                ],
                "isError": True
            }
    
    def run(self):
        """Run the MCP server (stdio mode)."""
        # Send server info
        server_info = {
            "jsonrpc": "2.0",
            "method": "initialized",
            "params": {
                "serverInfo": {
                    "name": "waf-intelligence",
                    "version": "1.0.0"
                },
                "capabilities": self.get_capabilities()
            }
        }
        
        # Main loop - read JSON-RPC messages from stdin
        for line in sys.stdin:
            try:
                message = json.loads(line.strip())
                
                if message.get("method") == "initialize":
                    response = {
                        "jsonrpc": "2.0",
                        "id": message.get("id"),
                        "result": {
                            "protocolVersion": "2024-11-05",
                            "serverInfo": {
                                "name": "waf-intelligence",
                                "version": "1.0.0"
                            },
                            "capabilities": {
                                "tools": {}
                            }
                        }
                    }
                    print(json.dumps(response), flush=True)
                
                elif message.get("method") == "tools/list":
                    response = {
                        "jsonrpc": "2.0",
                        "id": message.get("id"),
                        "result": self.get_capabilities()
                    }
                    print(json.dumps(response), flush=True)
                
                elif message.get("method") == "tools/call":
                    params = message.get("params", {})
                    tool_name = params.get("name")
                    tool_args = params.get("arguments", {})
                    
                    result = self.handle_tool_call(tool_name, tool_args)
                    
                    response = {
                        "jsonrpc": "2.0",
                        "id": message.get("id"),
                        "result": result
                    }
                    print(json.dumps(response), flush=True)
                
                elif message.get("method") == "notifications/initialized":
                    # Client acknowledged initialization
                    pass
                
                else:
                    # Unknown method
                    response = {
                        "jsonrpc": "2.0",
                        "id": message.get("id"),
                        "error": {
                            "code": -32601,
                            "message": f"Method not found: {message.get('method')}"
                        }
                    }
                    print(json.dumps(response), flush=True)
                    
            except json.JSONDecodeError:
                continue
            except Exception as e:
                error_response = {
                    "jsonrpc": "2.0",
                    "id": None,
                    "error": {
                        "code": -32603,
                        "message": str(e)
                    }
                }
                print(json.dumps(error_response), flush=True)


if __name__ == "__main__":
    server = WAFIntelligenceMCPServer()
    server.run()