327 lines
12 KiB
Python
Executable File
327 lines
12 KiB
Python
Executable File
#!/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
|