havoc-remaster/tests/test_havoc_mcp_server.py

import json
import socket
import threading
import pytest
from starlette.testclient import TestClient

from tools.havoc_mcp_server import (
    DebugAgentClient,
    DebugAgentError,
    build_asgi_app,
    build_mcp_server,
)


class _StubAgentServer:
    """A minimal in-process stand-in for tools/havoc_debug_agent_server.py, so this
    test never touches ctypes or a real Windows process."""

    def __init__(self, response: dict):
        self._response = response
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._sock.bind(("127.0.0.1", 0))
        self._sock.listen(1)
        self.port = self._sock.getsockname()[1]
        self._thread = threading.Thread(target=self._serve_one, daemon=True)
        self._thread.start()

    def _serve_one(self):
        conn, _ = self._sock.accept()
        conn.recv(4096)
        conn.sendall((json.dumps(self._response) + "\n").encode("utf-8"))
        conn.close()


def test_client_returns_result_on_success():
    stub = _StubAgentServer({"ok": True, "result": {"status": "running"}})
    client = DebugAgentClient(host="127.0.0.1", port=stub.port)
    result = client.call("get_status", {})
    assert result == {"status": "running"}


def test_client_raises_on_error_response():
    stub = _StubAgentServer({"ok": False, "error": "all 4 hardware debug registers (DR0-DR3) are in use"})
    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"})


class _FakeCreateConnectionSocket:
    """Stands in for the object socket.create_connection() returns, so tests can
    capture the timeout it was called with without opening a real connection or
    sleeping. Supports only what DebugAgentClient.call() actually uses."""

    def __init__(self, response: dict):
        self._unread = (json.dumps(response) + "\n").encode("utf-8")

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        return False

    def sendall(self, data):
        pass

    def recv(self, bufsize):
        chunk, self._unread = self._unread, b""
        return chunk


def test_client_call_uses_a_sane_default_socket_timeout(monkeypatch):
    # Regression guard for the previous hardcoded `timeout=60` on every op: normal,
    # fast ops (get_status, set_breakpoint, ...) should still get some bounded default
    # even though continue_execution now needs to ask for something much longer.
    captured = {}

    def fake_create_connection(address, timeout=None):
        captured["timeout"] = timeout
        return _FakeCreateConnectionSocket({"ok": True, "result": {}})

    monkeypatch.setattr(socket, "create_connection", fake_create_connection)
    client = DebugAgentClient(host="127.0.0.1", port=1)
    client.call("get_status", {})
    assert captured["timeout"] == 60


def test_client_call_honors_socket_timeout_override_for_long_waits(monkeypatch):
    # continue_execution can be asked to wait far longer than 60s (e.g. waiting for a
    # user to alt-tab). The client must be able to use a socket timeout derived from
    # that requested wait, not the fixed 60s default.
    captured = {}

    def fake_create_connection(address, timeout=None):
        captured["timeout"] = timeout
        return _FakeCreateConnectionSocket({"ok": True, "result": {}})

    monkeypatch.setattr(socket, "create_connection", fake_create_connection)
    client = DebugAgentClient(host="127.0.0.1", port=1)
    client.call("continue_execution", {"timeout": 300}, socket_timeout=330)
    assert captured["timeout"] == 330


class _RecordingDebugAgentClient:
    """Fake DebugAgentClient that just records how it was called, so tests can verify
    the continue_execution MCP tool derives the right socket_timeout without needing a
    real debug agent or a real socket."""

    def __init__(self):
        self.calls = []

    def call(self, op, params, socket_timeout=60):
        self.calls.append((op, params, socket_timeout))
        return {}


def test_continue_execution_tool_derives_socket_timeout_from_requested_wait():
    fake_client = _RecordingDebugAgentClient()
    server = build_mcp_server(fake_client)
    # FastMCP's @tool() decorator returns the wrapped function unchanged, so the
    # underlying callable is reachable straight off the tool manager for a direct,
    # synchronous call -- no need to go through the async MCP session machinery.
    continue_execution_fn = server._tool_manager.get_tool("continue_execution").fn

    continue_execution_fn(timeout=300)

    assert fake_client.calls == [("continue_execution", {"timeout": 300}, 330)]


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

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