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):
|
if not header_value.startswith(prefix):
|
||||||
return False
|
return False
|
||||||
return header_value[len(prefix):] == expected
|
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