pyr0ball
6baefe8e8a
Adds DebugAgent.screenshot() using Pillow/ImageGrab -- locates the game window
via FindWindowA, grabs its rect, returns PNG bytes. Adds tools/havoc_debug_agent_server.py:
a 127.0.0.1-only ThreadingTCPServer that serializes all ops behind a single lock and
dispatches NDJSON requests to DebugAgent. Covers all debug-agent ops (launch,
breakpoints, watchpoints, continue_execution, read/write_memory, get_registers,
screenshot, get_status). Acceptance test confirmed: get_status returns
{"ok": true, "result": {"status": "not_launched"}} from a fresh server. Adds
tools/requirements-windows.txt (pillow).
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TN4Ytn3gdWRonNmHpisWQv
136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
"""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)
|