From 0388488c7464177038859130e19832b4f9e023f7 Mon Sep 17 00:00:00 2001 From: pyr0ball Date: Sat, 4 Jul 2026 07:13:36 -0700 Subject: [PATCH] feat: wire MCP tool definitions and bearer-token auth into server Adds requirements.txt (mcp==1.28.1, pytest==9.1.1) since no Python dependency-management convention existed in this repo yet. --- requirements.txt | 5 ++ tools/havoc_mcp_server.py | 124 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100644 requirements.txt diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..dc2414e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +# Runtime-observation MCP tooling (tools/havoc_mcp_server.py, tools/havoc_debug_protocol.py). +# No existing Python dependency-management convention existed in this repo (tools/ is a +# collection of standalone RE scripts); this file was introduced for Task 5. +mcp==1.28.1 +pytest==9.1.1 diff --git a/tools/havoc_mcp_server.py b/tools/havoc_mcp_server.py index ea3e069..7e71733 100644 --- a/tools/havoc_mcp_server.py +++ b/tools/havoc_mcp_server.py @@ -43,3 +43,127 @@ def check_bearer_token(header_value: str | None, expected: str) -> bool: if not header_value.startswith(prefix): return False return header_value[len(prefix):] == expected + + +# --- MCP tool wiring (mcp SDK 1.28.1) -------------------------------------------- +# +# FastMCP.tool() is a plain decorator: `@mcp.tool()` registers the wrapped function, +# using its name, type hints, and docstring to build the tool's JSON schema -- this +# matches the brief's example code as-is under the installed SDK version. +# +# FastMCP's own built-in auth (`mcp.server.auth`, `AuthSettings`, `TokenVerifier`) is a +# full OAuth2-resource-server flow (issuer URLs, scopes, token introspection) meant for +# public/multi-tenant deployments. That is the wrong shape for this tool: a single +# shared secret guarding one LAN-only debug endpoint. So `check_bearer_token` is instead +# wired in below `FastMCP`, as a raw ASGI middleware placed in front of the Starlette +# app that `FastMCP.streamable_http_app()` returns -- every HTTP request is checked +# there before it ever reaches the MCP session manager. + +from mcp.server.fastmcp import FastMCP + + +def build_mcp_server(client: DebugAgentClient, token: str) -> FastMCP: + mcp = FastMCP("havoc-runtime-observation") + + @mcp.tool() + def launch(path: str, cwd: str) -> dict: + return client.call("launch", {"path": path, "cwd": cwd}) + + @mcp.tool() + def set_breakpoint(addr: str, condition: dict | None = None) -> dict: + return client.call("set_breakpoint", {"addr": addr, "condition": condition}) + + @mcp.tool() + def remove_breakpoint(addr: str) -> dict: + return client.call("remove_breakpoint", {"addr": addr}) + + @mcp.tool() + def list_breakpoints() -> dict: + return client.call("list_breakpoints", {}) + + @mcp.tool() + def set_watchpoint(addr: str, size: int, mode: str) -> dict: + return client.call("set_watchpoint", {"addr": addr, "size": size, "mode": mode}) + + @mcp.tool() + def remove_watchpoint(slot: int) -> dict: + return client.call("remove_watchpoint", {"slot": slot}) + + @mcp.tool() + def list_watchpoints() -> dict: + return client.call("list_watchpoints", {}) + + @mcp.tool() + def continue_execution(timeout: float) -> dict: + return client.call("continue_execution", {"timeout": timeout}) + + @mcp.tool() + def read_memory(addr: str, length: int) -> dict: + return client.call("read_memory", {"addr": addr, "length": length}) + + @mcp.tool() + def write_memory(addr: str, data_base64: str) -> dict: + return client.call("write_memory", {"addr": addr, "data": data_base64}) + + @mcp.tool() + def get_registers() -> dict: + return client.call("get_registers", {}) + + @mcp.tool() + def screenshot() -> dict: + return client.call("screenshot", {}) + + @mcp.tool() + def get_status() -> dict: + return client.call("get_status", {}) + + return mcp + + +class _BearerAuthMiddleware: + """Raw ASGI middleware that enforces `check_bearer_token` on every HTTP request + before it reaches the wrapped app. Non-HTTP scopes (e.g. "lifespan") pass through + untouched. This is the actual auth enforcement point for the running server -- + the built-in `mcp.server.auth` OAuth stack is deliberately not used here (see the + module-level note above `build_mcp_server`).""" + + def __init__(self, app, expected_token: str) -> None: + self._app = app + self._expected_token = expected_token + + async def __call__(self, scope, receive, send) -> None: + if scope["type"] != "http": + await self._app(scope, receive, send) + return + + header_value = None + for name, value in scope.get("headers", []): + if name.lower() == b"authorization": + header_value = value.decode("latin-1") + 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 + + await self._app(scope, receive, send) + + +def build_asgi_app(server: FastMCP, token: str): + """Wrap the FastMCP Streamable HTTP app with bearer-token enforcement. This is + what actually gets served -- see the __main__ block below.""" + return _BearerAuthMiddleware(server.streamable_http_app(), expected_token=token) + + +if __name__ == "__main__": + import uvicorn + + token = os.environ["HAVOC_MCP_TOKEN"] # required, no default -- fail loudly if unset + client = DebugAgentClient(host="127.0.0.1", port=47474) + server = build_mcp_server(client, token) + 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)