prooflayer-runtime 0.1.0__py3-none-any.whl

prooflayer/runtime/transport.py

@@ -0,0 +1,306 @@
+"""
+ProofLayer Transport Proxy
+===========================
+
+HTTP reverse proxy that intercepts MCP JSON-RPC tool calls
+for security scanning. Designed for Rick Spencer's simple-mcp
+(Go-based MCP server that speaks HTTP).
+
+Usage:
+    proxy = ProofLayerTransportProxy(listen_port=8080, backend_port=8081)
+    proxy.start()  # blocking
+    # or
+    proxy.start_background()  # returns Thread
+"""
+
+import json
+import logging
+import threading
+import time
+from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
+from typing import Optional, Dict, Any
+
+import httpx
+
+from ..detection.engine import DetectionEngine
+from ..detection.models import ScanResult
+from ..reporting.reporter import SecurityReporter
+from ..response.actions import ResponseAction, ThreatAction
+
+logger = logging.getLogger(__name__)
+
+
+class ProofLayerTransportProxy:
+    """HTTP reverse proxy with MCP tool call security scanning."""
+
+    def __init__(
+        self,
+        listen_port: int = 8080,
+        backend_port: int = 8081,
+        backend_host: str = "127.0.0.1",
+        detection_engine: Optional[DetectionEngine] = None,
+        reporter: Optional[SecurityReporter] = None,
+        response_action: Optional[ResponseAction] = None,
+        rules_dir: Optional[str] = None,
+        report_dir: str = "./security-reports",
+        action_on_threat: str = "block",
+        kill_on_threat: bool = False,
+    ):
+        self.listen_port = listen_port
+        self.backend_host = backend_host
+        self.backend_port = backend_port
+        self.kill_on_threat = kill_on_threat
+
+        # Initialize components (use provided or create defaults)
+        self.engine = detection_engine or DetectionEngine(rules_dir=rules_dir)
+        self.reporter = reporter or SecurityReporter(report_dir=report_dir)
+        self.response_action = response_action or ResponseAction(
+            default_action=action_on_threat, reporter=self.reporter
+        )
+
+        self._server: Optional[ThreadingHTTPServer] = None
+        self._thread: Optional[threading.Thread] = None
+        self._client = httpx.Client(
+            base_url=f"http://{backend_host}:{backend_port}",
+            timeout=30.0,
+        )
+
+    def start(self):
+        """Start the proxy server (blocking)."""
+        handler = self._make_handler()
+        self._server = ThreadingHTTPServer(("0.0.0.0", self.listen_port), handler)
+        logger.info(
+            "ProofLayer proxy listening on :%d, forwarding to %s:%d",
+            self.listen_port, self.backend_host, self.backend_port,
+        )
+        try:
+            self._server.serve_forever()
+        except KeyboardInterrupt:
+            self.stop()
+
+    def start_background(self) -> threading.Thread:
+        """Start the proxy in a background thread. Returns the thread."""
+        handler = self._make_handler()
+        self._server = ThreadingHTTPServer(("0.0.0.0", self.listen_port), handler)
+        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
+        self._thread.start()
+        logger.info(
+            "ProofLayer proxy started in background on :%d -> %s:%d",
+            self.listen_port, self.backend_host, self.backend_port,
+        )
+        return self._thread
+
+    def stop(self):
+        """Stop the proxy server."""
+        if self._server:
+            self._server.shutdown()
+            self._server = None
+        if self._client:
+            self._client.close()
+        logger.info("ProofLayer proxy stopped")
+
+    def _make_handler(self):
+        """Create the request handler class with access to proxy instance."""
+        proxy = self
+
+        class ProxyHandler(BaseHTTPRequestHandler):
+            def do_POST(self):
+                content_length = int(self.headers.get("Content-Length", 0))
+                body = self.rfile.read(content_length) if content_length else b""
+
+                # Try to parse as JSON-RPC
+                try:
+                    payload = json.loads(body) if body else None
+                except (json.JSONDecodeError, UnicodeDecodeError):
+                    payload = None
+
+                # Handle batch JSON-RPC (array of requests)
+                if isinstance(payload, list):
+                    batch_responses = [None] * len(payload)
+                    items_to_forward = []  # (index, item) pairs
+
+                    for i, item in enumerate(payload):
+                        blocked, response = proxy._check_tool_call(item)
+                        if blocked:
+                            batch_responses[i] = response
+                        else:
+                            items_to_forward.append((i, item))
+
+                    # Forward non-blocked items individually to backend
+                    for idx, item in items_to_forward:
+                        try:
+                            item_body = json.dumps(item).encode("utf-8")
+                            backend_response = proxy._client.post(
+                                self.path,
+                                content=item_body,
+                                headers={
+                                    k: v for k, v in self.headers.items()
+                                    if k.lower() not in ("host", "content-length")
+                                },
+                            )
+                            try:
+                                batch_responses[idx] = json.loads(backend_response.content)
+                            except (json.JSONDecodeError, UnicodeDecodeError):
+                                batch_responses[idx] = {
+                                    "jsonrpc": "2.0",
+                                    "error": {"code": -32000, "message": "Invalid backend response"},
+                                    "id": item.get("id") if isinstance(item, dict) else None,
+                                }
+                        except httpx.ConnectError:
+                            batch_responses[idx] = {
+                                "jsonrpc": "2.0",
+                                "error": {"code": -32000, "message": "Backend unavailable"},
+                                "id": item.get("id") if isinstance(item, dict) else None,
+                            }
+                        except httpx.TimeoutException:
+                            batch_responses[idx] = {
+                                "jsonrpc": "2.0",
+                                "error": {"code": -32000, "message": "Backend timeout"},
+                                "id": item.get("id") if isinstance(item, dict) else None,
+                            }
+
+                    self._send_json(200, batch_responses)
+                    return
+
+                elif isinstance(payload, dict):
+                    blocked, response = proxy._check_tool_call(payload)
+                    if blocked:
+                        self._send_json(200, response)
+                        return
+
+                # Forward the request to backend
+                try:
+                    backend_response = proxy._client.post(
+                        self.path,
+                        content=body,
+                        headers={
+                            k: v for k, v in self.headers.items()
+                            if k.lower() not in ("host", "content-length")
+                        },
+                    )
+                    self.send_response(backend_response.status_code)
+                    for key, value in backend_response.headers.items():
+                        if key.lower() not in ("transfer-encoding", "content-encoding", "content-length"):
+                            self.send_header(key, value)
+                    response_body = backend_response.content
+                    self.send_header("Content-Length", str(len(response_body)))
+                    self.end_headers()
+                    self.wfile.write(response_body)
+                except httpx.ConnectError:
+                    self._send_json(502, {
+                        "jsonrpc": "2.0",
+                        "error": {"code": -32000, "message": "Backend unavailable"},
+                        "id": payload.get("id") if isinstance(payload, dict) else None,
+                    })
+                except httpx.TimeoutException:
+                    self._send_json(504, {
+                        "jsonrpc": "2.0",
+                        "error": {"code": -32000, "message": "Backend timeout"},
+                        "id": payload.get("id") if isinstance(payload, dict) else None,
+                    })
+
+            def do_GET(self):
+                """Forward GET requests transparently."""
+                try:
+                    backend_response = proxy._client.get(
+                        self.path,
+                        headers={
+                            k: v for k, v in self.headers.items()
+                            if k.lower() not in ("host",)
+                        },
+                    )
+                    self.send_response(backend_response.status_code)
+                    for key, value in backend_response.headers.items():
+                        if key.lower() not in ("transfer-encoding", "content-encoding", "content-length"):
+                            self.send_header(key, value)
+                    response_body = backend_response.content
+                    self.send_header("Content-Length", str(len(response_body)))
+                    self.end_headers()
+                    self.wfile.write(response_body)
+                except httpx.ConnectError:
+                    self.send_response(502)
+                    self.end_headers()
+
+            def _send_json(self, status_code, data):
+                body = json.dumps(data).encode("utf-8")
+                self.send_response(status_code)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(body)))
+                self.end_headers()
+                self.wfile.write(body)
+
+            def log_message(self, format, *args):
+                """Route access logs through Python logging."""
+                logger.debug("Proxy: %s", format % args)
+
+        return ProxyHandler
+
+    def _check_tool_call(self, payload: dict) -> tuple:
+        """
+        Check if a JSON-RPC payload is a tools/call and scan it.
+
+        Returns:
+            (blocked: bool, response: dict or None)
+        """
+        # Only intercept tools/call JSON-RPC method
+        method = payload.get("method", "")
+        if method != "tools/call":
+            return False, None
+
+        params = payload.get("params", {})
+        tool_name = params.get("name", "")
+        arguments = params.get("arguments", {})
+        request_id = payload.get("id")
+
+        if not tool_name:
+            return False, None
+
+        # Run detection
+        result = self.engine.scan(tool_name=tool_name, arguments=arguments)
+
+        if result.score >= self.engine.score_threshold["block"][0]:
+            # THREAT - block
+            threat_type = "unknown"
+            if result.matched_rules:
+                top_rule = max(result.matched_rules, key=lambda r: r.score)
+                threat_type = top_rule.category
+
+            self.reporter.generate_report(
+                threat_type=threat_type,
+                tool_name=tool_name,
+                arguments=arguments,
+                risk_score=result.score,
+                matched_rules=result.matched_rules,
+                action="BLOCK",
+                scan_result=result,
+            )
+
+            logger.warning(
+                "BLOCKED tool call via proxy: %s (score=%d, rules=%s)",
+                tool_name, result.score,
+                [r.id for r in result.matched_rules],
+            )
+
+            blocked_response = {
+                "jsonrpc": "2.0",
+                "result": {
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": f"Tool call blocked by ProofLayer: {tool_name} "
+                                    f"(risk score: {result.score})",
+                        }
+                    ],
+                    "isError": True,
+                },
+                "id": request_id,
+            }
+            return True, blocked_response
+
+        if result.score >= self.engine.score_threshold["warn"][0]:
+            logger.warning(
+                "SUSPICIOUS tool call via proxy: %s (score=%d)",
+                tool_name, result.score,
+            )
+
+        return False, None

