agent_os_kernel 3.1.0__py3-none-any.whl

agent_os/mcp_message_signer.py

@@ -0,0 +1,273 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""HMAC signing and replay protection for MCP messages."""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import logging
+import secrets
+import threading
+import uuid
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Callable
+
+from agent_os.mcp_protocols import InMemoryNonceStore, MCPNonceStore
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class MCPSignedEnvelope:
+    """A signed MCP message envelope."""
+
+    payload: str
+    nonce: str
+    timestamp: datetime
+    signature: str
+    sender_id: str | None = None
+
+
+@dataclass(frozen=True)
+class MCPVerificationResult:
+    """The result of verifying a signed MCP envelope."""
+
+    is_valid: bool
+    payload: str | None = None
+    sender_id: str | None = None
+    failure_reason: str | None = None
+
+    @classmethod
+    def success(cls, payload: str, sender_id: str | None) -> "MCPVerificationResult":
+        return cls(is_valid=True, payload=payload, sender_id=sender_id)
+
+    @classmethod
+    def failed(cls, reason: str) -> "MCPVerificationResult":
+        return cls(is_valid=False, failure_reason=reason)
+
+
+def _utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class MCPMessageSigner:
+    """Sign and verify MCP messages with replay protection.
+
+    The signer produces tamper-evident envelopes and tracks nonces inside a
+    replay window so previously accepted messages cannot be replayed
+    indefinitely. Persistence and nonce generation are injectable to support
+    deterministic tests and external storage backends.
+    """
+
+    def __init__(
+        self,
+        signing_key: bytes,
+        *,
+        replay_window: timedelta = timedelta(minutes=5),
+        nonce_cache_cleanup_interval: timedelta = timedelta(minutes=10),
+        max_nonce_cache_size: int = 10_000,
+        nonce_store: MCPNonceStore | None = None,
+        nonce_generator: Callable[[], str] | None = None,
+    ) -> None:
+        """Initialize the signer and replay-protection settings.
+
+        Args:
+            signing_key: Shared secret used to compute HMAC signatures.
+            replay_window: Maximum permitted clock skew and replay horizon.
+            nonce_cache_cleanup_interval: Minimum interval between automatic
+                nonce cleanup passes.
+            max_nonce_cache_size: Maximum number of nonces retained by the
+                default in-memory nonce store.
+            nonce_store: Optional nonce persistence backend. Defaults to the
+                in-memory LRU nonce store.
+            nonce_generator: Optional nonce factory for deterministic testing
+                or custom nonce generation strategies.
+        """
+        if signing_key is None:
+            raise ValueError("signing_key must not be None")
+        if len(signing_key) < 32:
+            raise ValueError("signing_key must be at least 32 bytes")
+        if replay_window <= timedelta(0):
+            raise ValueError("replay_window must be positive")
+        if nonce_cache_cleanup_interval <= timedelta(0):
+            raise ValueError("nonce_cache_cleanup_interval must be positive")
+        if max_nonce_cache_size <= 0:
+            raise ValueError("max_nonce_cache_size must be positive")
+
+        self._signing_key = signing_key
+        self.replay_window = replay_window
+        self.nonce_cache_cleanup_interval = nonce_cache_cleanup_interval
+        self.max_nonce_cache_size = max_nonce_cache_size
+        self._lock = threading.Lock()
+        self._nonce_store = nonce_store or InMemoryNonceStore(
+            max_entries=max_nonce_cache_size,
+        )
+        self._nonce_generator = nonce_generator or (lambda: uuid.uuid4().hex)
+        self._last_cleanup = _utcnow()
+
+    @classmethod
+    def from_base64_key(cls, base64_key: str) -> "MCPMessageSigner":
+        """Build a signer from a base64-encoded shared secret.
+
+        Args:
+            base64_key: Shared secret encoded as ASCII base64 text.
+
+        Returns:
+            A configured ``MCPMessageSigner`` instance using the decoded key
+            bytes.
+        """
+        if not base64_key or not base64_key.strip():
+            raise ValueError("base64_key must not be empty")
+        decoded = base64.b64decode(base64_key.encode("ascii"), validate=True)
+        return cls(decoded)
+
+    @staticmethod
+    def generate_key() -> bytes:
+        """Return a new 256-bit signing key.
+
+        Returns:
+            A cryptographically random 32-byte signing key.
+        """
+        return secrets.token_bytes(32)
+
+    @property
+    def cached_nonce_count(self) -> int:
+        """Return the number of tracked nonces.
+
+        Returns:
+            The number of replay-protection nonces currently retained by the
+            backing store when the store exposes a ``count`` helper.
+        """
+        with self._lock:
+            count = getattr(self._nonce_store, "count", None)
+            return int(count()) if callable(count) else 0
+
+    def sign_message(self, payload: str, sender_id: str | None = None) -> MCPSignedEnvelope:
+        """Create a signed message envelope.
+
+        Args:
+            payload: Serialized MCP payload to sign.
+            sender_id: Optional sender identifier included in the signature.
+
+        Returns:
+            A signed envelope containing the payload, nonce, timestamp, and
+            computed signature.
+        """
+        if payload is None:
+            raise ValueError("payload must not be None")
+        if not payload.strip():
+            raise ValueError("payload must not be empty")
+
+        timestamp = _utcnow()
+        nonce = self._nonce_generator()
+        signature = self._compute_signature(
+            nonce=nonce,
+            timestamp=timestamp,
+            sender_id=sender_id,
+            payload=payload,
+        )
+        return MCPSignedEnvelope(
+            payload=payload,
+            nonce=nonce,
+            timestamp=timestamp,
+            sender_id=sender_id,
+            signature=signature,
+        )
+
+    def verify_message(self, envelope: MCPSignedEnvelope) -> MCPVerificationResult:
+        """Verify an envelope's signature and replay metadata.
+
+        Args:
+            envelope: Signed message envelope to validate.
+
+        Returns:
+            ``MCPVerificationResult`` describing whether the envelope is valid
+            and, on success, exposing the verified payload and sender identity.
+            Invalid signatures, replayed nonces, or unexpected errors fail
+            closed.
+        """
+        if envelope is None:
+            raise ValueError("envelope must not be None")
+
+        try:
+            now = _utcnow()
+            age = now - envelope.timestamp
+            if age > self.replay_window or age < -self.replay_window:
+                return MCPVerificationResult.failed("Message timestamp outside replay window.")
+
+            with self._lock:
+                self._maybe_cleanup_locked(now)
+                if self._nonce_store.has(envelope.nonce):
+                    logger.warning("Duplicate MCP nonce detected: %s", envelope.nonce)
+                    return MCPVerificationResult.failed("Duplicate nonce (replay detected).")
+                self._nonce_store.add(
+                    envelope.nonce,
+                    envelope.timestamp + self.replay_window,
+                )
+
+            expected_signature = self._compute_signature(
+                nonce=envelope.nonce,
+                timestamp=envelope.timestamp,
+                sender_id=envelope.sender_id,
+                payload=envelope.payload,
+            )
+            if not hmac.compare_digest(expected_signature, envelope.signature):
+                return MCPVerificationResult.failed("Invalid signature.")
+
+            return MCPVerificationResult.success(envelope.payload, envelope.sender_id)
+        except Exception as exc:
+            logger.error("MCP signature verification failed closed", exc_info=True)
+            return MCPVerificationResult.failed(f"Verification error (fail-closed): {exc}")
+
+    def cleanup_nonce_cache(self) -> int:
+        """Remove expired nonces and return the number removed.
+
+        Returns:
+            The number of expired nonces removed from the backing store.
+        """
+        with self._lock:
+            return self._cleanup_nonce_cache_locked(_utcnow())
+
+    def _compute_signature(
+        self,
+        *,
+        nonce: str,
+        timestamp: datetime,
+        sender_id: str | None,
+        payload: str,
+    ) -> str:
+        canonical = self._build_canonical_string(
+            nonce=nonce,
+            timestamp=timestamp,
+            sender_id=sender_id,
+            payload=payload,
+        )
+        digest = hmac.new(
+            self._signing_key,
+            canonical.encode("utf-8"),
+            hashlib.sha256,
+        ).digest()
+        return base64.b64encode(digest).decode("ascii")
+
+    @staticmethod
+    def _build_canonical_string(
+        *,
+        nonce: str,
+        timestamp: datetime,
+        sender_id: str | None,
+        payload: str,
+    ) -> str:
+        timestamp_ms = int(timestamp.timestamp() * 1000)
+        return f"{nonce}|{timestamp_ms}|{sender_id or ''}|{payload}"
+
+    def _maybe_cleanup_locked(self, now: datetime) -> None:
+        if now - self._last_cleanup >= self.nonce_cache_cleanup_interval:
+            self._cleanup_nonce_cache_locked(now)
+
+    def _cleanup_nonce_cache_locked(self, now: datetime) -> int:
+        expired = self._nonce_store.cleanup()
+        self._last_cleanup = now
+        return expired

agent_os/mcp_protocols.py

@@ -0,0 +1,161 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Persistence protocols and in-memory defaults for MCP governance components."""
+
+from __future__ import annotations
+
+import threading
+from collections import OrderedDict
+from datetime import datetime, timezone
+from typing import Any, Callable, Protocol
+
+
+def _utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class MCPSessionStore(Protocol):
+    """Persistence contract for authenticated MCP sessions."""
+
+    def get(self, session_id: str) -> Any | None:
+        """Return the stored session for *session_id*, if present."""
+
+    def set(self, session: Any) -> None:
+        """Persist *session* keyed by its session identifier."""
+
+    def delete(self, session_id: str) -> bool:
+        """Delete *session_id* and return ``True`` when it existed."""
+
+
+class MCPNonceStore(Protocol):
+    """Persistence contract for replay-protection nonces."""
+
+    def has(self, nonce: str) -> bool:
+        """Return ``True`` when *nonce* is still tracked."""
+
+    def add(self, nonce: str, expires_at: datetime) -> None:
+        """Track *nonce* until *expires_at*."""
+
+    def cleanup(self) -> int:
+        """Remove expired entries and return the number removed."""
+
+
+class MCPRateLimitStore(Protocol):
+    """Persistence contract for per-agent rate-limit buckets."""
+
+    def get_bucket(self, agent_id: str) -> Any | None:
+        """Return the bucket state stored for *agent_id*, if present."""
+
+    def set_bucket(self, agent_id: str, bucket: Any) -> None:
+        """Persist *bucket* for *agent_id*."""
+
+
+class MCPAuditSink(Protocol):
+    """Persistence contract for structured audit records."""
+
+    def record(self, entry: dict[str, Any]) -> None:
+        """Persist a structured audit *entry*."""
+
+
+class InMemorySessionStore:
+    """Thread-safe in-memory session storage."""
+
+    def __init__(self) -> None:
+        self._sessions: dict[str, Any] = {}
+        self._lock = threading.Lock()
+
+    def get(self, session_id: str) -> Any | None:
+        with self._lock:
+            return self._sessions.get(session_id)
+
+    def set(self, session: Any) -> None:
+        session_id = getattr(session, "token", None)
+        if not isinstance(session_id, str) or not session_id:
+            raise ValueError("session must provide a non-empty string token")
+        with self._lock:
+            self._sessions[session_id] = session
+
+    def delete(self, session_id: str) -> bool:
+        with self._lock:
+            return self._sessions.pop(session_id, None) is not None
+
+
+class InMemoryNonceStore:
+    """Thread-safe in-memory nonce storage with TTL cleanup and eviction."""
+
+    def __init__(
+        self,
+        *,
+        clock: Callable[[], datetime] = _utcnow,
+        max_entries: int = 10_000,
+    ) -> None:
+        if max_entries <= 0:
+            raise ValueError("max_entries must be positive")
+        self._clock = clock
+        self._max_entries = max_entries
+        self._nonces: OrderedDict[str, datetime] = OrderedDict()
+        self._lock = threading.Lock()
+
+    def has(self, nonce: str) -> bool:
+        with self._lock:
+            expires_at = self._nonces.get(nonce)
+            if expires_at is None:
+                return False
+            if expires_at <= self._clock():
+                self._nonces.pop(nonce, None)
+                return False
+            self._nonces.move_to_end(nonce)
+            return True
+
+    def add(self, nonce: str, expires_at: datetime) -> None:
+        with self._lock:
+            self._nonces[nonce] = expires_at
+            self._nonces.move_to_end(nonce)
+            while len(self._nonces) > self._max_entries:
+                self._nonces.popitem(last=False)
+
+    def cleanup(self) -> int:
+        removed = 0
+        now = self._clock()
+        with self._lock:
+            expired = [nonce for nonce, expires_at in self._nonces.items() if expires_at <= now]
+            for nonce in expired:
+                self._nonces.pop(nonce, None)
+            removed = len(expired)
+        return removed
+
+    def count(self) -> int:
+        with self._lock:
+            return len(self._nonces)
+
+
+class InMemoryRateLimitStore:
+    """Thread-safe in-memory rate-limit bucket storage."""
+
+    def __init__(self) -> None:
+        self._buckets: dict[str, Any] = {}
+        self._lock = threading.Lock()
+
+    def get_bucket(self, agent_id: str) -> Any | None:
+        with self._lock:
+            return self._buckets.get(agent_id)
+
+    def set_bucket(self, agent_id: str, bucket: Any) -> None:
+        with self._lock:
+            self._buckets[agent_id] = bucket
+
+
+class InMemoryAuditSink:
+    """Thread-safe in-memory audit sink for structured records."""
+
+    def __init__(self) -> None:
+        self._entries: list[dict[str, Any]] = []
+        self._lock = threading.Lock()
+
+    def record(self, entry: dict[str, Any]) -> None:
+        with self._lock:
+            self._entries.append(dict(entry))
+
+    def entries(self) -> list[dict[str, Any]]:
+        with self._lock:
+            return [dict(entry) for entry in self._entries]

