havoc-remaster/tools/havoc_debug_agent_server.py

"""Local-only (127.0.0.1) JSON-over-TCP server wrapping a DebugAgent. Never bind 0.0.0.0
here -- the MCP server is the only LAN-reachable piece (see spec Security section)."""

from __future__ import annotations

import base64
import json
import queue
import socketserver
import threading

from tools.havoc_debug_agent import DebugAgent
from tools.havoc_debug_protocol import decode_message, WatchpointLimitError

_OPS = {}


def _op(name):
    def register(fn):
        _OPS[name] = fn
        return fn
    return register


@_op("launch")
def _op_launch(agent: DebugAgent, params: dict) -> dict:
    agent.launch(params["path"], params["cwd"])
    return {}


@_op("launch_with_breakpoints")
def _op_launch_with_breakpoints(agent: DebugAgent, params: dict) -> dict:
    """Launch and immediately plant breakpoints before the WOW64 loader hardens page protection."""
    agent.launch(params["path"], params["cwd"])
    for addr_hex in params.get("breakpoints", []):
        agent.set_breakpoint(int(addr_hex, 16), None)
    return {}


@_op("set_breakpoint")
def _op_set_breakpoint(agent: DebugAgent, params: dict) -> dict:
    agent.set_breakpoint(int(params["addr"], 16), params.get("condition"))
    return {}


@_op("remove_breakpoint")
def _op_remove_breakpoint(agent: DebugAgent, params: dict) -> dict:
    agent.remove_breakpoint(int(params["addr"], 16))
    return {}


@_op("list_breakpoints")
def _op_list_breakpoints(agent: DebugAgent, params: dict) -> dict:
    return {"breakpoints": [hex(addr) for addr in agent.list_breakpoints()]}


@_op("set_watchpoint")
def _op_set_watchpoint(agent: DebugAgent, params: dict) -> dict:
    slot = agent.set_watchpoint(int(params["addr"], 16), params["size"], params["mode"])
    return {"slot": slot}


@_op("remove_watchpoint")
def _op_remove_watchpoint(agent: DebugAgent, params: dict) -> dict:
    agent.remove_watchpoint(params["slot"])
    return {}


@_op("list_watchpoints")
def _op_list_watchpoints(agent: DebugAgent, params: dict) -> dict:
    return {"watchpoints": {slot: hex(addr) for slot, addr in agent.list_watchpoints().items()}}


@_op("continue_execution")
def _op_continue_execution(agent: DebugAgent, params: dict) -> dict:
    return agent.continue_execution(params["timeout"])


@_op("read_memory")
def _op_read_memory(agent: DebugAgent, params: dict) -> dict:
    data = agent.read_memory(int(params["addr"], 16), params["length"])
    return {"data": base64.b64encode(data).decode("ascii")}


@_op("write_memory")
def _op_write_memory(agent: DebugAgent, params: dict) -> dict:
    agent.write_memory(int(params["addr"], 16), base64.b64decode(params["data"]))
    return {}


@_op("get_registers")
def _op_get_registers(agent: DebugAgent, params: dict) -> dict:
    return {"registers": {name: hex(value) for name, value in agent.get_registers().items()}}


@_op("screenshot")
def _op_screenshot(agent: DebugAgent, params: dict) -> dict:
    png_bytes = agent.screenshot()
    return {"png_base64": base64.b64encode(png_bytes).decode("ascii")}


@_op("get_status")
def _op_get_status(agent: DebugAgent, params: dict) -> dict:
    return agent.get_status()


class _DebugWorker(threading.Thread):
    """Single dedicated thread that owns ALL debug agent calls.
    Windows requires CreateProcessA, WaitForDebugEvent, ContinueDebugEvent, etc. to
    all be called from the same OS thread. ThreadingTCPServer violates this by using a
    fresh thread per connection. This worker serializes every op on one thread via a queue."""

    def __init__(self, agent: DebugAgent) -> None:
        super().__init__(daemon=True, name="debug-worker")
        self._agent = agent
        self._q: queue.Queue = queue.Queue()

    def run(self) -> None:
        while True:
            op_name, params, result_q = self._q.get()
            try:
                if op_name not in _OPS:
                    raise ValueError(f"unknown op: {op_name!r}")
                result = _OPS[op_name](self._agent, params)
                result_q.put(("ok", result))
            except WatchpointLimitError as exc:
                result_q.put(("err", str(exc)))
            except Exception as exc:  # noqa: BLE001
                result_q.put(("err", f"{type(exc).__name__}: {exc}"))

    def call(self, op_name: str, params: dict, timeout: float = 120.0):
        result_q: queue.Queue = queue.Queue()
        self._q.put((op_name, params, result_q))
        return result_q.get(timeout=timeout)


class _Handler(socketserver.StreamRequestHandler):
    def handle(self) -> None:
        worker: _DebugWorker = self.server.worker  # type: ignore[attr-defined]
        for line in self.rfile:
            if not line.strip():
                continue
            request = decode_message(line)
            op_name = request["op"]
            params = request.get("params", {})
            # continue_execution can block for its full timeout plus processing overhead
            op_timeout = params.get("timeout", 30) + 10 if op_name == "continue_execution" else 30.0
            status, payload = worker.call(op_name, params, timeout=op_timeout)
            if status == "ok":
                response = {"ok": True, "result": payload}
            else:
                response = {"ok": False, "error": payload}
            self.wfile.write((json.dumps(response) + "\n").encode("utf-8"))


class _Server(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True


def run_agent_server(agent: DebugAgent, port: int) -> None:
    """Blocks forever, serving the local-only debug agent API on 127.0.0.1:port."""
    worker = _DebugWorker(agent)
    worker.start()
    server = _Server(("127.0.0.1", port), _Handler)
    server.worker = worker  # type: ignore[attr-defined]
    server.serve_forever()


if __name__ == "__main__":
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 47474
    run_agent_server(DebugAgent(), port)