prooflayer/runtime/wrapper.py

@@ -0,0 +1,265 @@
+"""
+ProofLayer Runtime Wrapper
+===========================
+
+Wraps MCP servers with runtime security monitoring.
+"""
+
+import os
+import sys
+import json
+import logging
+from typing import Any, Dict, Optional, Callable, Tuple
+from pathlib import Path
+
+from ..detection.engine import DetectionEngine
+from ..response.actions import ResponseAction, ThreatAction
+from ..reporting.reporter import SecurityReporter
+from ..config.loader import ConfigLoader
+from ..utils.logging import configure_logging
+from ..metrics import metrics, start_metrics_server
+
+
+logger = logging.getLogger(__name__)
+
+
+class ProofLayerRuntime:
+    """
+    Runtime security wrapper for MCP servers.
+
+    Usage:
+        runtime = ProofLayerRuntime(config_path="prooflayer.yaml")
+        protected_server = runtime.wrap(mcp_server)
+        protected_server.run()
+    """
+
+    def __init__(
+        self,
+        config_path: Optional[str] = None,
+        detection_rules: Optional[str] = None,
+        action_on_threat: str = "kill",
+        report_dir: Optional[str] = None,
+        score_threshold: Optional[Dict[str, Tuple[int, ...]]] = None
+    ):
+        """
+        Initialize ProofLayer Runtime.
+
+        Args:
+            config_path: Path to YAML config file
+            detection_rules: Rules to load ("prompt-injection", "all")
+            action_on_threat: Action on threat detection ("allow", "warn", "block", "kill")
+            report_dir: Directory for security reports
+            score_threshold: Dict with allow/warn/block ranges
+        """
+        self.config = self._load_config(config_path)
+
+        # Override config with explicit parameters
+        if detection_rules:
+            self.config["detection"]["rules"] = detection_rules
+        if action_on_threat:
+            self.config["response"]["on_threat"] = action_on_threat
+        if report_dir:
+            self.config["response"]["report_dir"] = report_dir
+        if score_threshold:
+            self.config["detection"]["score_threshold"] = score_threshold
+
+        # Initialize components
+        self.detection_engine = DetectionEngine(
+            rules_dir=self.config["detection"].get("rules_dir"),
+            score_threshold=self.config["detection"].get("score_threshold"),
+            fail_closed=self.config["detection"].get("fail_closed", True)
+        )
+
+        self.reporter = SecurityReporter(
+            report_dir=self.config["response"].get("report_dir", "./security-reports")
+        )
+
+        self.response_action = ResponseAction(
+            default_action=self.config["response"].get("on_threat", "warn"),
+            reporter=self.reporter
+        )
+
+        # Apply structured logging from config
+        log_cfg = self.config.get("logging", {})
+        configure_logging(
+            level=log_cfg.get("level", "INFO"),
+            log_format=log_cfg.get("format", "text"),
+        )
+
+        # Enable metrics if configured
+        metrics_cfg = self.config.get("metrics", {})
+        if metrics_cfg.get("enabled", False):
+            metrics.enabled = True
+            metrics_port = metrics_cfg.get("port", 9090)
+            start_metrics_server(port=metrics_port)
+
+        logger.info(f"ProofLayer Runtime v0.1.0 initialized")
+        logger.info(f"Detection rules loaded: {len(self.detection_engine.rules)}")
+        logger.info(f"Default action on threat: {self.response_action.default_action}")
+
+    def _load_config(self, config_path: Optional[str]) -> Dict[str, Any]:
+        """Load configuration from file or use defaults."""
+        if config_path and os.path.exists(config_path):
+            return ConfigLoader.load(config_path)
+
+        # Default configuration
+        return {
+            "detection": {
+                "enabled": True,
+                "rules_dir": None,  # Will use packaged rules
+                "score_threshold": {
+                    "allow": (0, 29),
+                    "warn": (30, 69),
+                    "block": (70, 100)
+                }
+            },
+            "response": {
+                "on_threat": "warn",  # Conservative default
+                "report_dir": "./security-reports",
+                "alert_webhook": None
+            },
+            "performance": {
+                "max_latency_ms": 10,
+                "cache_rules": True
+            },
+            "logging": {
+                "level": "INFO",
+                "format": "json"
+            }
+        }
+
+    def wrap(self, mcp_server: Any) -> "ProtectedMCPServer":
+        """
+        Wrap an MCP server with ProofLayer security.
+
+        Args:
+            mcp_server: Original MCP server instance
+
+        Returns:
+            Protected MCP server with runtime security
+        """
+        return ProtectedMCPServer(
+            mcp_server=mcp_server,
+            detection_engine=self.detection_engine,
+            response_action=self.response_action,
+            reporter=self.reporter
+        )
+
+    def scan_tool_call(
+        self,
+        tool_name: str,
+        arguments: Dict[str, Any],
+        context: Optional[Dict[str, Any]] = None
+    ) -> Tuple[int, ThreatAction, Dict[str, Any]]:
+        """
+        Scan a single MCP tool call for threats.
+
+        Args:
+            tool_name: Name of the MCP tool
+            arguments: Tool call arguments
+            context: Additional context (message ID, timestamp, etc.)
+
+        Returns:
+            Tuple of (risk_score, action, detection_details)
+        """
+        # Run detection
+        risk_score, matched_rules = self.detection_engine.scan(
+            tool_name=tool_name,
+            arguments=arguments
+        )
+
+        # Determine action based on score
+        if risk_score <= self.config["detection"]["score_threshold"]["allow"][1]:
+            action = ThreatAction.ALLOW
+        elif risk_score <= self.config["detection"]["score_threshold"]["warn"][1]:
+            action = ThreatAction.WARN
+        else:
+            action = ThreatAction.BLOCK
+
+        detection_details = {
+            "risk_score": risk_score,
+            "matched_rules": [rule.id for rule in matched_rules],
+            "confidence": "HIGH" if risk_score > 70 else "MEDIUM" if risk_score > 30 else "LOW"
+        }
+
+        return risk_score, action, detection_details
+
+
+class ProtectedMCPServer:
+    """
+    MCP server wrapped with ProofLayer runtime security.
+    """
+
+    def __init__(
+        self,
+        mcp_server: Any,
+        detection_engine: DetectionEngine,
+        response_action: ResponseAction,
+        reporter: SecurityReporter
+    ):
+        self.mcp_server = mcp_server
+        self.detection_engine = detection_engine
+        self.response_action = response_action
+        self.reporter = reporter
+
+        # Wrap the original tool call handler
+        self._wrap_tool_handlers()
+
+    def _wrap_tool_handlers(self):
+        """Intercept MCP tool calls and inject security scanning."""
+        original_call_tool = getattr(self.mcp_server, "call_tool", None)
+
+        if original_call_tool:
+            def wrapped_call_tool(tool_name: str, arguments: Dict[str, Any]):
+                # Security scan before execution
+                risk_score, matched_rules = self.detection_engine.scan(
+                    tool_name=tool_name,
+                    arguments=arguments
+                )
+
+                # Determine action
+                action = self.response_action.decide_action(risk_score)
+
+                if action == ThreatAction.BLOCK or action == ThreatAction.KILL:
+                    # Extract threat type from highest-scoring matched rule
+                    threat_type = "unknown"
+                    if matched_rules:
+                        top_rule = max(matched_rules, key=lambda r: r.score)
+                        threat_type = top_rule.category
+
+                    # Generate security report
+                    report = self.reporter.generate_report(
+                        threat_type=threat_type,
+                        tool_name=tool_name,
+                        arguments=arguments,
+                        risk_score=risk_score,
+                        matched_rules=matched_rules,
+                        action=action.value
+                    )
+
+                    logger.error(f"THREAT DETECTED: {tool_name} - Score: {risk_score} - Action: {action.value}")
+
+                    if action == ThreatAction.KILL:
+                        # Kill the MCP server
+                        self.response_action.kill_server(report)
+
+                    raise SecurityError(f"Tool call blocked: {tool_name} (risk score: {risk_score})")
+
+                elif action == ThreatAction.WARN:
+                    logger.warning(f"SUSPICIOUS: {tool_name} - Score: {risk_score}")
+
+                # Allow the call to proceed
+                return original_call_tool(tool_name, arguments)
+
+            # Replace the original method
+            setattr(self.mcp_server, "call_tool", wrapped_call_tool)
+
+    def run(self):
+        """Run the protected MCP server."""
+        logger.info("Starting ProofLayer-protected MCP server")
+        return self.mcp_server.run()
+
+
+class SecurityError(Exception):
+    """Raised when a security threat is detected."""
+    pass

