diff --git a/tools/havoc_debug_agent.py b/tools/havoc_debug_agent.py index d057dd2..e48386c 100644 --- a/tools/havoc_debug_agent.py +++ b/tools/havoc_debug_agent.py @@ -318,6 +318,21 @@ class DebugAgent: return None + def screenshot(self) -> bytes: + """Capture the game window as PNG bytes. Requires Pillow (pip install pillow).""" + from PIL import ImageGrab # local import: only needed on Caspian, not for pure-logic tests + import io + + hwnd = ctypes.windll.user32.FindWindowA(None, b"HAVOC(tm) by Reality Bytes") + if not hwnd: + raise OSError("game window not found -- is HAVOC_NOCD.EXE running and visible?") + rect = wintypes.RECT() + ctypes.windll.user32.GetWindowRect(hwnd, byref(rect)) + img = ImageGrab.grab(bbox=(rect.left, rect.top, rect.right, rect.bottom)) + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + def get_status(self) -> dict: if self._exited is not None: return {"status": "exited", "exit_code": self._exited} diff --git a/tools/havoc_debug_agent_server.py b/tools/havoc_debug_agent_server.py new file mode 100644 index 0000000..6e4e79e --- /dev/null +++ b/tools/havoc_debug_agent_server.py @@ -0,0 +1,136 @@ +"""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 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("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 _Handler(socketserver.StreamRequestHandler): + def handle(self) -> None: + agent: DebugAgent = self.server.agent # type: ignore[attr-defined] + lock: threading.Lock = self.server.lock # 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", {}) + try: + with lock: + if op_name not in _OPS: + raise ValueError(f"unknown op: {op_name!r}") + result = _OPS[op_name](agent, params) + response = {"ok": True, "result": result} + except WatchpointLimitError as exc: + response = {"ok": False, "error": str(exc)} + except Exception as exc: # noqa: BLE001 + response = {"ok": False, "error": f"{type(exc).__name__}: {exc}"} + 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.""" + server = _Server(("127.0.0.1", port), _Handler) + server.agent = agent # type: ignore[attr-defined] + server.lock = threading.Lock() # 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) diff --git a/tools/requirements-windows.txt b/tools/requirements-windows.txt new file mode 100644 index 0000000..3868fb1 --- /dev/null +++ b/tools/requirements-windows.txt @@ -0,0 +1 @@ +pillow