havoc-remaster/tests/test_havoc_mcp_server.py

40 lines
1.5 KiB
Python

import json
import socket
import threading
import pytest
from tools.havoc_mcp_server import DebugAgentClient, DebugAgentError
class _StubAgentServer:
"""A minimal in-process stand-in for tools/havoc_debug_agent_server.py, so this
test never touches ctypes or a real Windows process."""
def __init__(self, response: dict):
self._response = response
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._sock.bind(("127.0.0.1", 0))
self._sock.listen(1)
self.port = self._sock.getsockname()[1]
self._thread = threading.Thread(target=self._serve_one, daemon=True)
self._thread.start()
def _serve_one(self):
conn, _ = self._sock.accept()
conn.recv(4096)
conn.sendall((json.dumps(self._response) + "\n").encode("utf-8"))
conn.close()
def test_client_returns_result_on_success():
stub = _StubAgentServer({"ok": True, "result": {"status": "running"}})
client = DebugAgentClient(host="127.0.0.1", port=stub.port)
result = client.call("get_status", {})
assert result == {"status": "running"}
def test_client_raises_on_error_response():
stub = _StubAgentServer({"ok": False, "error": "all 4 hardware debug registers (DR0-DR3) are in use"})
client = DebugAgentClient(host="127.0.0.1", port=stub.port)
with pytest.raises(DebugAgentError, match="DR0-DR3"):
client.call("set_watchpoint", {"addr": "0x5000", "size": 4, "mode": "write"})