55 lines
2 KiB
Python
55 lines
2 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"})
|
|
|
|
|
|
from tools.havoc_mcp_server import check_bearer_token
|
|
|
|
def test_check_bearer_token_accepts_matching_token():
|
|
assert check_bearer_token("Bearer secret123", expected="secret123") is True
|
|
|
|
def test_check_bearer_token_rejects_wrong_token():
|
|
assert check_bearer_token("Bearer wrong", expected="secret123") is False
|
|
|
|
def test_check_bearer_token_rejects_missing_header():
|
|
assert check_bearer_token(None, expected="secret123") is False
|
|
|
|
def test_check_bearer_token_rejects_malformed_header():
|
|
assert check_bearer_token("secret123", expected="secret123") is False # missing "Bearer " prefix
|