prooflayer/utils/__init__.py

@@ -0,0 +1,21 @@
+"""Utility functions."""
+
+from .entropy import calculate_shannon_entropy
+from .masking import mask_sensitive_data
+from .encoding import (
+    decode_hex_escapes,
+    decode_octal_escapes,
+    decode_unicode_escapes,
+    decode_url_encoding,
+    decode_base64_payloads,
+)
+
+__all__ = [
+    "calculate_shannon_entropy",
+    "mask_sensitive_data",
+    "decode_hex_escapes",
+    "decode_octal_escapes",
+    "decode_unicode_escapes",
+    "decode_url_encoding",
+    "decode_base64_payloads",
+]

prooflayer/utils/encoding.py

@@ -0,0 +1,87 @@
+"""
+Encoding Utilities
+==================
+
+Decode various encoding evasion techniques used to bypass detection.
+"""
+
+import re
+import base64
+import logging
+from typing import List
+
+logger = logging.getLogger(__name__)
+
+# Pre-compiled regex patterns for encoding detection
+HEX_ESCAPE_RE = re.compile(r"\\x([0-9a-fA-F]{2})")
+OCTAL_ESCAPE_RE = re.compile(r"\\([0-7]{1,3})")
+UNICODE_ESCAPE_RE = re.compile(r"\\u([0-9a-fA-F]{4})")
+URL_ENCODE_RE = re.compile(r"%([0-9a-fA-F]{2})")
+BASE64_RE = re.compile(r"(?:[A-Za-z0-9+/]{4}){1,}(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?")
+
+
+def decode_hex_escapes(text: str) -> str:
+    """Decode hex escape sequences like \\x63\\x75\\x72\\x6c -> curl."""
+    def _replace(match):
+        try:
+            return chr(int(match.group(1), 16))
+        except (ValueError, OverflowError):
+            return match.group(0)
+    return HEX_ESCAPE_RE.sub(_replace, text)
+
+
+def decode_octal_escapes(text: str) -> str:
+    """Decode octal escape sequences like \\143\\165\\162\\154 -> curl."""
+    def _replace(match):
+        try:
+            value = int(match.group(1), 8)
+            if value <= 0x10FFFF:
+                return chr(value)
+        except (ValueError, OverflowError):
+            pass
+        return match.group(0)
+    return OCTAL_ESCAPE_RE.sub(_replace, text)
+
+
+def decode_unicode_escapes(text: str) -> str:
+    """Decode unicode escape sequences like \\u0063\\u0075\\u0072\\u006c -> curl."""
+    def _replace(match):
+        try:
+            return chr(int(match.group(1), 16))
+        except (ValueError, OverflowError):
+            return match.group(0)
+    return UNICODE_ESCAPE_RE.sub(_replace, text)
+
+
+def decode_url_encoding(text: str) -> str:
+    """Decode URL-encoded sequences like %63%75%72%6c -> curl."""
+    def _replace(match):
+        try:
+            return chr(int(match.group(1), 16))
+        except (ValueError, OverflowError):
+            return match.group(0)
+    return URL_ENCODE_RE.sub(_replace, text)
+
+
+def decode_base64_payloads(text: str) -> str:
+    """
+    Detect and decode base64-encoded segments within text.
+
+    Prepends decoded content before the original so both are scanned.
+    Only decodes segments that produce valid UTF-8 text.
+    """
+    decoded_parts: List[str] = []
+    for match in BASE64_RE.finditer(text):
+        candidate = match.group(0)
+        if len(candidate) < 8:
+            continue
+        try:
+            decoded = base64.b64decode(candidate).decode("utf-8", errors="strict")
+            if any(c.isalpha() for c in decoded):
+                decoded_parts.append(decoded)
+        except Exception:
+            continue
+
+    if decoded_parts:
+        return " ".join(decoded_parts) + " " + text
+    return text