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