#!/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 from layer0 import layer0_entry from layer0.shadow_classifier import ShadowEvalResult 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", {}) # Layer 0: pre-boot Shadow Eval gate before handling tool calls. routing_action, shadow = layer0_entry(_shadow_query_repr(tool_name, tool_args)) if routing_action != "HANDOFF_TO_LAYER1": response = _layer0_mcp_response(routing_action, shadow, message.get("id")) print(json.dumps(response), flush=True) continue 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() def _shadow_query_repr(tool_name: str, tool_args: dict) -> str: """Build a textual representation of the tool call for Layer 0 classification.""" try: return f"{tool_name}: {json.dumps(tool_args, sort_keys=True)}" except TypeError: return f"{tool_name}: {str(tool_args)}" def _layer0_mcp_response(routing_action: str, shadow: ShadowEvalResult, msg_id: Any) -> dict: """ Map Layer 0 outcomes to MCP responses. Catastrophic/forbidden/ambiguous short-circuit with minimal disclosure. """ base = {"jsonrpc": "2.0", "id": msg_id} if routing_action == "FAIL_CLOSED": base["error"] = {"code": -32000, "message": "Layer 0: cannot comply with this request."} return base if routing_action == "HANDOFF_TO_GUARDRAILS": reason = shadow.reason or "governance_violation" base["error"] = { "code": -32001, "message": f"Layer 0: governance violation detected ({reason}).", } return base if routing_action == "PROMPT_FOR_CLARIFICATION": base["error"] = { "code": -32002, "message": "Layer 0: request is ambiguous. Please clarify and retry.", } return base base["error"] = {"code": -32099, "message": "Layer 0: unrecognized routing action; refusing."} return base