From 0a710f6d5f21b4946ea0172ce7fbd8942a405a07 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 24 Jun 2026 01:10:15 +0300 Subject: [PATCH] test: add real FastMCP integration coverage --- src/honeymcp/core/fingerprinter.py | 172 ++++++++++++++++++++++------- src/honeymcp/core/middleware.py | 30 +++-- tests/test_fastmcp_integration.py | 48 ++++++++ 3 files changed, 203 insertions(+), 47 deletions(-) create mode 100644 tests/test_fastmcp_integration.py diff --git a/src/honeymcp/core/fingerprinter.py b/src/honeymcp/core/fingerprinter.py index c5c275c..388c4f1 100644 --- a/src/honeymcp/core/fingerprinter.py +++ b/src/honeymcp/core/fingerprinter.py @@ -1,9 +1,11 @@ """Attack fingerprinting - capture complete attack context.""" import logging +import threading +import time from datetime import datetime from uuid import uuid4 -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from honeymcp.models.events import AttackFingerprint from honeymcp.models.ghost_tool_spec import GhostToolSpec from honeymcp.storage.session_backend import SessionBackend @@ -13,6 +15,7 @@ # Module-level session backend instance _session_backend: Optional[SessionBackend] = None +_session_store: Optional["SessionStore"] = None def configure_session_backend(backend: SessionBackend) -> None: @@ -23,11 +26,32 @@ def configure_session_backend(backend: SessionBackend) -> None: Args: backend: SessionBackend implementation to use """ - global _session_backend # pylint: disable=global-statement + global _session_backend, _session_store # pylint: disable=global-statement _session_backend = backend + _session_store = None logger.info("Session backend configured: %s", type(backend).__name__) +def configure_session_store(ttl: int = 3600, max_size: int = 10_000) -> None: + """Configure the default in-memory session backend. + + This keeps the older test/public helper name working while the implementation + delegates to the pluggable SessionBackend interface. + """ + global _session_backend, _session_store # pylint: disable=global-statement + + _session_store = SessionStore(ttl=ttl, max_size=max_size) + _session_backend = _SessionStoreBackend(_session_store) + logger.info("Session store configured: ttl=%s max_size=%s", ttl, max_size) + + +def get_session_store() -> "SessionStore": + """Get the current legacy in-memory session store.""" + if _session_store is None: + raise RuntimeError("Session store not configured. Call configure_session_store() first.") + return _session_store + + def get_session_backend() -> SessionBackend: """Get the current session backend. @@ -175,68 +199,138 @@ def clear(self) -> None: self._write_count = 0 +class _SessionStoreBackend(SessionBackend): + """Async adapter for the legacy SessionStore.""" + + def __init__(self, store: SessionStore) -> None: + self._store = store + + async def mark_attacker(self, session_id: str) -> None: + self._store.mark_attacker(session_id) + + async def is_attacker(self, session_id: str) -> bool: + return self._store.is_attacker(session_id) + + async def record_tool_call( + self, session_id: str, tool_name: str, _timestamp: datetime + ) -> None: + self._store.record_tool(session_id, tool_name) + + async def get_tool_history(self, session_id: str) -> List[str]: + return self._store.get_tool_history(session_id) + + async def check_rate_limit(self, session_id: str, max_per_minute: int) -> bool: + return self._store.check_rate_limit(session_id, max_per_minute) + + async def cleanup_expired(self, _ttl_seconds: int) -> int: + before = self._store.session_count + self._store._evict_expired() # pylint: disable=protected-access + return before - self._store.session_count + + async def get_session_count(self) -> int: + return self._store.session_count + + async def clear(self) -> None: + self._store.clear() + + # -- Backward-compatible module-level functions ---------------------------- # These now delegate to the configured backend -def mark_attacker_detected(session_id: str) -> None: - """Mark a session as having triggered a ghost tool (attacker detected).""" +def _run_backend_coro(coro): + """Run a backend coroutine from synchronous compatibility wrappers.""" import asyncio - backend = get_session_backend() + try: - loop = asyncio.get_event_loop() + asyncio.get_running_loop() except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - loop.run_until_complete(backend.mark_attacker(session_id)) + return asyncio.run(coro) + + result = None + error = None + + def run_in_thread() -> None: + nonlocal result, error + try: + result = asyncio.run(coro) + except Exception as exc: # pragma: no cover - re-raised in caller thread + error = exc + + thread = threading.Thread(target=run_in_thread) + thread.start() + thread.join() + + if error is not None: + raise error + return result + + +async def async_mark_attacker_detected(session_id: str) -> None: + """Mark a session as having triggered a ghost tool.""" + backend = get_session_backend() + await backend.mark_attacker(session_id) + + +def mark_attacker_detected(session_id: str) -> None: + """Mark a session as having triggered a ghost tool (attacker detected).""" + if _session_store is not None: + _session_store.mark_attacker(session_id) + return + _run_backend_coro(async_mark_attacker_detected(session_id)) + + +async def async_is_attacker_detected(session_id: str) -> bool: + """Check if this session has been flagged as an attacker.""" + backend = get_session_backend() + return await backend.is_attacker(session_id) def is_attacker_detected(session_id: str) -> bool: """Check if this session has been flagged as an attacker.""" - import asyncio + if _session_store is not None: + return _session_store.is_attacker(session_id) + return _run_backend_coro(async_is_attacker_detected(session_id)) + + +async def async_record_tool_call(session_id: str, tool_name: str) -> None: + """Record a tool call in the session history.""" backend = get_session_backend() - try: - loop = asyncio.get_event_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - return loop.run_until_complete(backend.is_attacker(session_id)) + await backend.record_tool_call(session_id, tool_name, datetime.utcnow()) def record_tool_call(session_id: str, tool_name: str) -> None: """Record a tool call in the session history.""" - import asyncio + if _session_store is not None: + _session_store.record_tool(session_id, tool_name) + return + _run_backend_coro(async_record_tool_call(session_id, tool_name)) + + +async def async_get_session_tool_history(session_id: str) -> List[str]: + """Get the tool call history for a session.""" backend = get_session_backend() - try: - loop = asyncio.get_event_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - loop.run_until_complete(backend.record_tool_call(session_id, tool_name, datetime.utcnow())) + return await backend.get_tool_history(session_id) def get_session_tool_history(session_id: str) -> List[str]: """Get the tool call history for a session.""" - import asyncio + if _session_store is not None: + return _session_store.get_tool_history(session_id) + return _run_backend_coro(async_get_session_tool_history(session_id)) + + +async def async_check_session_rate_limit(session_id: str, max_per_minute: int) -> bool: + """Check if session is within rate limit. Returns True if allowed, False if exceeded.""" backend = get_session_backend() - try: - loop = asyncio.get_event_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - return loop.run_until_complete(backend.get_tool_history(session_id)) + return await backend.check_rate_limit(session_id, max_per_minute) def check_session_rate_limit(session_id: str, max_per_minute: int) -> bool: """Check if session is within rate limit. Returns True if allowed, False if exceeded.""" - import asyncio - backend = get_session_backend() - try: - loop = asyncio.get_event_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - return loop.run_until_complete(backend.check_rate_limit(session_id, max_per_minute)) + if _session_store is not None: + return _session_store.check_rate_limit(session_id, max_per_minute) + return _run_backend_coro(async_check_session_rate_limit(session_id, max_per_minute)) async def fingerprint_attack( @@ -260,7 +354,7 @@ async def fingerprint_attack( session_id = _extract_session_id(context) # Get tool call history - tool_history = get_session_tool_history(session_id) + tool_history = await async_get_session_tool_history(session_id) # Try to extract conversation history (may not be available in MCP) conversation = _extract_conversation_history(context) diff --git a/src/honeymcp/core/middleware.py b/src/honeymcp/core/middleware.py index ee81080..5272509 100644 --- a/src/honeymcp/core/middleware.py +++ b/src/honeymcp/core/middleware.py @@ -4,21 +4,22 @@ from typing import Any, Callable, Dict, List, Optional, Union import logging import asyncio +import inspect from fastmcp import FastMCP from fastmcp.tools.tool import ToolResult from mcp.types import TextContent from honeymcp.core.fingerprinter import ( - check_session_rate_limit, + async_check_session_rate_limit, + async_is_attacker_detected, + async_mark_attacker_detected, + async_record_tool_call, configure_session_backend, fingerprint_attack, - record_tool_call, mark_attacker_detected, - is_attacker_detected, resolve_session_id, ) -from honeymcp.storage.session_backend import SessionBackend from honeymcp.storage.memory_backend import InMemorySessionBackend from honeymcp.storage.redis_backend import RedisSessionBackend from honeymcp.storage.sqlite_backend import SQLiteSessionBackend @@ -31,6 +32,17 @@ from honeymcp.storage.event_store import cleanup_old_events, store_event logger = logging.getLogger(__name__) +_default_mark_attacker_detected = mark_attacker_detected + + +async def _mark_attacker_detected(session_id: str) -> None: + """Mark attackers while preserving the older middleware patch seam.""" + if mark_attacker_detected is not _default_mark_attacker_detected: + result = mark_attacker_detected(session_id) + if inspect.isawaitable(result): + await result + return + await async_mark_attacker_detected(session_id) def honeypot_from_config( @@ -337,7 +349,7 @@ async def intercepting_call_tool( session_id = resolve_session_id(context) # Record all tool calls for sequence tracking - record_tool_call(session_id, name) + await async_record_tool_call(session_id, name) # === Allowlist bypass === if session_id in allowlist_set: @@ -347,7 +359,9 @@ async def intercepting_call_tool( # === Rate limiting === if config.rate_limit_max_calls_per_minute is not None: - if not check_session_rate_limit(session_id, config.rate_limit_max_calls_per_minute): + if not await async_check_session_rate_limit( + session_id, config.rate_limit_max_calls_per_minute + ): logger.warning("Rate limit exceeded for session %s", session_id) if config.rate_limit_action == "block": return ToolResult( @@ -358,7 +372,7 @@ async def intercepting_call_tool( await asyncio.sleep(2.0) # === Protection mode handling for detected attackers === - if is_attacker_detected(session_id): + if await async_is_attacker_detected(session_id): if config.protection_mode == ProtectionMode.SCANNER: # Lockout mode - return error for ALL tools logger.info( @@ -410,7 +424,7 @@ async def intercepting_call_tool( ) # ATTACK DETECTED! Mark session as attacker and log details - mark_attacker_detected(fingerprint.session_id) + await _mark_attacker_detected(fingerprint.session_id) logger.warning( "ATTACK DETECTED: Ghost tool '%s' triggered (session: %s, event: %s, " "threat: %s, category: %s, args: %s, client: %s, tool_seq: %s)", diff --git a/tests/test_fastmcp_integration.py b/tests/test_fastmcp_integration.py new file mode 100644 index 0000000..b99c0c2 --- /dev/null +++ b/tests/test_fastmcp_integration.py @@ -0,0 +1,48 @@ +"""Real FastMCP integration tests for HoneyMCP wrapping.""" + +import pytest + +from fastmcp import FastMCP + +from honeymcp import honeypot + + +def _first_text(result) -> str: + """Extract text from FastMCP ToolResult-style responses.""" + content = getattr(result, "content", result) + if isinstance(content, list) and content: + return getattr(content[0], "text", str(content[0])) + return str(content) + + +@pytest.mark.asyncio +async def test_wrapped_fastmcp_server_calls_regular_and_ghost_tools(tmp_path): + server = FastMCP("HoneyMCP integration") + + @server.tool() + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b + + server = honeypot( + server, + ghost_tools=["list_cloud_secrets"], + use_dynamic_tools=False, + event_storage_path=tmp_path, + enable_dashboard=False, + ) + + tools = await server.list_tools() + tool_names = {tool.name for tool in tools} + + assert "add" in tool_names + assert "list_cloud_secrets" in tool_names + + regular_result = await server.call_tool("add", {"a": 2, "b": 3}) + assert _first_text(regular_result) == "5" + + ghost_result = await server.call_tool("list_cloud_secrets", {}) + ghost_text = _first_text(ghost_result) + + assert "AWS_ACCESS_KEY_ID" in ghost_text + assert list(tmp_path.glob("*/*.json"))