traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/governance/labeler.py

@@ -0,0 +1,250 @@
+"""GovernanceLabeler — Phase 2 of the governance pipeline.
+
+Side-effect-free enrichment. ``capability`` and ``structure`` are unioned on
+top of the existing Classification dimensions (labels are cumulative). By
+contrast ``source_labels`` are *dynamic* IFC clearance labels — they are
+computed FRESH for the current event and replace (never union with) whatever
+the base classification carried, matching their exclusion from the canonical
+action hash. Governance risk bonuses are accumulated through
+``_RiskModifiersBuilder`` and bounded to their documented caps on ``freeze()``.
+Returns GovernanceResult.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from traceforge.classify.core import Classification
+    from traceforge.governance.budget import BudgetThresholds
+    from traceforge.governance.drift import DriftAssessment, DriftDetector, DriftResult
+    from traceforge.governance.ifc import IFCChecker
+    from traceforge.governance.integrity import IntegrityVerifier
+    from traceforge.governance.mcp_drift import MCPIntegrityScanner
+    from traceforge.governance.pii import PIIScanner
+    from traceforge.governance.risk_wrapper import RiskModifiers
+    from traceforge.governance.types import EnrichmentContext
+
+
+@dataclass(frozen=True)
+class GovernanceResult:
+    """Output of Phase 2 — enriched classification + risk modifiers."""
+
+    classification: "Classification"
+    risk_modifiers: "RiskModifiers"
+    drift_result: "DriftResult | None" = None
+    mcp_alerts: tuple = ()  # tuple[MCPIntegrityAlert, ...]
+    mcp_deferred_writes: tuple = ()  # tuple[MCPDeferredWrite, ...] — committed by pipeline after finalization
+    integrity_deferred_writes: tuple = ()  # tuple[IntegrityWrite, ...] — committed by pipeline after finalization
+
+
+def _bounded(value: int, cap: int) -> int:
+    """Clamp ``value`` to the inclusive range ``[0, cap]``."""
+    if value < 0:
+        return 0
+    return cap if value > cap else value
+
+
+class _RiskModifiersBuilder:
+    """Accumulator that bounds governance risk bonuses to their documented caps.
+
+    Each governance signal contributes additively to a modifier; ``freeze()``
+    bounds every accumulated value to the cap documented on the corresponding
+    :class:`~traceforge.governance.risk_wrapper.RiskModifiers` field, then
+    returns an immutable ``RiskModifiers``. Centralizing the caps here means no
+    individual accumulation site can emit a value that violates the
+    ``RiskModifiers`` contract — e.g. the drift detector may return a phase
+    bonus up to 25, but ``phase_drift_bonus`` is documented "up to 20", so the
+    ceiling is enforced exactly once, on freeze.
+    """
+
+    # Documented caps — mirror the RiskModifiers field docstrings.
+    PHASE_DRIFT_CAP = 20  # phase_drift_bonus: +10 per anomaly, up to 20
+    MCP_DRIFT_CAP = 40  # mcp_drift_bonus: severity-weighted, up to 40
+    IFC_VIOLATIONS_CAP = 3  # ifc_violations: +10 per violation up to 30 => 3 counts
+    INTEGRITY_CAP = 10  # integrity_bonus: +10 when integrity_unverified
+    BUDGET_CAP = 5  # budget_bonus: +5 under budget pressure
+
+    __slots__ = ("_phase_drift", "_mcp_drift", "_ifc_violations", "_integrity", "_budget")
+
+    def __init__(self) -> None:
+        self._phase_drift = 0
+        self._mcp_drift = 0
+        self._ifc_violations = 0
+        self._integrity = 0
+        self._budget = 0
+
+    def add_phase_drift(self, bonus: int) -> None:
+        """Accumulate a phase-drift bonus (bounded to the cap on freeze)."""
+        self._phase_drift += bonus
+
+    def add_mcp_drift(self, bonus: int) -> None:
+        """Accumulate an MCP fingerprint-drift bonus (bounded on freeze)."""
+        self._mcp_drift += bonus
+
+    def add_ifc_violation(self, count: int = 1) -> None:
+        """Record IFC violation(s); the running count is bounded on freeze."""
+        self._ifc_violations += count
+
+    def set_integrity_unverified(self) -> None:
+        """Flag content integrity as unverified (fixed integrity bonus)."""
+        self._integrity = self.INTEGRITY_CAP
+
+    def set_budget_pressure(self) -> None:
+        """Flag budget pressure (fixed budget bonus)."""
+        self._budget = self.BUDGET_CAP
+
+    def freeze(self) -> "RiskModifiers":
+        """Bound every accumulated modifier to its cap and return ``RiskModifiers``.
+
+        Pure: does not mutate the builder, so it may be called more than once.
+        """
+        from traceforge.governance.risk_wrapper import RiskModifiers
+
+        return RiskModifiers(
+            phase_drift_bonus=_bounded(self._phase_drift, self.PHASE_DRIFT_CAP),
+            mcp_drift_bonus=_bounded(self._mcp_drift, self.MCP_DRIFT_CAP),
+            ifc_violations=_bounded(self._ifc_violations, self.IFC_VIOLATIONS_CAP),
+            integrity_bonus=_bounded(self._integrity, self.INTEGRITY_CAP),
+            budget_bonus=_bounded(self._budget, self.BUDGET_CAP),
+        )
+
+
+class GovernanceLabeler:
+    """Phase 2 labeler. Stateless — all state comes via EnrichmentContext."""
+
+    def __init__(
+        self,
+        pii_scanner: "PIIScanner | None" = None,
+        integrity_verifier: "IntegrityVerifier | None" = None,
+        mcp_scanner: "MCPIntegrityScanner | None" = None,
+        ifc_checker: "IFCChecker | None" = None,
+        drift_detector: "DriftDetector | None" = None,
+        budget_thresholds: "BudgetThresholds | None" = None,
+    ) -> None:
+        self._pii = pii_scanner
+        self._integrity = integrity_verifier
+        self._mcp = mcp_scanner
+        self._ifc = ifc_checker
+        self._drift = drift_detector
+        self._budget_thresholds = budget_thresholds
+
+    def label(self, ctx: "EnrichmentContext") -> GovernanceResult:
+        """Enrich classification with governance labels. Side-effect-free."""
+        cap: set[str] = set()
+        struct: set[str] = set()
+        src_labels: set[str] = set()
+        modifiers = _RiskModifiersBuilder()
+
+        # PII scanning
+        if self._pii:
+            self._pii.scan(ctx, cap, struct)
+
+        # Content integrity — CHECK against the existing baseline (read-only), then
+        # collect deferred (re)baseline prescriptions. The check runs first so drift
+        # from a prior baseline is flagged before the monitor commits the new baseline.
+        integrity_deferred_writes: tuple = ()
+        if self._integrity:
+            self._integrity.check_event(ctx, cap)
+            integrity_deferred_writes = tuple(self._integrity.pending_writes(ctx))
+
+        # MCP fingerprint drift — scan returns typed MCPScanResult
+        mcp_alerts: tuple = ()
+        mcp_deferred_writes: tuple = ()
+        if self._mcp:
+            scan_result = self._mcp.scan(ctx, cap)
+            mcp_alerts = scan_result.alerts
+            mcp_deferred_writes = scan_result.deferred_writes
+            severity_map = {"critical": 20, "warning": 10, "info": 5}
+            modifiers.add_mcp_drift(sum(severity_map[a.severity] for a in mcp_alerts))
+            if mcp_alerts:
+                struct.add("semantic_drift")
+
+        # IFC source labels
+        if self._ifc and ctx.session_state:
+            self._ifc_label_only(ctx, src_labels)
+            all_caps = ctx.base_classification.capability | frozenset(cap)
+            # Filter out current event's taint to prevent self-violation
+            prior_taints = [
+                t
+                for t in ctx.session_state.taint_ledger
+                if t.source_event_key != ctx.event.source_event_key
+            ]
+            if prior_taints and ctx.base_classification.effect in ("mutating", "destructive"):
+                struct.add("ifc_violation")
+                modifiers.add_ifc_violation()
+            elif any("ifc:" in l for l in src_labels) and "network_outbound" in all_caps:
+                struct.add("ifc_violation")
+                modifiers.add_ifc_violation()
+
+        # Phase drift — returns typed DriftAssessment with risk_bonus field
+        drift_result: "DriftAssessment | None" = None
+        if self._drift and ctx.session_state:
+            drift_result = self._drift.detect(ctx, ctx.session_state, cap)
+            if drift_result:
+                modifiers.add_phase_drift(drift_result.risk_bonus)
+                if drift_result.anomaly:
+                    struct.add("phase_anomaly")
+
+        # Budget pressure check (read-only from snapshot — no thresholds dependency)
+        if ctx.session_state and ctx.session_state.budget.pressure:
+            cap.add("budget_pressure")
+            modifiers.set_budget_pressure()
+
+        # Integrity bonus — surfaced by the integrity check above via the cap label.
+        if "integrity_unverified" in cap:
+            modifiers.set_integrity_unverified()
+
+        # Merge labels into classification. capability/structure are cumulative
+        # (unioned with the base); source_labels are dynamic IFC clearance labels
+        # set FRESH for this event (never carried forward from the base).
+        enriched = dataclasses.replace(
+            ctx.base_classification,
+            capability=ctx.base_classification.capability | frozenset(cap),
+            structure=ctx.base_classification.structure | frozenset(struct),
+            source_labels=frozenset(src_labels),
+        )
+
+        return GovernanceResult(
+            classification=enriched,
+            risk_modifiers=modifiers.freeze(),
+            drift_result=drift_result,
+            mcp_alerts=mcp_alerts,
+            mcp_deferred_writes=mcp_deferred_writes,
+            integrity_deferred_writes=integrity_deferred_writes,
+        )
+
+    def _ifc_label_only(self, ctx: "EnrichmentContext", src_labels: set[str]) -> None:
+        """Read-only IFC label assignment from snapshot state."""
+        if not ctx.session_state:
+            return
+        # If there are existing taints and current event writes, propagate label
+        if ctx.session_state.taint_ledger:
+            max_clearance = max(
+                (t.clearance for t in ctx.session_state.taint_ledger),
+                key=lambda c: {"public": 0, "internal": 1, "confidential": 2, "secret": 3}.get(
+                    c, 0
+                ),
+                default=None,
+            )
+            if max_clearance:
+                src_labels.add(f"ifc:{max_clearance}")
+
+    # ── Public facade methods for pipeline (eliminate private member access) ──
+
+    @property
+    def has_ifc(self) -> bool:
+        """Whether IFC checker is configured."""
+        return self._ifc is not None
+
+    def check_ifc(self, ctx: "EnrichmentContext", src_labels: set[str], state) -> None:
+        """Run IFC check (Phase 1 taint propagation). Delegates to IFCChecker."""
+        if self._ifc:
+            self._ifc.check(ctx, src_labels, state)
+
+    @property
+    def has_mcp_scanner(self) -> bool:
+        """Whether MCP integrity scanner is configured."""
+        return self._mcp is not None

traceforge/governance/mcp_drift.py

@@ -0,0 +1,328 @@
+"""MCP tool fingerprint drift detection with per-dimension semantic analysis."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from traceforge.classify.core import Classification
+    from traceforge.governance.persistence import SystemStore
+    from traceforge.governance.types import EnrichmentContext
+
+
+@dataclass(frozen=True)
+class MCPToolProfile:
+    """Frozen fingerprint of an MCP tool at first-seen time."""
+
+    tool_name: str
+    server_namespace: str
+    description_hash: str
+    schema_hash: str
+    registered_effect: str | None
+    registered_role: frozenset[str]
+    registered_capabilities: frozenset[str]
+    registered_scope: frozenset[str]
+    clearance: str | None
+    first_seen: datetime
+    last_seen: datetime
+
+
+@dataclass(frozen=True)
+class MCPIntegrityAlert:
+    """Individual drift alert with severity."""
+
+    tool_name: str
+    server: str
+    alert_type: Literal[
+        "effect_escalation",
+        "capability_gain",
+        "scope_expansion",
+        "description_change",
+        "schema_change",
+        "adversarial_pattern",
+    ]
+    previous: str
+    current: str
+    severity: Literal["info", "warning", "critical"]
+    timestamp: datetime
+
+
+# Effect precedence for escalation detection
+_EFFECT_ORDER = {"read_only": 0, "informational": 1, "mutating": 2, "destructive": 3}
+# Dangerous capabilities that trigger alerts when gained
+_DANGEROUS_CAPS = frozenset(
+    {"network_outbound", "elevated_privilege", "arbitrary_execution", "credential_exposure"}
+)
+# Scope escalation (higher = broader)
+_SCOPE_ORDER = {"file": 0, "project": 1, "repository": 2, "host": 3, "network": 4}
+
+# Adversarial patterns in tool descriptions
+_ADVERSARIAL_PATTERNS = [
+    re.compile(r"[\u200b-\u200f\u2028-\u202f\u2060-\u206f\ufeff]"),  # Invisible unicode
+    re.compile(
+        r"(?:ignore|disregard|forget)\s+(?:previous|prior|above)\s+(?:instructions?|rules?|constraints?)",
+        re.IGNORECASE,
+    ),  # Prompt injection
+    re.compile(r"(?:you\s+are|act\s+as|pretend|role[-\s]?play)", re.IGNORECASE),  # Role override
+    re.compile(
+        r"(?<![A-Fa-f0-9])[A-Za-z0-9+/]{40,}(?:={1,2})(?![A-Za-z0-9+/=])"
+    ),  # Base64 payload (requires padding, excludes hex-only)
+    re.compile(r"<!--.*?-->", re.DOTALL),  # Hidden HTML comments
+    re.compile(r"\{%.*?%\}", re.DOTALL),  # Template injection
+]
+
+
+@dataclass(frozen=True)
+class MCPDeferredWrite:
+    """Immutable deferred DB write — committed only after pipeline finalization."""
+
+    kind: Literal["upsert", "last_seen"]
+    server: str
+    tool_name: str
+    payload: str  # JSON for upsert, ISO timestamp for last_seen
+
+
+@dataclass(frozen=True)
+class MCPScanResult:
+    """Complete scan output — alerts, novelty flag, and deferred writes."""
+
+    alerts: tuple[MCPIntegrityAlert, ...]
+    is_new: bool
+    deferred_writes: tuple[MCPDeferredWrite, ...]
+
+
+class MCPIntegrityScanner:
+    """Full MCP integrity scanning: fingerprint drift + semantic analysis + adversarial detection."""
+
+    def __init__(self, store: "SystemStore") -> None:
+        self._store = store
+
+    def scan(self, ctx: "EnrichmentContext", cap: set[str]) -> MCPScanResult:
+        """Full scan: fingerprint comparison + semantic drift + adversarial patterns.
+
+        Returns MCPScanResult with alerts, is_new flag, and deferred writes.
+        Deferred writes MUST only be committed after pipeline finalization.
+        """
+        from traceforge.governance.types import ToolCallEvent
+
+        if not isinstance(ctx.event, ToolCallEvent):
+            return MCPScanResult(alerts=(), is_new=False, deferred_writes=())
+
+        server = ctx.event.mcp_server_name or ""
+        tool_name = ctx.event.tool_name or ""
+        if not server or not tool_name:
+            return MCPScanResult(alerts=(), is_new=False, deferred_writes=())
+
+        alerts: list[MCPIntegrityAlert] = []
+        deferred: list[MCPDeferredWrite] = []
+        now = datetime.now(timezone.utc)
+
+        # Compute current fingerprints
+        desc_hash = self._hash(ctx.event.tool_description or "")
+        schema_hash = self._hash(ctx.event.tool_schema_json or "")
+
+        # Check stored profile
+        stored = self._store.get_mcp_profile(server, tool_name)
+
+        if stored is None:
+            # First time — defer registration until pipeline finalization
+            cls = ctx.base_classification
+            deferred.append(
+                MCPDeferredWrite(
+                    kind="upsert",
+                    server=server,
+                    tool_name=tool_name,
+                    payload=json.dumps(
+                        {
+                            "description_hash": desc_hash,
+                            "schema_hash": schema_hash,
+                            "registered_effect": cls.effect,
+                            "role": sorted(cls.role),
+                            "capability": sorted(cls.capability),
+                            "scope": sorted(cls.scope),
+                            "clearance": None,
+                            "first_seen": now.isoformat(),
+                            "last_seen": now.isoformat(),
+                        }
+                    ),
+                )
+            )
+            # Still scan description for adversarial content on first-seen
+            desc_alerts = self.scan_description(
+                ctx.event.tool_description or "", tool_name, server, now
+            )
+            if desc_alerts:
+                alerts.extend(desc_alerts)
+                cap.add("mcp_drift")
+            return MCPScanResult(
+                alerts=tuple(alerts),
+                is_new=True,
+                deferred_writes=tuple(deferred),
+            )
+
+        # ── Fingerprint comparison ──
+        if stored["description_hash"] != desc_hash:
+            alerts.append(
+                MCPIntegrityAlert(
+                    tool_name=tool_name,
+                    server=server,
+                    alert_type="description_change",
+                    previous=stored["description_hash"][:16],
+                    current=desc_hash[:16],
+                    severity="warning",
+                    timestamp=now,
+                )
+            )
+
+        if stored["schema_hash"] != schema_hash:
+            alerts.append(
+                MCPIntegrityAlert(
+                    tool_name=tool_name,
+                    server=server,
+                    alert_type="schema_change",
+                    previous=stored["schema_hash"][:16],
+                    current=schema_hash[:16],
+                    severity="critical",
+                    timestamp=now,
+                )
+            )
+
+        # ── Semantic drift: per-dimension comparison ──
+        semantic_alerts = self.check_semantic_drift(
+            tool_name, server, ctx.base_classification, stored, now
+        )
+        alerts.extend(semantic_alerts)
+
+        # ── Adversarial pattern scanning ──
+        desc_alerts = self.scan_description(
+            ctx.event.tool_description or "", tool_name, server, now
+        )
+        alerts.extend(desc_alerts)
+
+        # Apply labels based on alert severity
+        if alerts:
+            max_severity = max(a.severity for a in alerts)
+            if max_severity in ("warning", "critical"):
+                cap.add("mcp_drift")
+
+        # Defer last_seen update until pipeline finalization
+        deferred.append(
+            MCPDeferredWrite(
+                kind="last_seen",
+                server=server,
+                tool_name=tool_name,
+                payload=now.isoformat(),
+            )
+        )
+
+        return MCPScanResult(
+            alerts=tuple(alerts),
+            is_new=False,
+            deferred_writes=tuple(deferred),
+        )
+
+    def check_semantic_drift(
+        self,
+        tool: str,
+        server: str,
+        current: "Classification",
+        stored: dict,
+        now: datetime,
+    ) -> list[MCPIntegrityAlert]:
+        """Per-dimension comparison against registered profile."""
+        alerts: list[MCPIntegrityAlert] = []
+
+        # Effect escalation
+        reg_effect = stored.get("registered_effect")
+        if reg_effect and current.effect:
+            prev_order = _EFFECT_ORDER.get(reg_effect, 0)
+            curr_order = _EFFECT_ORDER.get(current.effect, 0)
+            if curr_order > prev_order:
+                severity = "critical" if curr_order >= 3 else "warning"
+                alerts.append(
+                    MCPIntegrityAlert(
+                        tool_name=tool,
+                        server=server,
+                        alert_type="effect_escalation",
+                        previous=reg_effect,
+                        current=current.effect,
+                        severity=severity,
+                        timestamp=now,
+                    )
+                )
+
+        # Capability gain
+        reg_caps = frozenset(stored.get("capability") or ())
+        new_caps = current.capability - reg_caps
+        dangerous_new = new_caps & _DANGEROUS_CAPS
+        if dangerous_new:
+            alerts.append(
+                MCPIntegrityAlert(
+                    tool_name=tool,
+                    server=server,
+                    alert_type="capability_gain",
+                    previous=",".join(sorted(reg_caps)),
+                    current=",".join(sorted(dangerous_new)),
+                    severity="critical",
+                    timestamp=now,
+                )
+            )
+
+        # Scope expansion
+        reg_scope = frozenset(stored.get("scope") or ())
+        new_scope = current.scope - reg_scope
+        if new_scope:
+            max_prev = max((_SCOPE_ORDER.get(s, 0) for s in reg_scope), default=0)
+            max_curr = max((_SCOPE_ORDER.get(s, 0) for s in new_scope), default=0)
+            if max_curr > max_prev:
+                alerts.append(
+                    MCPIntegrityAlert(
+                        tool_name=tool,
+                        server=server,
+                        alert_type="scope_expansion",
+                        previous=",".join(sorted(reg_scope)),
+                        current=",".join(sorted(new_scope)),
+                        severity="warning" if max_curr < 3 else "critical",
+                        timestamp=now,
+                    )
+                )
+
+        return alerts
+
+    def scan_description(
+        self,
+        description: str,
+        tool_name: str,
+        server: str,
+        now: datetime,
+    ) -> list[MCPIntegrityAlert]:
+        """Detect adversarial patterns: invisible unicode, prompt injection, base64, hidden markup."""
+        if not description:
+            return []
+
+        alerts: list[MCPIntegrityAlert] = []
+        for pattern in _ADVERSARIAL_PATTERNS:
+            match = pattern.search(description)
+            if match:
+                alerts.append(
+                    MCPIntegrityAlert(
+                        tool_name=tool_name,
+                        server=server,
+                        alert_type="adversarial_pattern",
+                        previous="clean",
+                        current=f"pattern:{pattern.pattern[:40]}",
+                        severity="critical",
+                        timestamp=now,
+                    )
+                )
+                break  # One alert is enough
+
+        return alerts
+
+    def _hash(self, content: str) -> str:
+        return hashlib.sha256(content.encode()).hexdigest()