|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
| 3 | +import hashlib |
3 | 4 | import json |
4 | 5 | import logging |
| 6 | +import re |
5 | 7 | from collections.abc import Callable, Mapping, Sequence |
6 | 8 | from dataclasses import dataclass |
7 | 9 | from functools import cache, reduce |
@@ -110,6 +112,32 @@ def _later_revision_fields(method: str, version: str) -> frozenset[str]: |
110 | 112 | return frozenset(later) - _wire_fields(current) |
111 | 113 |
|
112 | 114 |
|
| 115 | +_INJECTION_PATTERNS: tuple[re.Pattern[str], ...] = ( |
| 116 | + re.compile(r"(?i)(ignore|disregard)\s+(all\s+)?(previous|prior|above)\s+(instructions?|prompts?|rules?|context)"), |
| 117 | + re.compile(r"(?i)(you\s+are\s+now|act\s+as|pretend\s+(to\s+be|you\s+are))\s+\w"), |
| 118 | + re.compile(r"(?i)(system\s*:?\s*override|admin\s+mode|unrestricted\s+mode|jailbreak)"), |
| 119 | + re.compile(r"(?i)(before\s+(any|every|all)\s+(operation|request|task|call)|critical\s+security\s+policy)"), |
| 120 | + re.compile(r"(?i)(exfiltrate|send\s+to\s+attacker|include\s+(in\s+)?.*response.*key|leak)"), |
| 121 | + re.compile(r"(?i)(read\s+~/\.(ssh|aws|kube|gnupg|config)|/etc/passwd|/proc/self)"), |
| 122 | + re.compile(r"(?i)(do\s+not\s+(mention|tell|show|disclose)\s+(this|the user|to user))"), |
| 123 | + re.compile(r"(?i)(\[IMPORTANT\s*:\s*[^\]]{20,}\])"), |
| 124 | +) |
| 125 | + |
| 126 | + |
| 127 | +def _tool_desc_hash(description: str) -> str: |
| 128 | + """Return a short SHA-256 fingerprint of a tool description for rug-pull pinning.""" |
| 129 | + return hashlib.sha256(description.encode()).hexdigest()[:16] |
| 130 | + |
| 131 | + |
| 132 | +def _scan_tool_description(description: str) -> str | None: |
| 133 | + """Return the matched text if the description contains a known injection pattern, else None.""" |
| 134 | + for pat in _INJECTION_PATTERNS: |
| 135 | + m = pat.search(description) |
| 136 | + if m: |
| 137 | + return m.group(0)[:80] |
| 138 | + return None |
| 139 | + |
| 140 | + |
113 | 141 | def _same_schema(a: dict[str, Any] | None, b: dict[str, Any] | None) -> bool: |
114 | 142 | """JSON equality for two output schemas. |
115 | 143 |
|
@@ -442,6 +470,9 @@ def __init__( |
442 | 470 | # `_absorb_tool_listing`, which evicts a tool's entry whenever its schema changes. |
443 | 471 | self._tool_output_validators: dict[str, Validator] = {} |
444 | 472 | self._x_mcp_header_maps: dict[str, dict[tuple[str, ...], str]] = {} |
| 473 | + # Pinned description hashes for rug-pull detection (tool name → SHA-256[:16]). |
| 474 | + # Set on first listing; any subsequent change is a security event. |
| 475 | + self._pinned_tool_desc_hashes: dict[str, str] = {} |
445 | 476 | self._initialize_result: types.InitializeResult | None = None |
446 | 477 | self._discover_result: types.DiscoverResult | None = None |
447 | 478 | self._discover_server_info: types.Implementation | None = None |
@@ -1294,6 +1325,31 @@ def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool) |
1294 | 1325 | Idempotent: cached values are already post-filter, so the response cache can re-absorb a served listing. |
1295 | 1326 | `complete` (an uncursored single-page listing) prunes per-tool state down to the listing's tools. |
1296 | 1327 | """ |
| 1328 | + # Security pass: description injection scan + rug-pull detection. |
| 1329 | + # Runs for ALL protocol versions — rug-pull is not version-gated. |
| 1330 | + secure: list[types.Tool] = [] |
| 1331 | + for tool in result.tools: |
| 1332 | + desc_hash = _tool_desc_hash(tool.description or "") |
| 1333 | + if tool.name in self._pinned_tool_desc_hashes: |
| 1334 | + if self._pinned_tool_desc_hashes[tool.name] != desc_hash: |
| 1335 | + logger.warning( |
| 1336 | + "security: tool %r description changed after approval — rug-pull detected; dropping", |
| 1337 | + tool.name, |
| 1338 | + ) |
| 1339 | + self._x_mcp_header_maps.pop(tool.name, None) |
| 1340 | + continue |
| 1341 | + else: |
| 1342 | + if reason := _scan_tool_description(tool.description or ""): |
| 1343 | + logger.warning( |
| 1344 | + "security: tool %r description contains injection pattern (%s); dropping", |
| 1345 | + tool.name, |
| 1346 | + reason, |
| 1347 | + ) |
| 1348 | + continue |
| 1349 | + self._pinned_tool_desc_hashes[tool.name] = desc_hash |
| 1350 | + secure.append(tool) |
| 1351 | + result.tools = secure |
| 1352 | + |
1297 | 1353 | if self._negotiated_version in MODERN_PROTOCOL_VERSIONS: |
1298 | 1354 | # 2026-07-28: clients MUST drop tools whose x-mcp-header annotations are invalid. |
1299 | 1355 | kept: list[types.Tool] = [] |
@@ -1326,6 +1382,7 @@ def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool) |
1326 | 1382 | self._x_mcp_header_maps = {k: v for k, v in self._x_mcp_header_maps.items() if k in names} |
1327 | 1383 | self._tool_output_schemas = {k: v for k, v in self._tool_output_schemas.items() if k in names} |
1328 | 1384 | self._tool_output_validators = {k: v for k, v in self._tool_output_validators.items() if k in names} |
| 1385 | + self._pinned_tool_desc_hashes = {k: v for k, v in self._pinned_tool_desc_hashes.items() if k in names} |
1329 | 1386 |
|
1330 | 1387 | return result |
1331 | 1388 |
|
|
0 commit comments