From b2887753e93778131d182616b21fdcb9c7dfe3b5 Mon Sep 17 00:00:00 2001 From: pyr0ball Date: Sat, 4 Jul 2026 07:08:21 -0700 Subject: [PATCH] feat: add bearer token auth check for MCP server --- tests/test_havoc_mcp_server.py | 15 +++++++++++++++ tools/havoc_mcp_server.py | 9 +++++++++ 2 files changed, 24 insertions(+) diff --git a/tests/test_havoc_mcp_server.py b/tests/test_havoc_mcp_server.py index 816fd83..3ed46d8 100644 --- a/tests/test_havoc_mcp_server.py +++ b/tests/test_havoc_mcp_server.py @@ -38,3 +38,18 @@ def test_client_raises_on_error_response(): client = DebugAgentClient(host="127.0.0.1", port=stub.port) with pytest.raises(DebugAgentError, match="DR0-DR3"): client.call("set_watchpoint", {"addr": "0x5000", "size": 4, "mode": "write"}) + + +from tools.havoc_mcp_server import check_bearer_token + +def test_check_bearer_token_accepts_matching_token(): + assert check_bearer_token("Bearer secret123", expected="secret123") is True + +def test_check_bearer_token_rejects_wrong_token(): + assert check_bearer_token("Bearer wrong", expected="secret123") is False + +def test_check_bearer_token_rejects_missing_header(): + assert check_bearer_token(None, expected="secret123") is False + +def test_check_bearer_token_rejects_malformed_header(): + assert check_bearer_token("secret123", expected="secret123") is False # missing "Bearer " prefix diff --git a/tools/havoc_mcp_server.py b/tools/havoc_mcp_server.py index 659ba21..ea3e069 100644 --- a/tools/havoc_mcp_server.py +++ b/tools/havoc_mcp_server.py @@ -34,3 +34,12 @@ class DebugAgentClient: if not response.get("ok"): raise DebugAgentError(response.get("error", "unknown debug agent error")) return response.get("result", {}) + + +def check_bearer_token(header_value: str | None, expected: str) -> bool: + if header_value is None: + return False + prefix = "Bearer " + if not header_value.startswith(prefix): + return False + return header_value[len(prefix):] == expected