agent_os/mcp_response_scanner.py

@@ -0,0 +1,232 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Scan MCP tool responses for prompt injection and data leakage."""
+
+from __future__ import annotations
+
+import logging
+import re
+from dataclasses import dataclass, field
+
+from agent_os.credential_redactor import CredentialRedactor
+
+logger = logging.getLogger(__name__)
+
+_INSTRUCTION_TAG_PATTERNS: tuple[re.Pattern[str], ...] = (
+    re.compile(
+        r"<(?:important|system|instruction|instructions|hidden|inject|admin|override|prompt|context|role)\b[^>]*>",
+        re.IGNORECASE,
+    ),
+    re.compile(r"\[(?:system|admin|instructions?)\]", re.IGNORECASE),
+)
+_IMPERATIVE_PATTERNS: tuple[re.Pattern[str], ...] = (
+    re.compile(r"ignore\s+(?:all\s+)?previous\s+(?:instructions?|context|rules?)", re.IGNORECASE),
+    re.compile(
+        r"(?:forget|disregard|override)\s+(?:all\s+)?(?:previous|above|prior|earlier)",
+        re.IGNORECASE,
+    ),
+    re.compile(r"\bexecute\s+this\b", re.IGNORECASE),
+    re.compile(r"\byou\s+are\s+now\b", re.IGNORECASE),
+    re.compile(r"\bnew\s+(?:role|instruction|directive|persona)\s*:", re.IGNORECASE),
+    re.compile(r"\bfrom\s+now\s+on\b", re.IGNORECASE),
+    re.compile(r"\bdo\s+not\s+(?:follow|obey|listen)\b", re.IGNORECASE),
+)
+_URL_PATTERN = re.compile(r"https?://[^\s<>'\"]+", re.IGNORECASE)
+_EXFILTRATION_URL_PATTERN = re.compile(
+    r"(?i)(?:\b(?:api[_-]?key|token|secret|payload|data|dump|upload|exfil|webhook)\b|webhook\.site|requestbin|pastebin|ngrok|transfer\.sh)"
+)
+
+
+@dataclass(frozen=True)
+class MCPResponseThreat:
+    """A threat detected in tool output."""
+
+    category: str
+    description: str
+    matched_pattern: str | None = None
+    details: dict[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class MCPResponseScanResult:
+    """Result of scanning an MCP tool response."""
+
+    is_safe: bool
+    tool_name: str
+    threats: list[MCPResponseThreat] = field(default_factory=list)
+
+    @classmethod
+    def safe(cls, tool_name: str) -> "MCPResponseScanResult":
+        return cls(is_safe=True, tool_name=tool_name, threats=[])
+
+    @classmethod
+    def unsafe(
+        cls,
+        tool_name: str,
+        *,
+        reason: str,
+        category: str = "error",
+    ) -> "MCPResponseScanResult":
+        return cls(
+            is_safe=False,
+            tool_name=tool_name,
+            threats=[MCPResponseThreat(category=category, description=reason)],
+        )
+
+
+class MCPResponseScanner:
+    """Scan tool responses before they are returned to an LLM.
+
+    The scanner looks for prompt-injection markers, credential leaks, and
+    likely exfiltration URLs, returning structured findings that callers can
+    use to block or sanitize unsafe tool output.
+    """
+
+    def scan_response(
+        self,
+        response_content: str | None,
+        tool_name: str = "unknown",
+    ) -> MCPResponseScanResult:
+        """Scan tool output and return a structured safety result.
+
+        Args:
+            response_content: Raw tool output to inspect.
+            tool_name: Human-readable tool name for reporting.
+
+        Returns:
+            An ``MCPResponseScanResult`` marked safe when no threats are found.
+            Unexpected scanner failures return an unsafe result.
+        """
+        try:
+            if not response_content:
+                return MCPResponseScanResult.safe(tool_name)
+
+            threats: list[MCPResponseThreat] = []
+            threats.extend(
+                self._scan_patterns(
+                    response_content,
+                    patterns=_INSTRUCTION_TAG_PATTERNS,
+                    category="instruction_injection",
+                    description="Instruction tag detected in tool response.",
+                )
+            )
+            threats.extend(
+                self._scan_patterns(
+                    response_content,
+                    patterns=_IMPERATIVE_PATTERNS,
+                    category="prompt_injection",
+                    description="Imperative instruction detected in tool response.",
+                )
+            )
+            threats.extend(self._scan_credential_leaks(response_content))
+            threats.extend(self._scan_exfiltration_urls(response_content))
+
+            if not threats:
+                return MCPResponseScanResult.safe(tool_name)
+
+            logger.warning(
+                "MCP response scan found %s issue(s) in tool %s",
+                len(threats),
+                tool_name,
+            )
+            return MCPResponseScanResult(is_safe=False, tool_name=tool_name, threats=threats)
+        except Exception:
+            logger.error("MCP response scanning failed closed", exc_info=True)
+            return MCPResponseScanResult.unsafe(
+                tool_name,
+                reason="Response scanner error (fail-closed).",
+            )
+
+    def sanitize_response(
+        self,
+        response_content: str | None,
+        tool_name: str = "unknown",
+    ) -> tuple[str, list[MCPResponseThreat]]:
+        """Strip instruction tags from tool output and report stripped threats.
+
+        Args:
+            response_content: Raw tool output to sanitize.
+            tool_name: Human-readable tool name for reporting.
+
+        Returns:
+            A tuple of ``(sanitized_content, stripped_threats)``. On failure the
+            method returns an empty string and a fail-closed error finding.
+        """
+        try:
+            if not response_content:
+                return "", []
+
+            sanitized = response_content
+            stripped: list[MCPResponseThreat] = []
+            for pattern in _INSTRUCTION_TAG_PATTERNS:
+                for match in pattern.finditer(sanitized):
+                    stripped.append(
+                        MCPResponseThreat(
+                            category="instruction_injection",
+                            description="Instruction tag stripped from tool response.",
+                            matched_pattern=match.group(0),
+                        )
+                    )
+                sanitized = pattern.sub("", sanitized)
+
+            return sanitized, stripped
+        except Exception:
+            logger.error("MCP response sanitization failed closed", exc_info=True)
+            return "", [
+                MCPResponseThreat(
+                    category="error",
+                    description=(
+                        f"Response sanitization failed for tool '{tool_name}' (fail-closed)."
+                    ),
+                )
+            ]
+
+    @staticmethod
+    def _scan_patterns(
+        content: str,
+        *,
+        patterns: tuple[re.Pattern[str], ...],
+        category: str,
+        description: str,
+    ) -> list[MCPResponseThreat]:
+        threats: list[MCPResponseThreat] = []
+        for pattern in patterns:
+            for match in pattern.finditer(content):
+                threats.append(
+                    MCPResponseThreat(
+                        category=category,
+                        description=description,
+                        matched_pattern=match.group(0),
+                    )
+                )
+        return threats
+
+    @staticmethod
+    def _scan_credential_leaks(content: str) -> list[MCPResponseThreat]:
+        threats: list[MCPResponseThreat] = []
+        for credential_match in CredentialRedactor.find_matches(content):
+            threats.append(
+                MCPResponseThreat(
+                    category="credential_leak",
+                    description=f"{credential_match.name} detected in tool response.",
+                    matched_pattern=credential_match.name,
+                    details={"credential_type": credential_match.name},
+                )
+            )
+        return threats
+
+    @staticmethod
+    def _scan_exfiltration_urls(content: str) -> list[MCPResponseThreat]:
+        threats: list[MCPResponseThreat] = []
+        for match in _URL_PATTERN.finditer(content):
+            url = match.group(0)
+            if not _EXFILTRATION_URL_PATTERN.search(url):
+                continue
+            threats.append(
+                MCPResponseThreat(
+                    category="data_exfiltration",
+                    description="Potential data exfiltration URL detected in tool response.",
+                    matched_pattern=url,
+                )
+            )
+        return threats