prooflayer-runtime 0.1.0__py3-none-any.whl

prooflayer/metrics.py

@@ -0,0 +1,266 @@
+"""
+Prometheus Metrics
+==================
+
+Lightweight metrics collection and exposition for ProofLayer.
+Uses only stdlib (no prometheus_client dependency).
+
+Exposes metrics at /metrics in Prometheus text exposition format via http.server.
+
+Usage:
+    from prooflayer.metrics import metrics, start_metrics_server
+
+    # Record a scan
+    metrics.record_scan(action="block", duration=0.003, risk_score=85,
+                        matched_rules=[("cmd-inject-curl", "command_injection")])
+
+    # Start HTTP endpoint
+    start_metrics_server(port=9090)
+"""
+
+import threading
+import time
+import math
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from typing import Dict, List, Tuple, Optional
+
+
+class _Counter:
+    """Thread-safe counter with optional labels."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._values: Dict[tuple, float] = {}
+
+    def inc(self, labels: Optional[Dict[str, str]] = None, value: float = 1.0):
+        key = tuple(sorted((labels or {}).items()))
+        with self._lock:
+            self._values[key] = self._values.get(key, 0.0) + value
+
+    def collect(self) -> Dict[tuple, float]:
+        with self._lock:
+            return dict(self._values)
+
+
+class _Gauge:
+    """Thread-safe gauge."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._values: Dict[tuple, float] = {}
+
+    def set(self, value: float, labels: Optional[Dict[str, str]] = None):
+        key = tuple(sorted((labels or {}).items()))
+        with self._lock:
+            self._values[key] = value
+
+    def collect(self) -> Dict[tuple, float]:
+        with self._lock:
+            return dict(self._values)
+
+
+class _HistogramLike:
+    """Tracks min, max, sum, count, and sorted observations for percentile computation."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._observations: List[float] = []
+        self._sum: float = 0.0
+        self._count: int = 0
+        self._min: float = float("inf")
+        self._max: float = float("-inf")
+
+    def observe(self, value: float):
+        with self._lock:
+            self._observations.append(value)
+            self._sum += value
+            self._count += 1
+            if value < self._min:
+                self._min = value
+            if value > self._max:
+                self._max = value
+
+    def collect(self) -> Dict[str, float]:
+        with self._lock:
+            if self._count == 0:
+                return {
+                    "count": 0, "sum": 0, "min": 0, "max": 0,
+                    "avg": 0, "p50": 0, "p95": 0, "p99": 0,
+                }
+            sorted_obs = sorted(self._observations)
+            return {
+                "count": self._count,
+                "sum": self._sum,
+                "min": self._min,
+                "max": self._max,
+                "avg": self._sum / self._count,
+                "p50": _percentile(sorted_obs, 0.50),
+                "p95": _percentile(sorted_obs, 0.95),
+                "p99": _percentile(sorted_obs, 0.99),
+            }
+
+
+def _percentile(sorted_data: List[float], pct: float) -> float:
+    """Compute percentile from sorted list using nearest-rank method."""
+    if not sorted_data:
+        return 0.0
+    idx = int(math.ceil(pct * len(sorted_data))) - 1
+    return sorted_data[max(0, idx)]
+
+
+class MetricsCollector:
+    """Central metrics collector for ProofLayer."""
+
+    def __init__(self):
+        self.scans_total = _Counter()
+        self.scan_duration = _HistogramLike()
+        self.rules_matched_total = _Counter()
+        self.risk_score_last = _Gauge()
+        self.risk_score_avg = _Gauge()
+        self.active_rules = _Gauge()
+
+        self._lock = threading.Lock()
+        self._score_sum: float = 0.0
+        self._score_count: int = 0
+        self.enabled: bool = False
+
+    def record_scan(
+        self,
+        action: str,
+        duration: float,
+        risk_score: int,
+        matched_rules: Optional[List[Tuple[str, str]]] = None,
+    ):
+        """
+        Record a completed scan.
+
+        Args:
+            action: The action taken (allow, warn, block, kill)
+            duration: Scan duration in seconds
+            risk_score: Computed risk score (0-100)
+            matched_rules: List of (rule_id, category) tuples
+        """
+        if not self.enabled:
+            return
+
+        self.scans_total.inc(labels={"action": action.lower()})
+        self.scan_duration.observe(duration)
+
+        self.risk_score_last.set(float(risk_score))
+        with self._lock:
+            self._score_sum += risk_score
+            self._score_count += 1
+            self.risk_score_avg.set(self._score_sum / self._score_count)
+
+        for rule_id, category in (matched_rules or []):
+            self.rules_matched_total.inc(
+                labels={"rule_id": rule_id, "category": category}
+            )
+
+    def set_active_rules(self, count: int):
+        """Set the number of currently active rules."""
+        if not self.enabled:
+            return
+        self.active_rules.set(float(count))
+
+    def render_prometheus(self) -> str:
+        """Render all metrics in Prometheus text exposition format."""
+        lines: List[str] = []
+
+        # prooflayer_scans_total
+        lines.append("# HELP prooflayer_scans_total Total number of scans performed.")
+        lines.append("# TYPE prooflayer_scans_total counter")
+        for labels, value in self.scans_total.collect().items():
+            label_str = _labels_to_str(labels)
+            lines.append(f"prooflayer_scans_total{{{label_str}}} {value}")
+
+        # prooflayer_scan_duration_seconds
+        dur = self.scan_duration.collect()
+        lines.append("# HELP prooflayer_scan_duration_seconds Scan duration statistics.")
+        lines.append("# TYPE prooflayer_scan_duration_seconds summary")
+        lines.append(f'prooflayer_scan_duration_seconds{{quantile="0.5"}} {dur["p50"]}')
+        lines.append(f'prooflayer_scan_duration_seconds{{quantile="0.95"}} {dur["p95"]}')
+        lines.append(f'prooflayer_scan_duration_seconds{{quantile="0.99"}} {dur["p99"]}')
+        lines.append(f'prooflayer_scan_duration_seconds_count {dur["count"]}')
+        lines.append(f'prooflayer_scan_duration_seconds_sum {dur["sum"]}')
+        lines.append(f'prooflayer_scan_duration_seconds_min {dur["min"]}')
+        lines.append(f'prooflayer_scan_duration_seconds_max {dur["max"]}')
+        lines.append(f'prooflayer_scan_duration_seconds_avg {dur["avg"]}')
+
+        # prooflayer_rules_matched_total
+        lines.append("# HELP prooflayer_rules_matched_total Total rules matched.")
+        lines.append("# TYPE prooflayer_rules_matched_total counter")
+        for labels, value in self.rules_matched_total.collect().items():
+            label_str = _labels_to_str(labels)
+            lines.append(f"prooflayer_rules_matched_total{{{label_str}}} {value}")
+
+        # prooflayer_risk_score
+        lines.append("# HELP prooflayer_risk_score Risk score gauge.")
+        lines.append("# TYPE prooflayer_risk_score gauge")
+        for labels, value in self.risk_score_last.collect().items():
+            lines.append(f'prooflayer_risk_score{{stat="last"}} {value}')
+        for labels, value in self.risk_score_avg.collect().items():
+            lines.append(f'prooflayer_risk_score{{stat="avg"}} {value}')
+
+        # prooflayer_active_rules
+        lines.append("# HELP prooflayer_active_rules Number of active detection rules.")
+        lines.append("# TYPE prooflayer_active_rules gauge")
+        for labels, value in self.active_rules.collect().items():
+            lines.append(f"prooflayer_active_rules {value}")
+
+        lines.append("")  # trailing newline
+        return "\n".join(lines)
+
+
+def _labels_to_str(labels: tuple) -> str:
+    """Convert label tuple to Prometheus label string."""
+    parts = [f'{k}="{v}"' for k, v in labels]
+    return ",".join(parts)
+
+
+# Singleton collector
+metrics = MetricsCollector()
+
+
+class _MetricsHandler(BaseHTTPRequestHandler):
+    """HTTP handler that serves Prometheus metrics."""
+
+    def do_GET(self):
+        if self.path == "/metrics":
+            body = metrics.render_prometheus().encode("utf-8")
+            self.send_response(200)
+            self.send_header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+    def log_message(self, format, *args):
+        # Suppress default access log noise
+        pass
+
+
+_server_instance: Optional[HTTPServer] = None
+
+
+def start_metrics_server(port: int = 9090) -> HTTPServer:
+    """
+    Start the Prometheus metrics HTTP server on a background thread.
+
+    Args:
+        port: TCP port to listen on (default 9090)
+
+    Returns:
+        HTTPServer instance (call .shutdown() to stop)
+    """
+    global _server_instance
+    if _server_instance is not None:
+        return _server_instance
+
+    server = HTTPServer(("0.0.0.0", port), _MetricsHandler)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    _server_instance = server
+    return server

prooflayer/reporting/__init__.py

@@ -0,0 +1,5 @@
+"""Security report generation."""
+
+from .reporter import SecurityReporter
+
+__all__ = ["SecurityReporter"]

prooflayer/reporting/reporter.py

@@ -0,0 +1,190 @@
+"""
+Security Reporter
+=================
+
+Generates security reports in JSON and SARIF formats.
+"""
+
+import json
+import os
+import threading
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Dict, List, Any
+
+from ..utils.masking import mask_sensitive_data
+
+
+class SecurityReporter:
+    """
+    Generates security reports for detected threats.
+    """
+
+    def __init__(self, report_dir: str = "./security-reports"):
+        """
+        Initialize security reporter.
+
+        Args:
+            report_dir: Directory to write reports
+        """
+        self._lock = threading.Lock()
+        self.report_dir = Path(report_dir)
+        self.report_dir.mkdir(parents=True, exist_ok=True)
+        # Restrict report directory to owner only
+        os.chmod(self.report_dir, 0o700)
+
+    def generate_report(
+        self,
+        threat_type: str,
+        tool_name: str,
+        arguments: Dict[str, Any],
+        risk_score: int,
+        matched_rules: List[Any],
+        action: str,
+        scan_result: Any = None,
+    ) -> Dict[str, Any]:
+        """
+        Generate a security report.
+
+        Args:
+            threat_type: Type of threat detected
+            tool_name: MCP tool name
+            arguments: Tool arguments
+            risk_score: Calculated risk score
+            matched_rules: List of matched detection rules
+            action: Action taken (ALLOW/WARN/BLOCK/KILL)
+            scan_result: Optional ScanResult to auto-populate extended fields
+
+        Returns:
+            Dict containing report data and file path
+        """
+        timestamp = datetime.now(timezone.utc)
+        timestamp_str = timestamp.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+        # Mask sensitive data before writing to report
+        masked_arguments = mask_sensitive_data(arguments) if arguments else {}
+
+        # Build scoring breakdown from scan_result if available
+        scoring_breakdown: Dict[str, Any] = {}
+        owasp_mapping: List[str] = []
+        latency_ms = 0.0
+        scan_timestamp = timestamp_str
+        if scan_result is not None:
+            scoring_breakdown = getattr(scan_result, "scoring_breakdown", {})
+            owasp_mapping = getattr(scan_result, "owasp_mapping", [])
+            latency_ms = getattr(scan_result, "latency_ms", 0.0)
+            scan_timestamp = getattr(scan_result, "timestamp", timestamp_str)
+
+        report = {
+            "prooflayer_version": "0.1.0",
+            "timestamp": timestamp_str,
+            "threat": {
+                "type": threat_type,
+                "tool": tool_name,
+                "arguments": masked_arguments,
+                "risk_score": risk_score,
+                "action": action
+            },
+            "detection": {
+                "rules_matched": [
+                    {
+                        "id": rule.id,
+                        "severity": rule.severity,
+                        "message": rule.message,
+                        "category": rule.category
+                    }
+                    for rule in matched_rules
+                ],
+                "confidence": "HIGH" if risk_score > 70 else "MEDIUM" if risk_score > 30 else "LOW",
+                "scoring_breakdown": scoring_breakdown,
+            },
+            "context": {
+                "timestamp_detected": timestamp_str,
+                "timestamp_received": scan_timestamp,
+                "latency_ms": round(latency_ms, 3),
+                "mcp_message_id": None,
+                "mcp_session_id": None,
+                "server": os.environ.get("MCP_SERVER_NAME", "unknown")
+            },
+            "owasp_mapping": owasp_mapping,
+        }
+
+        # Write JSON report with restricted permissions
+        report_filename = f"prooflayer-{timestamp.strftime('%Y%m%d-%H%M%S')}-{threat_type}.json"
+        report_path = self.report_dir / report_filename
+
+        with self._lock:
+            with open(report_path, "w") as f:
+                json.dump(report, f, indent=2)
+            os.chmod(report_path, 0o600)
+
+        report["report_path"] = str(report_path)
+
+        return report
+
+    def generate_sarif_report(
+        self,
+        threats: List[Dict[str, Any]]
+    ) -> str:
+        """
+        Generate SARIF format report for CI/CD integration.
+
+        Args:
+            threats: List of threat detections
+
+        Returns:
+            Path to SARIF report file
+        """
+        sarif = {
+            "$schema": "https://raw.githubusercontent.com/oasis-tcs/sarif-spec/master/Schemata/sarif-schema-2.1.0.json",
+            "version": "2.1.0",
+            "runs": [
+                {
+                    "tool": {
+                        "driver": {
+                            "name": "ProofLayer Runtime Security",
+                            "version": "0.1.0",
+                            "informationUri": "https://www.proof-layer.com"
+                        }
+                    },
+                    "results": [
+                        {
+                            "ruleId": threat["threat"]["type"],
+                            "level": self._sarif_level(threat["threat"]["risk_score"]),
+                            "message": {
+                                "text": f"Security threat detected in {threat['threat']['tool']}"
+                            },
+                            "locations": [
+                                {
+                                    "physicalLocation": {
+                                        "artifactLocation": {
+                                            "uri": threat["threat"]["tool"]
+                                        }
+                                    }
+                                }
+                            ]
+                        }
+                        for threat in threats
+                    ]
+                }
+            ]
+        }
+
+        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
+        sarif_path = self.report_dir / f"sarif-report-{timestamp}.sarif"
+
+        with self._lock:
+            with open(sarif_path, "w") as f:
+                json.dump(sarif, f, indent=2)
+            os.chmod(sarif_path, 0o600)
+
+        return str(sarif_path)
+
+    def _sarif_level(self, risk_score: int) -> str:
+        """Convert risk score to SARIF severity level."""
+        if risk_score >= 70:
+            return "error"
+        elif risk_score >= 30:
+            return "warning"
+        else:
+            return "note"

prooflayer/response/__init__.py

@@ -0,0 +1,6 @@
+"""Response actions for threat detection."""
+
+from .actions import ResponseAction, ThreatAction, SecurityViolation
+from .killer import ServerKiller
+
+__all__ = ["ResponseAction", "ThreatAction", "SecurityViolation", "ServerKiller"]

prooflayer/response/actions.py

@@ -0,0 +1,152 @@
+"""
+Response Actions
+================
+
+Handles ALLOW/WARN/BLOCK/KILL actions based on threat detection.
+"""
+
+import os
+import sys
+import signal
+import logging
+from enum import Enum
+from typing import Dict, Any, Optional
+
+from .killer import ServerKiller
+
+logger = logging.getLogger(__name__)
+
+
+class ThreatAction(Enum):
+    """Possible actions in response to a detected threat."""
+    ALLOW = "ALLOW"
+    WARN = "WARN"
+    BLOCK = "BLOCK"
+    KILL = "KILL"
+
+
+class ResponseAction:
+    """
+    Manages response actions for detected threats.
+    """
+
+    def __init__(
+        self,
+        default_action: str = "warn",
+        reporter: Optional[Any] = None
+    ):
+        """
+        Initialize response action handler.
+
+        Args:
+            default_action: Default action ("allow", "warn", "block", "kill")
+            reporter: SecurityReporter instance
+        """
+        self.default_action = default_action.upper()
+        self.reporter = reporter
+
+    def decide_action(self, risk_score: int) -> ThreatAction:
+        """
+        Decide action based on risk score.
+
+        Args:
+            risk_score: Risk score (0-100)
+
+        Returns:
+            ThreatAction to take
+        """
+        if risk_score <= 29:
+            return ThreatAction.ALLOW
+        elif risk_score <= 69:
+            return ThreatAction.WARN
+        elif risk_score <= 89:
+            return ThreatAction.BLOCK
+        else:
+            # Score 90-100 = KILL
+            if self.default_action == "KILL":
+                return ThreatAction.KILL
+            else:
+                return ThreatAction.BLOCK
+
+    def execute(
+        self,
+        action: ThreatAction,
+        context: Dict[str, Any]
+    ) -> None:
+        """
+        Execute the chosen action.
+
+        Args:
+            action: Action to execute
+            context: Context information (tool, arguments, risk_score, etc.)
+        """
+        if action == ThreatAction.ALLOW:
+            logger.info(f"ALLOW: {context.get('tool_name')} (score: {context.get('risk_score')})")
+
+        elif action == ThreatAction.WARN:
+            logger.warning(
+                f"WARN: {context.get('tool_name')} "
+                f"(score: {context.get('risk_score')}) - Suspicious activity detected"
+            )
+            if self.reporter:
+                self.reporter.generate_report(
+                    threat_type="suspicious_activity",
+                    tool_name=context.get("tool_name"),
+                    arguments=context.get("arguments", {}),
+                    risk_score=context.get("risk_score", 0),
+                    matched_rules=context.get("matched_rules", []),
+                    action="WARN"
+                )
+
+        elif action == ThreatAction.BLOCK:
+            logger.error(
+                f"BLOCK: {context.get('tool_name')} "
+                f"(score: {context.get('risk_score')}) - Threat blocked"
+            )
+            if self.reporter:
+                report = self.reporter.generate_report(
+                    threat_type="prompt_injection",
+                    tool_name=context.get("tool_name"),
+                    arguments=context.get("arguments", {}),
+                    risk_score=context.get("risk_score", 0),
+                    matched_rules=context.get("matched_rules", []),
+                    action="BLOCK"
+                )
+            raise SecurityViolation(
+                f"Tool call blocked: {context.get('tool_name')} "
+                f"(risk score: {context.get('risk_score')})"
+            )
+
+        elif action == ThreatAction.KILL:
+            logger.critical(
+                f"KILL: {context.get('tool_name')} "
+                f"(score: {context.get('risk_score')}) - Terminating MCP server"
+            )
+            if self.reporter:
+                report = self.reporter.generate_report(
+                    threat_type="critical_threat",
+                    tool_name=context.get("tool_name"),
+                    arguments=context.get("arguments", {}),
+                    risk_score=context.get("risk_score", 0),
+                    matched_rules=context.get("matched_rules", []),
+                    action="SERVER_KILLED"
+                )
+                logger.critical(f"Security report written: {report['report_path']}")
+
+            self.kill_server(context)
+
+    def kill_server(self, context: Dict[str, Any]) -> None:
+        """
+        Terminate the MCP server process.
+
+        Delegates to ServerKiller.kill().
+
+        Args:
+            context: Context information for logging
+        """
+        ServerKiller().kill(context)
+
+
+class SecurityViolation(Exception):
+    """Raised when a security threat is detected and blocked."""
+    pass

prooflayer/response/killer.py

@@ -0,0 +1,73 @@
+"""
+Server Killer
+=============
+
+Emergency MCP server termination for critical threats.
+"""
+
+import os
+import sys
+import signal
+import logging
+import time
+from typing import Dict, Any
+
+logger = logging.getLogger(__name__)
+
+
+class ServerKiller:
+    """Handles emergency termination of MCP server processes."""
+
+    def kill(self, context: Dict[str, Any]) -> None:
+        """
+        Terminate the MCP server process.
+
+        Writes an emergency log, sends SIGTERM, then SIGKILL if needed.
+
+        Args:
+            context: Context information (tool_name, risk_score, arguments, etc.)
+        """
+        logger.critical("=" * 80)
+        logger.critical("PROOFLAYER RUNTIME SECURITY: CRITICAL THREAT DETECTED")
+        logger.critical("Tool: %s", context.get("tool_name"))
+        logger.critical("Risk Score: %s", context.get("risk_score"))
+        logger.critical("Action: MCP SERVER TERMINATED")
+        logger.critical("=" * 80)
+
+        # Write emergency log with restricted permissions
+        emergency_path = "/tmp/prooflayer-emergency.log"
+        try:
+            fd = os.open(
+                emergency_path,
+                os.O_WRONLY | os.O_CREAT | os.O_APPEND,
+                0o600,
+            )
+            with os.fdopen(fd, "a") as f:
+                f.write(f"\n{'=' * 80}\n")
+                f.write(f"TIMESTAMP: {context.get('timestamp', 'N/A')}\n")
+                f.write(f"TOOL: {context.get('tool_name')}\n")
+                f.write(f"RISK SCORE: {context.get('risk_score')}\n")
+                f.write(f"ARGUMENTS: {context.get('arguments')}\n")
+                f.write(f"{'=' * 80}\n")
+            os.chmod(emergency_path, 0o600)
+        except Exception as e:
+            logger.error("Failed to write emergency log: %s", e)
+
+        # Kill the process
+        pid = os.getpid()
+        logger.critical("Sending SIGTERM to process %d", pid)
+
+        sys.stdout.flush()
+        sys.stderr.flush()
+
+        try:
+            os.killpg(os.getpgid(pid), signal.SIGTERM)
+        except (PermissionError, ProcessLookupError):
+            os.kill(pid, signal.SIGTERM)
+
+        # If still running after 1 second, force kill
+        time.sleep(1)
+        try:
+            os.killpg(os.getpgid(pid), signal.SIGKILL)
+        except (PermissionError, ProcessLookupError):
+            os.kill(pid, signal.SIGKILL)