chore: pre-migration snapshot

Layer0, MCP servers, Terraform consolidation

mcp/waf_intelligence/mcp_server.py

@@ -0,0 +1,632 @@
+from __future__ import annotations
+
+import glob
+import json
+import os
+import sys
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+from cloudflare.layer0 import layer0_entry
+from cloudflare.layer0.shadow_classifier import ShadowEvalResult
+
+from .orchestrator import ThreatAssessment, WAFInsight, WAFIntelligence
+
+MAX_BYTES_DEFAULT = 32_000
+
+
+def _cloudflare_root() -> Path:
+    # mcp_server.py -> waf_intelligence -> mcp -> cloudflare
+    return Path(__file__).resolve().parents[2]
+
+
+def _max_bytes() -> int:
+    raw = (os.getenv("VM_MCP_MAX_BYTES") or "").strip()
+    if not raw:
+        return MAX_BYTES_DEFAULT
+    try:
+        return max(4_096, int(raw))
+    except ValueError:
+        return MAX_BYTES_DEFAULT
+
+
+def _redact(obj: Any) -> Any:
+    sensitive_keys = ("token", "secret", "password", "private", "key", "certificate")
+
+    if isinstance(obj, dict):
+        out: Dict[str, Any] = {}
+        for k, v in obj.items():
+            if any(s in str(k).lower() for s in sensitive_keys):
+                out[k] = "<REDACTED>"
+            else:
+                out[k] = _redact(v)
+        return out
+    if isinstance(obj, list):
+        return [_redact(v) for v in obj]
+    if isinstance(obj, str):
+        if obj.startswith("ghp_") or obj.startswith("github_pat_"):
+            return "<REDACTED>"
+        return obj
+    return obj
+
+
+def _safe_json(payload: Dict[str, Any]) -> str:
+    payload = _redact(payload)
+    raw = json.dumps(payload, ensure_ascii=False, separators=(",", ":"), default=str)
+    if len(raw.encode("utf-8")) <= _max_bytes():
+        return json.dumps(payload, ensure_ascii=False, indent=2, default=str)
+
+    truncated = {
+        "ok": payload.get("ok", True),
+        "truncated": True,
+        "summary": payload.get("summary", "Response exceeded max size; truncated."),
+        "next_steps": payload.get(
+            "next_steps",
+            [
+                "request fewer files/insights (limit=...)",
+                "use higher min_severity to reduce output",
+            ],
+        ),
+    }
+    return json.dumps(truncated, ensure_ascii=False, indent=2, default=str)
+
+
+def _mcp_text_result(
+    payload: Dict[str, Any], *, is_error: bool = False
+) -> Dict[str, Any]:
+    result: Dict[str, Any] = {
+        "content": [{"type": "text", "text": _safe_json(payload)}]
+    }
+    if is_error:
+        result["isError"] = True
+    return result
+
+
+def _insight_to_dict(insight: WAFInsight) -> Dict[str, Any]:
+    return asdict(insight)
+
+
+def _assessment_to_dict(assessment: ThreatAssessment) -> Dict[str, Any]:
+    violations = []
+    if assessment.analysis_result and getattr(
+        assessment.analysis_result, "violations", None
+    ):
+        violations = list(assessment.analysis_result.violations)
+
+    severity_counts = {"error": 0, "warning": 0, "info": 0}
+    for v in violations:
+        sev = getattr(v, "severity", "info")
+        if sev in severity_counts:
+            severity_counts[sev] += 1
+
+    return {
+        "risk_score": assessment.risk_score,
+        "risk_level": assessment.risk_level,
+        "classification_summary": assessment.classification_summary,
+        "recommended_actions": assessment.recommended_actions,
+        "analysis": {
+            "has_config_analysis": assessment.analysis_result is not None,
+            "violations_total": len(violations),
+            "violations_by_severity": severity_counts,
+        },
+        "has_threat_intel": assessment.threat_report is not None,
+        "generated_at": str(assessment.generated_at),
+    }
+
+
+TOOLS: List[Dict[str, Any]] = [
+    {
+        "name": "waf_capabilities",
+        "description": "List available WAF Intelligence capabilities.",
+        "inputSchema": {"type": "object", "properties": {}},
+    },
+    {
+        "name": "analyze_waf",
+        "description": "Analyze Terraform WAF file(s) and return curated insights (legacy alias for waf_analyze).",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "file": {
+                    "type": "string",
+                    "description": "Single file path to analyze.",
+                },
+                "files": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                    "description": "List of file paths or glob patterns to analyze.",
+                },
+                "limit": {
+                    "type": "integer",
+                    "default": 3,
+                    "description": "Max insights per file.",
+                },
+                "severity_threshold": {
+                    "type": "string",
+                    "enum": ["info", "warning", "error"],
+                    "default": "warning",
+                    "description": "Minimum severity to include (alias for min_severity).",
+                },
+            },
+        },
+    },
+    {
+        "name": "waf_analyze",
+        "description": "Analyze Terraform WAF file(s) and return curated insights (requires file or files).",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "file": {
+                    "type": "string",
+                    "description": "Single file path to analyze.",
+                },
+                "files": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                    "description": "List of file paths or glob patterns to analyze.",
+                },
+                "limit": {
+                    "type": "integer",
+                    "default": 3,
+                    "description": "Max insights per file.",
+                },
+                "min_severity": {
+                    "type": "string",
+                    "enum": ["info", "warning", "error"],
+                    "default": "warning",
+                    "description": "Minimum severity to include.",
+                },
+            },
+        },
+    },
+    {
+        "name": "waf_assess",
+        "description": "Run a broader assessment (optionally includes threat intel collection).",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "waf_config_path": {
+                    "type": "string",
+                    "description": "Path to Terraform WAF config (default: terraform/waf.tf).",
+                },
+                "include_threat_intel": {
+                    "type": "boolean",
+                    "default": False,
+                    "description": "If true, attempt to collect threat intel (may require network and credentials).",
+                },
+            },
+        },
+    },
+    {
+        "name": "waf_generate_gitops_proposals",
+        "description": "Generate GitOps-ready rule proposals (best-effort; requires threat intel to produce output).",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "waf_config_path": {
+                    "type": "string",
+                    "description": "Path to Terraform WAF config (default: terraform/waf.tf).",
+                },
+                "include_threat_intel": {
+                    "type": "boolean",
+                    "default": True,
+                    "description": "Attempt to collect threat intel before proposing rules.",
+                },
+                "max_proposals": {
+                    "type": "integer",
+                    "default": 5,
+                    "description": "Maximum proposals to generate.",
+                },
+            },
+        },
+    },
+]
+
+
+class WafIntelligenceTools:
+    def __init__(self) -> None:
+        self.workspace_root = _cloudflare_root()
+        self.repo_root = self.workspace_root.parent
+        self.waf = WAFIntelligence(workspace_path=str(self.workspace_root))
+
+    def _resolve_path(self, raw: str) -> Path:
+        path = Path(raw)
+        if path.is_absolute():
+            return path
+
+        candidates = [
+            Path.cwd() / path,
+            self.workspace_root / path,
+            self.repo_root / path,
+        ]
+        for candidate in candidates:
+            if candidate.exists():
+                return candidate
+        return self.workspace_root / path
+
+    def waf_capabilities(self) -> Dict[str, Any]:
+        return {
+            "ok": True,
+            "summary": "WAF Intelligence capabilities.",
+            "data": {"capabilities": self.waf.capabilities},
+            "truncated": False,
+            "next_steps": [
+                "Call waf_analyze(file=..., limit=...) to analyze config.",
+                "Call waf_assess(include_threat_intel=true) for a broader assessment.",
+            ],
+        }
+
+    def waf_analyze(
+        self,
+        *,
+        file: Optional[str] = None,
+        files: Optional[List[str]] = None,
+        limit: int = 3,
+        min_severity: str = "warning",
+    ) -> Dict[str, Any]:
+        paths: List[str] = []
+        if files:
+            for pattern in files:
+                paths.extend(glob.glob(pattern))
+        if file:
+            paths.append(file)
+
+        seen = set()
+        unique_paths: List[str] = []
+        for p in paths:
+            if p not in seen:
+                seen.add(p)
+                unique_paths.append(p)
+
+        if not unique_paths:
+            return {
+                "ok": False,
+                "summary": "Provide 'file' or 'files' to analyze.",
+                "truncated": False,
+                "next_steps": ["Call waf_analyze(file='terraform/waf.tf')"],
+            }
+
+        results: List[Dict[str, Any]] = []
+        for p in unique_paths:
+            path = self._resolve_path(p)
+            if not path.exists():
+                results.append(
+                    {
+                        "file": str(path),
+                        "ok": False,
+                        "summary": "File not found.",
+                    }
+                )
+                continue
+
+            insights = self.waf.analyze_and_recommend(
+                str(path),
+                limit=limit,
+                min_severity=min_severity,
+            )
+            results.append(
+                {
+                    "file": str(path),
+                    "ok": True,
+                    "insights": [_insight_to_dict(i) for i in insights],
+                }
+            )
+
+        ok = all(r.get("ok") for r in results)
+        return {
+            "ok": ok,
+            "summary": f"Analyzed {len(results)} file(s).",
+            "data": {"results": results},
+            "truncated": False,
+            "next_steps": [
+                "Raise/lower min_severity or limit to tune output size.",
+            ],
+        }
+
+    def waf_assess(
+        self,
+        *,
+        waf_config_path: Optional[str] = None,
+        include_threat_intel: bool = False,
+    ) -> Dict[str, Any]:
+        waf_config_path_resolved = (
+            str(self._resolve_path(waf_config_path)) if waf_config_path else None
+        )
+        assessment = self.waf.full_assessment(
+            waf_config_path=waf_config_path_resolved,
+            include_threat_intel=include_threat_intel,
+        )
+        return {
+            "ok": True,
+            "summary": "WAF assessment complete.",
+            "data": _assessment_to_dict(assessment),
+            "truncated": False,
+            "next_steps": [
+                "Call waf_generate_gitops_proposals(...) to draft Terraform rule proposals (best-effort).",
+            ],
+        }
+
+    def waf_generate_gitops_proposals(
+        self,
+        *,
+        waf_config_path: Optional[str] = None,
+        include_threat_intel: bool = True,
+        max_proposals: int = 5,
+    ) -> Dict[str, Any]:
+        waf_config_path_resolved = (
+            str(self._resolve_path(waf_config_path)) if waf_config_path else None
+        )
+        assessment = self.waf.full_assessment(
+            waf_config_path=waf_config_path_resolved,
+            include_threat_intel=include_threat_intel,
+        )
+        proposals = self.waf.generate_gitops_proposals(
+            threat_report=assessment.threat_report,
+            max_proposals=max_proposals,
+        )
+        return {
+            "ok": True,
+            "summary": f"Generated {len(proposals)} proposal(s).",
+            "data": {
+                "assessment": _assessment_to_dict(assessment),
+                "proposals": proposals,
+            },
+            "truncated": False,
+            "next_steps": [
+                "If proposals are empty, enable threat intel and ensure required credentials/log sources exist.",
+            ],
+        }
+
+
+class StdioJsonRpc:
+    def __init__(self) -> None:
+        self._in = sys.stdin.buffer
+        self._out = sys.stdout.buffer
+        self._mode: str | None = None  # "headers" | "line"
+
+    def read_message(self) -> Optional[Dict[str, Any]]:
+        while True:
+            if self._mode == "line":
+                line = self._in.readline()
+                if not line:
+                    return None
+                raw = line.decode("utf-8", "replace").strip()
+                if not raw:
+                    continue
+                try:
+                    msg = json.loads(raw)
+                except Exception:
+                    continue
+                if isinstance(msg, dict):
+                    return msg
+                continue
+
+            first = self._in.readline()
+            if not first:
+                return None
+
+            if first in (b"\r\n", b"\n"):
+                continue
+
+            # Auto-detect newline-delimited JSON framing.
+            if self._mode is None and first.lstrip().startswith(b"{"):
+                try:
+                    msg = json.loads(first.decode("utf-8", "replace"))
+                except Exception:
+                    msg = None
+                if isinstance(msg, dict):
+                    self._mode = "line"
+                    return msg
+
+            headers: Dict[str, str] = {}
+            try:
+                text = first.decode("utf-8", "replace").strip()
+            except Exception:
+                continue
+            if ":" not in text:
+                continue
+            k, v = text.split(":", 1)
+            headers[k.lower().strip()] = v.strip()
+
+            while True:
+                line = self._in.readline()
+                if not line:
+                    return None
+                if line in (b"\r\n", b"\n"):
+                    break
+                try:
+                    text = line.decode("utf-8", "replace").strip()
+                except Exception:
+                    continue
+                if ":" not in text:
+                    continue
+                k, v = text.split(":", 1)
+                headers[k.lower().strip()] = v.strip()
+
+            if "content-length" not in headers:
+                return None
+            try:
+                length = int(headers["content-length"])
+            except ValueError:
+                return None
+            body = self._in.read(length)
+            if not body:
+                return None
+            self._mode = "headers"
+            msg = json.loads(body.decode("utf-8", "replace"))
+            if isinstance(msg, dict):
+                return msg
+            return None
+
+    def write_message(self, message: Dict[str, Any]) -> None:
+        if self._mode == "line":
+            payload = json.dumps(
+                message, ensure_ascii=False, separators=(",", ":"), default=str
+            ).encode("utf-8")
+            self._out.write(payload + b"\n")
+            self._out.flush()
+            return
+
+        body = json.dumps(
+            message, ensure_ascii=False, separators=(",", ":"), default=str
+        ).encode("utf-8")
+        header = f"Content-Length: {len(body)}\r\n\r\n".encode("utf-8")
+        self._out.write(header)
+        self._out.write(body)
+        self._out.flush()
+
+
+def main() -> None:
+    tools = WafIntelligenceTools()
+    rpc = StdioJsonRpc()
+
+    handlers: Dict[str, Callable[[Dict[str, Any]], Dict[str, Any]]] = {
+        "waf_capabilities": lambda a: tools.waf_capabilities(),
+        "analyze_waf": lambda a: tools.waf_analyze(
+            file=a.get("file"),
+            files=a.get("files"),
+            limit=int(a.get("limit", 3)),
+            min_severity=str(a.get("severity_threshold", "warning")),
+        ),
+        "waf_analyze": lambda a: tools.waf_analyze(**a),
+        "waf_assess": lambda a: tools.waf_assess(**a),
+        "waf_generate_gitops_proposals": lambda a: tools.waf_generate_gitops_proposals(
+            **a
+        ),
+    }
+
+    while True:
+        msg = rpc.read_message()
+        if msg is None:
+            return
+
+        method = msg.get("method")
+        msg_id = msg.get("id")
+        params = msg.get("params") or {}
+
+        try:
+            if method == "initialize":
+                result = {
+                    "protocolVersion": "2024-11-05",
+                    "serverInfo": {"name": "waf_intelligence", "version": "0.1.0"},
+                    "capabilities": {"tools": {}},
+                }
+                rpc.write_message({"jsonrpc": "2.0", "id": msg_id, "result": result})
+                continue
+
+            if method == "tools/list":
+                rpc.write_message(
+                    {"jsonrpc": "2.0", "id": msg_id, "result": {"tools": TOOLS}}
+                )
+                continue
+
+            if method == "tools/call":
+                tool_name = str(params.get("name") or "")
+                args = params.get("arguments") or {}
+
+                routing_action, shadow = layer0_entry(
+                    _shadow_query_repr(tool_name, args)
+                )
+                if routing_action != "HANDOFF_TO_LAYER1":
+                    rpc.write_message(
+                        {
+                            "jsonrpc": "2.0",
+                            "id": msg_id,
+                            "result": _mcp_text_result(
+                                _layer0_payload(routing_action, shadow), is_error=True
+                            ),
+                        }
+                    )
+                    continue
+
+                handler = handlers.get(tool_name)
+                if not handler:
+                    rpc.write_message(
+                        {
+                            "jsonrpc": "2.0",
+                            "id": msg_id,
+                            "result": _mcp_text_result(
+                                {
+                                    "ok": False,
+                                    "summary": f"Unknown tool: {tool_name}",
+                                    "data": {"known_tools": sorted(handlers.keys())},
+                                    "truncated": False,
+                                    "next_steps": ["Call tools/list"],
+                                },
+                                is_error=True,
+                            ),
+                        }
+                    )
+                    continue
+
+                payload = handler(args)
+                is_error = (
+                    not bool(payload.get("ok", True))
+                    if isinstance(payload, dict)
+                    else False
+                )
+                rpc.write_message(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": msg_id,
+                        "result": _mcp_text_result(payload, is_error=is_error),
+                    }
+                )
+                continue
+
+            # Ignore notifications.
+            if msg_id is None:
+                continue
+
+            rpc.write_message(
+                {
+                    "jsonrpc": "2.0",
+                    "id": msg_id,
+                    "result": _mcp_text_result(
+                        {"ok": False, "summary": f"Unsupported method: {method}"},
+                        is_error=True,
+                    ),
+                }
+            )
+        except Exception as e:  # noqa: BLE001
+            if msg_id is not None:
+                rpc.write_message(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": msg_id,
+                        "result": _mcp_text_result(
+                            {"ok": False, "summary": f"fatal error: {e}"},
+                            is_error=True,
+                        ),
+                    }
+                )
+
+
+def _shadow_query_repr(tool_name: str, tool_args: Dict[str, Any]) -> str:
+    if tool_name == "waf_capabilities":
+        return "List WAF Intelligence capabilities."
+    try:
+        return f"{tool_name}: {json.dumps(tool_args, sort_keys=True, default=str)}"
+    except Exception:
+        return f"{tool_name}: {str(tool_args)}"
+
+
+def _layer0_payload(routing_action: str, shadow: ShadowEvalResult) -> Dict[str, Any]:
+    if routing_action == "FAIL_CLOSED":
+        return {"ok": False, "summary": "Layer 0: cannot comply with this request."}
+    if routing_action == "HANDOFF_TO_GUARDRAILS":
+        reason = shadow.reason or "governance_violation"
+        return {
+            "ok": False,
+            "summary": f"Layer 0: governance violation detected ({reason}).",
+        }
+    if routing_action == "PROMPT_FOR_CLARIFICATION":
+        return {
+            "ok": False,
+            "summary": "Layer 0: request is ambiguous. Please clarify and retry.",
+        }
+    return {"ok": False, "summary": "Layer 0: unrecognized routing action; refusing."}
+
+
+if __name__ == "__main__":
+    main()