diff --git a/tests/test_havoc_mcp_server.py b/tests/test_havoc_mcp_server.py index 3ed46d8..dae26f3 100644 --- a/tests/test_havoc_mcp_server.py +++ b/tests/test_havoc_mcp_server.py @@ -2,8 +2,14 @@ import json import socket import threading import pytest +from starlette.testclient import TestClient -from tools.havoc_mcp_server import DebugAgentClient, DebugAgentError +from tools.havoc_mcp_server import ( + DebugAgentClient, + DebugAgentError, + build_asgi_app, + build_mcp_server, +) class _StubAgentServer: @@ -53,3 +59,52 @@ def test_check_bearer_token_rejects_missing_header(): def test_check_bearer_token_rejects_malformed_header(): assert check_bearer_token("secret123", expected="secret123") is False # missing "Bearer " prefix + +def test_check_bearer_token_rejects_empty_expected_token(): + # A set-but-empty HAVOC_MCP_TOKEN must not turn auth into a no-op. + assert check_bearer_token("Bearer ", expected="") is False + assert check_bearer_token("Bearer anything", expected="") is False + + +def _build_test_app(token: str = "secret123"): + """Build the real ASGI app (auth middleware wrapping the FastMCP Streamable HTTP + app) the same way __main__ does, but with a DebugAgentClient that never needs to + make a real connection -- these tests only exercise the auth middleware layer, + they never let a request reach the point of actually calling the debug agent.""" + client = DebugAgentClient(host="127.0.0.1", port=1) + server = build_mcp_server(client) + return build_asgi_app(server, token) + + +def test_asgi_app_rejects_request_with_no_auth_header(): + app = _build_test_app() + with TestClient(app) as test_client: + response = test_client.post("/mcp", json={}) + assert response.status_code == 401 + + +def test_asgi_app_rejects_request_with_wrong_token(): + app = _build_test_app() + with TestClient(app) as test_client: + response = test_client.post( + "/mcp", json={}, headers={"Authorization": "Bearer wrong"} + ) + assert response.status_code == 401 + + +def test_asgi_app_passes_through_request_with_correct_token(): + app = _build_test_app() + with TestClient(app) as test_client: + response = test_client.post( + "/mcp", + json={}, + headers={ + "Authorization": "Bearer secret123", + "Accept": "application/json, text/event-stream", + }, + ) + # The auth middleware must not 401 a correctly-authenticated request. Whatever + # status the underlying MCP session manager returns past that point (e.g. it may + # reject malformed MCP protocol bodies) is out of scope here -- a non-401 response + # is sufficient proof the middleware let the request through. + assert response.status_code != 401 diff --git a/tools/havoc_mcp_server.py b/tools/havoc_mcp_server.py index 7e71733..a53d4b7 100644 --- a/tools/havoc_mcp_server.py +++ b/tools/havoc_mcp_server.py @@ -5,10 +5,12 @@ running on the same Caspian box.""" from __future__ import annotations -import json +import hmac import os import socket +from starlette.responses import PlainTextResponse + from tools.havoc_debug_protocol import decode_message, encode_message @@ -37,12 +39,14 @@ class DebugAgentClient: def check_bearer_token(header_value: str | None, expected: str) -> bool: + if not expected: + return False if header_value is None: return False prefix = "Bearer " if not header_value.startswith(prefix): return False - return header_value[len(prefix):] == expected + return hmac.compare_digest(header_value[len(prefix):], expected) # --- MCP tool wiring (mcp SDK 1.28.1) -------------------------------------------- @@ -62,7 +66,7 @@ def check_bearer_token(header_value: str | None, expected: str) -> bool: from mcp.server.fastmcp import FastMCP -def build_mcp_server(client: DebugAgentClient, token: str) -> FastMCP: +def build_mcp_server(client: DebugAgentClient) -> FastMCP: mcp = FastMCP("havoc-runtime-observation") @mcp.tool() @@ -143,8 +147,6 @@ class _BearerAuthMiddleware: break if not check_bearer_token(header_value, expected=self._expected_token): - from starlette.responses import PlainTextResponse - response = PlainTextResponse("Unauthorized", status_code=401) await response(scope, receive, send) return @@ -162,8 +164,10 @@ if __name__ == "__main__": import uvicorn token = os.environ["HAVOC_MCP_TOKEN"] # required, no default -- fail loudly if unset + if not token: + raise ValueError("HAVOC_MCP_TOKEN must not be empty") client = DebugAgentClient(host="127.0.0.1", port=47474) - server = build_mcp_server(client, token) + server = build_mcp_server(client) app = build_asgi_app(server, token) # Bind 0.0.0.0 so the LAN can reach it. uvicorn.run(app, host="0.0.0.0", port=8765)