forked from Keelrleek/Harbor
59 lines
2.9 KiB
Python
59 lines
2.9 KiB
Python
"""next_functions, defer and progress against a tiny synthetic project."""
|
|
import json
|
|
import unittest
|
|
|
|
from helpers import HarborCase
|
|
from harbor.queue import FunctionQueue
|
|
from harbor.util import HarborError
|
|
|
|
|
|
class FunctionQueueTests(HarborCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
d = self.dir
|
|
image = bytearray(0x200)
|
|
image[0x00:0x10] = bytes.fromhex("8b44240485c07e03b801000000c3cccc") # 0x401000: stack argument
|
|
image[0x40:0x48] = bytes.fromhex("85c07e03b8010000") # 0x401040: reads EAX
|
|
image[0x48] = 0xC3
|
|
image[0x80:0x8C] = bytes.fromhex("8b442404c3cccccccccccccc") # 0x401080: readable already
|
|
(d / "image.bin").write_bytes(bytes(image))
|
|
facts = [{"address": "0x401000", "size": 14, "calls": []},
|
|
{"address": "0x401040", "size": 9, "calls": []},
|
|
{"address": "0x401080", "size": 5, "calls": []}]
|
|
(d / "facts.jsonl").write_text("\n".join(json.dumps(f) for f in facts), encoding="utf-8")
|
|
ann = d / "annotations"
|
|
ann.mkdir()
|
|
(ann / "names.csv").write_text("address,name,category,note\n0x00401000,First,game,x\n"
|
|
"0x00401040,Second,game,x\n0x00401080,Third,game,x\n", encoding="utf-8")
|
|
src = d / "src" / "game"
|
|
src.mkdir(parents=True)
|
|
(src / "401080_Third.c").write_text("/*\n * status: readable\n */\nint Third(int a) { return a; }\n",
|
|
encoding="utf-8")
|
|
(src / "401000_First.c").write_text("/*\n * status: raw\n */\n", encoding="utf-8")
|
|
self.h.project.reverse = {"image": str(d / "image.bin"), "base": "0x401000", "facts": str(d / "facts.jsonl"),
|
|
"annotations": str(ann), "src": str(d / "src")}
|
|
self.h.claim("Claude", "range:0x401000-0x401FFF", kind="band", on_behalf_of="DeepSeek")
|
|
|
|
def test_queue_skips_readable_and_deferred(self):
|
|
q = FunctionQueue(self.h)
|
|
q.min_size = 1
|
|
rows = q.next_functions("DeepSeek", 10)
|
|
self.assertEqual({r["address"] for r in rows}, {"0x401000", "0x401040"})
|
|
q.defer("DeepSeek", "0x401040", "register-args", "takes the value in EAX; prototypes.h row needed first")
|
|
self.assertEqual([r["address"] for r in q.next_functions("DeepSeek", 10)], ["0x401000"])
|
|
with self.assertRaises(HarborError):
|
|
q.defer("DeepSeek", "0x401000", "bored", "x")
|
|
with self.assertRaises(HarborError):
|
|
q.next_functions("Codex", 5) # no band
|
|
|
|
def test_progress_tiers(self):
|
|
q = FunctionQueue(self.h)
|
|
q.defer("DeepSeek", "0x401040", "x87", "precision differs")
|
|
tiers = q.progress()["tiers"]
|
|
self.assertEqual((tiers["readable"]["functions"], tiers["deferred"]["functions"], tiers["raw"]["functions"]),
|
|
(1, 1, 1))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|