havoc-remaster/tools/havoc_mcp_server.py

"""Thin MCP server translating MCP tool calls into calls to the local debug agent.
No Win32/ctypes dependency of its own -- runs and is tested on any platform, though at
runtime it must reach the real debug agent server (tools/havoc_debug_agent_server.py)
running on the same Caspian box."""

from __future__ import annotations

import hmac
import os
import socket

from starlette.responses import PlainTextResponse

from tools.havoc_debug_protocol import decode_message, encode_message


class DebugAgentError(Exception):
    """Raised when the debug agent reports {"ok": false}."""


class DebugAgentClient:
    def __init__(self, host: str, port: int) -> None:
        self._host = host
        self._port = port

    def call(self, op: str, params: dict, socket_timeout: float = 60) -> dict:
        """`socket_timeout` bounds both connect time and every subsequent recv() on
        this socket. The 60s default suits every op except continue_execution, which
        can legitimately be asked to wait far longer (e.g. for a user to alt-tab) --
        callers that need to wait longer than that must pass a larger override
        derived from the wait they actually requested."""
        with socket.create_connection((self._host, self._port), timeout=socket_timeout) as sock:
            sock.sendall(encode_message(op, params))
            buf = b""
            while not buf.endswith(b"\n"):
                chunk = sock.recv(65536)
                if not chunk:
                    break
                buf += chunk
        response = decode_message(buf)
        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 not expected:
        return False
    if header_value is None:
        return False
    prefix = "Bearer "
    if not header_value.startswith(prefix):
        return False
    return hmac.compare_digest(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) -> 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 launch_with_breakpoints(path: str, cwd: str, breakpoints: list) -> dict:
        """Launch the game and plant INT3 breakpoints atomically before the WOW64 loader
        hardens .text page protection. Use instead of launch() + set_breakpoint() when
        setting breakpoints at process start."""
        return client.call("launch_with_breakpoints",
                           {"path": path, "cwd": cwd, "breakpoints": breakpoints})

    @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:
        """Resume execution until a breakpoint/watchpoint hits or `timeout` seconds
        elapse. The debug agent is given `timeout + 10` seconds of headroom to
        actually wait that long, so the client socket needs a correspondingly longer
        timeout too -- otherwise a long wait (e.g. waiting for a user to alt-tab)
        would make this call's own recv() time out while the agent is still
        legitimately waiting."""
        return client.call(
            "continue_execution", {"timeout": timeout}, socket_timeout=timeout + 30
        )

    @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):
            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
    if not token:
        raise ValueError("HAVOC_MCP_TOKEN must not be empty")
    client = DebugAgentClient(host="127.0.0.1", port=47474)
    server = build_mcp_server(client)
    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)