HarborFork/tests/test_reverse.py

"""The shared reverse-engineering lookups, against a small synthetic image and annotation set."""
import json
import shutil
import struct
import unittest

from helpers import temp_dir
from harbor.reverse import Reverse
from harbor.util import HarborError


class ReverseTools(unittest.TestCase):
    def setUp(self):
        self.d = temp_dir()
        base = 0x401000
        image = bytearray(0x100)
        # 0x401000: mov ecx,[0x401080]; call 0x401040; fmul dword [0x401090]; ret
        code = bytes.fromhex("8b0d80104000") + b"\xe8" + struct.pack("<i", 0x401040 - 0x40100B) \
            + bytes.fromhex("d80d90104000") + b"\xc3"
        image[0:len(code)] = code
        image[0x40] = 0xC3                                       # 0x401040: ret
        image[0x90:0x94] = struct.pack("<f", 100.0)              # 0x401090: 100.0f
        image[0xA0:0xA5] = b".cue\0"
        (self.d / "image.bin").write_bytes(bytes(image))
        facts = [{"address": "0x401000", "size": len(code), "name": "FUN_00401000", "calls": ["0x401040"],
                  "data": ["0x401080"], "imports": [], "strings": []},
                 {"address": "0x401040", "size": 1, "name": "FUN_00401040", "calls": [], "data": [],
                  "imports": [], "strings": []}]
        (self.d / "facts.jsonl").write_text("\n".join(json.dumps(f) for f in facts), encoding="utf-8")
        ann = self.d / "annotations"
        ann.mkdir()
        (ann / "names.csv").write_text("address,name,category,note\n0x00401000,GetLevel,game,\"Reads it.\"\n"
                                       "0x00401040,Helper,game,x\n", encoding="utf-8")
        (ann / "globals.csv").write_text("address,name,type,note\n0x00401080,g_root,int,x\n", encoding="utf-8")
        (ann / "functions.csv").write_text("address,signature,note\n", encoding="utf-8")
        (ann / "prototypes.h").write_text("int GetLevel(void);  /* plain ret */\n", encoding="utf-8")
        self.r = Reverse({"image": str(self.d / "image.bin"), "base": "0x401000", "facts": str(self.d / "facts.jsonl"),
                          "annotations": str(ann)})

    def tearDown(self):
        shutil.rmtree(self.d, ignore_errors=True)

    def test_disassembly_names_calls_globals_and_constants(self):
        text = self.r.disassemble("GetLevel")
        self.assertIn("0x401080 g_root", text)
        self.assertIn("0x401040 Helper", text)
        self.assertIn("= 100.0f", text)

    def test_xrefs_bytes_and_lookup(self):
        x = self.r.xrefs("0x401040")
        self.assertEqual([c["name"] for c in x["callers"]], ["GetLevel"])
        self.assertEqual(self.r.bytes_at("0x401090", 4, "float")["float"], ["100.0"])
        self.assertEqual(self.r.bytes_at("0x4010A0", as_type="string")["string"], ".cue")
        info = self.r.lookup("GetLevel")
        self.assertIn("plain ret", info["prototype"])
        with self.assertRaises(HarborError):
            self.r.bytes_at("0x500000")
        with self.assertRaises(HarborError):
            Reverse({})


class GhidraForwarder(unittest.TestCase):
    """The ghidra() tool against a fake streamable-HTTP MCP server: a session, a polled async task, a refusal."""

    def setUp(self):
        import http.server
        import threading
        calls = self.calls = []

        class Fake(http.server.BaseHTTPRequestHandler):
            def log_message(self, *a):
                pass

            def do_POST(self):
                body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
                calls.append((body["method"], body.get("params", {}).get("name"), self.headers.get("Mcp-Session-Id")))
                if "id" not in body:
                    self.send_response(202)
                    self.end_headers()
                    return
                name = body.get("params", {}).get("name")
                if body["method"] == "initialize":
                    result = {"protocolVersion": "2025-06-18"}
                elif body["method"] == "tools/list":
                    result = {"tools": [{"name": "get_code", "description": "Code",
                                         "inputSchema": {"properties": {"function": {}}, "required": ["function"]}}]}
                elif name == "get_code":
                    result = {"content": [{"type": "text", "text": "Task submitted for async execution.\n\n"
                                           "Task ID: 11111111-2222-3333-4444-555555555555\nStatus: RUNNING"}]}
                elif name == "get_task_status":
                    result = {"content": [{"type": "text", "text": "[Context] Operating on: image.bin\n\nint F(void)"}]}
                else:
                    result = {"isError": True, "content": [{"type": "text", "text": "Refused: read-only"}]}
                data = ("data: " + json.dumps({"jsonrpc": "2.0", "id": body["id"], "result": result}) + "\n\n").encode()
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Mcp-Session-Id", "s1")
                self.send_header("Content-Length", str(len(data)))
                self.end_headers()
                self.wfile.write(data)

        self.server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Fake)
        threading.Thread(target=self.server.serve_forever, daemon=True).start()
        from harbor.reverse import Ghidra
        self.g = Ghidra({"ghidra_url": f"http://127.0.0.1:{self.server.server_port}/mcp"})

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()

    def test_tools_call_and_refusal(self):
        self.assertEqual(self.g.tools()[0]["required"], ["function"])
        self.assertEqual(self.g.call("get_code", {"function": "F"}), "int F(void)")
        self.assertIn(("tools/call", "get_task_status", "s1"), self.calls)
        with self.assertRaises(HarborError):
            self.g.call("rename_symbol", {})
        with self.assertRaises(HarborError):
            from harbor.reverse import Ghidra
            Ghidra({"ghidra_url": "http://127.0.0.1:9/mcp"}).tools()