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.
This commit is contained in:
parent
b2887753e9
commit
0388488c74
2 changed files with 129 additions and 0 deletions
5
requirements.txt
Normal file
5
requirements.txt
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue