Turn lift_bfs into a reusable function

2026-06-21 13:43:42 +00:00 · 2026-05-14 04:25:49 +02:00
parent 15a95377c2
commit 3c671c5cc2
7 changed files with 150 additions and 56 deletions
@@ -0,0 +1,37 @@
+from queue import Queue
+from llvm import Module, Value
+from striga import Semantics, Successor
+from container import Container
+
+
+def lift_bfs(module: Module, container: Container, start: int, *, verbose=True) -> Semantics:
+    sem = Semantics(module, verbose=verbose)
+    sem.begin(start)
+
+    queue: Queue[Successor] = Queue()
+    queue.put(Successor(0, sem.const64(start)))
+    # Keep destinations as LLVM Values instead of splitting constants into ints.
+    # This keeps the worklist uniform and matches later slicing/data-flow uses.
+    visited: set[Value] = set()
+    while not queue.empty():
+        src, dst = queue.get()
+
+        if not dst.is_constant:
+            if sem.verbose:
+                print(f"; non-constant branch destination: {hex(src)} -> {dst}")
+            continue
+
+        if dst in visited:
+            continue
+        visited.add(dst)
+
+        va = dst.const_zext_value
+        code = container.get_data(va, 15)
+        successors = sem.lift_bytes(va, code)
+        for successor in successors:
+            if successor.dst in visited:
+                continue
+            queue.put(successor)
+
+    sem.module.verify_or_raise()
+    return sem
@@ -1,9 +1,9 @@
 from pefile import PE
 from llvm import Module, create_context
-from lift import lift
+from lift import lift_pe


-def lift_crackme(module: Module, pe: PE):
+def lift_crackme(module: Module, filename: str):
    vm_handlers = [
        0x140016000,
        0x14001604E,
@@ -71,7 +71,7 @@ def lift_crackme(module: Module, pe: PE):
        0x14001676B,
    ]
    for handler in vm_handlers:
-        lifted = lift(module, pe, handler, verbose=False)
+        lifted = lift_pe(module, filename, handler, verbose=False)
        print(lifted.name)

    with open("tests/binaryshield.ll", "w") as f:
@@ -81,4 +81,4 @@ def lift_crackme(module: Module, pe: PE):
 if __name__ == "__main__":
    with create_context() as context:
        with context.create_module("binaryshield") as module:
-            lift_crackme(module, PE("tests/binaryshield.exe"))
+            lift_crackme(module, "tests/binaryshield.exe")
@@ -0,0 +1,32 @@
+from striga import Semantics, Successor
+from llvm import global_context, Module, Value
+from container import RawContainer
+from bfs import lift_bfs
+
+CODE = RawContainer(bytes.fromhex("48 C1 EF 03 48 8D 04 7F 48 01 F0 48 05 39 05 00 00 C3"), 0x1000)
+
+
+with global_context().create_module("blog") as module:
+    sem = lift_bfs(module, CODE, 0x1000)
+
+    types = module.context.types
+
+    ram = module.add_global(types.array(types.i8, 0), "RAM")
+
+    i64 = types.i64
+    lift2_ty = types.function(i64, [i64, i64])
+    lift2 = module.add_function("lift2", lift2_ty)
+    entry = lift2.append_basic_block("entry")
+    with entry.create_builder() as ir:
+        state = ir.alloca(sem.state_ty, "state")
+        reg_ptr = lambda name: ir.struct_gep(sem.state_ty, state, sem.reg_indices[name], name)
+        ir.store(lift2.get_param(0), reg_ptr("rdi"))
+        ir.store(lift2.get_param(1), reg_ptr("rsi"))
+
+        stack = ir.alloca(types.i8, i64.constant(4096), "stack")
+        stack_ptr = ir.gep(types.i8, stack, [i64.constant(4096 - 8)])
+        ir.store(stack_ptr, reg_ptr("rsp"))
+
+        ir.call(sem.function, [ram, state])
+        ir.ret(ir.load(i64, reg_ptr("rax")))
+    print(module)
@@ -0,0 +1,59 @@
+from pefile import PE
+
+
+class Container:
+    @property
+    def image_base(self) -> int:
+        raise NotImplementedError()
+
+    @property
+    def image_size(self) -> int:
+        raise NotImplementedError()
+
+    def in_range(self, va: int):
+        return va >= self.image_base and va < self.image_base + self.image_size
+
+    def get_data(self, va: int, size: int) -> bytes:
+        raise NotImplementedError()
+
+
+class RawContainer(Container):
+    def __init__(self, data: bytes, image_base=0):
+        self.data = data
+        self._image_base = image_base
+        self._image_size = len(data)
+
+    @property
+    def image_base(self) -> int:
+        return self._image_base
+
+    @property
+    def image_size(self) -> int:
+        return self._image_size
+
+    def get_data(self, va: int, size: int) -> bytes:
+        assert self.in_range(va)
+        rva = va - self.image_base
+        return self.data[rva : rva + size]
+
+
+class PEContainer(Container):
+    def __init__(self, pe: str | PE) -> None:
+        if isinstance(pe, str):
+            pe = PE(pe)
+        self.pe = pe
+        self._image_base = pe.OPTIONAL_HEADER.ImageBase  # pyright: ignore
+        self._image_size = pe.OPTIONAL_HEADER.SizeOfImage  # pyright: ignore
+
+    @property
+    def image_base(self) -> int:
+        return self._image_base
+
+    @property
+    def image_size(self) -> int:
+        return self._image_size
+
+    def get_data(self, va: int, size: int) -> bytes:
+        assert self.in_range(va)
+        rva = va - self.image_base
+        return self.pe.get_data(rva, size)
@@ -1,55 +1,20 @@
-from queue import Queue
-
-from striga import Semantics, Successor
-from pefile import PE
-from llvm import create_context, Value, Module
+from llvm import create_context, Module, Function
+from container import PEContainer
+from bfs import lift_bfs


-def lift(module: Module, pe: PE, start: int, *, verbose=True):
-    image_base = pe.OPTIONAL_HEADER.ImageBase  # pyright: ignore[reportOptionalMemberAccess, reportAttributeAccessIssue]
-    image_size = pe.OPTIONAL_HEADER.SizeOfImage  # pyright: ignore[reportOptionalMemberAccess, reportAttributeAccessIssue]
-    sem = Semantics(module, verbose=verbose)
-    lifted_fn = sem.begin(start)
-
-    queue: Queue[Successor] = Queue()
-    queue.put(Successor(0, sem.const64(start)))
-    # Keep destinations as LLVM Values instead of splitting constants into ints.
-    # This keeps the worklist uniform and matches later slicing/data-flow uses.
-    visited: set[Value] = set()
-    while not queue.empty():
-        src, dst = queue.get()
-
-        if not dst.is_constant:
-            if sem.verbose:
-                print(f"; non-constant branch destination: {hex(src)} -> {dst}")
-            # TODO: recover jump tables / returned-to callers
-            continue
-
-        if dst in visited:
-            continue
-        visited.add(dst)
-
-        va = dst.const_zext_value
-        assert va >= image_base and va < image_base + image_size
-        code = pe.get_data(va - image_base, 15)
-        successors = sem.lift_bytes(va, code)
-        for successor in successors:
-            if successor.dst in visited:
-                continue
-            queue.put(successor)
-
-    sem.module.verify_or_raise()
-    return lifted_fn
+def lift_pe(module: Module, filename: str, start: int, *, verbose=False) -> Function:
+    return lift_bfs(module, PEContainer(filename), start, verbose=verbose).function


 if __name__ == "__main__":
    with create_context() as context:
        with context.create_module("lifted") as module:
-            vm_entry = lift(module, PE("tests/binaryshield.exe"), 0x140017A41)
+            vm_entry = lift_pe(module, "tests/binaryshield.exe", 0x140017A41)
            print(vm_entry)
-            cfg = lift(module, PE("tests/cfg.exe"), 0x140001000)
+            cfg = lift_pe(module, "tests/cfg.exe", 0x140001000)
            print(cfg)
-            riscvm_run = lift(module, PE("tests/riscvm.exe"), 0x140001104)
+            riscvm_run = lift_pe(module, "tests/riscvm.exe", 0x140001104)
            print(riscvm_run)
-            themida = lift(module, PE("tests/example2-virt.bin"), 0x140001000)
+            themida = lift_pe(module, "tests/example2-virt.bin", 0x140001000)
            print(themida)
@@ -5,7 +5,7 @@ requires-python = ">=3.12"
 dependencies = [
    "capstone>=5.0.7",
    "icicle-emu>=0.0.11",
-    "llvm-nanobind",
+    "llvm-nanobind==21.1.6.3",
    "pefile>=2024.8.26",
 ]

@@ -35,14 +35,15 @@ wheels = [

 [[package]]
 name = "llvm-nanobind"
-version = "21.1.6.1"
+version = "21.1.6.3"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/66/b5/2255a7b42ddd154ea3247841f0d6610b634a376ad2e20cc49b91c8d83c22/llvm_nanobind-21.1.6.1.tar.gz", hash = "sha256:b46ba66cabf6f47b4fb988ee77ebf0fb9a52a94df7396400e27dad09c26d05e6", size = 633640, upload-time = "2026-05-12T16:22:33.865Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/78/6e/8501bb26482f231bc22815f4dde05fb05e0c52474cf5c455a211ed847474/llvm_nanobind-21.1.6.3.tar.gz", hash = "sha256:0629ce80418a94b2e1a0019d66372cba86e1327b6c9fd16d63a10aa3c09204e2", size = 645135, upload-time = "2026-05-14T01:06:01.871Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/9d/b0/e776e058669e8037b30b7327c7c243be626b67c3e12365be50b5a18846c3/llvm_nanobind-21.1.6.1-cp312-abi3-macosx_13_0_arm64.whl", hash = "sha256:e4263b5570e6877bb2d417c9f7dd10c10107a9f3b9f8847ec31da76a970307da", size = 43911001, upload-time = "2026-05-12T16:22:18.594Z" },
-    { url = "https://files.pythonhosted.org/packages/56/2c/660005a35afea084d76e9bbfaf4c3326de457dbfa7f630fa2a2667839d04/llvm_nanobind-21.1.6.1-cp312-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:5782bd8172cdfd0c0ba8b175137ee860a9fa9208ee9b652b6cfd2def5f6aacf7", size = 58457985, upload-time = "2026-05-12T16:22:22.491Z" },
-    { url = "https://files.pythonhosted.org/packages/9d/55/deb6aa6c2f7b2387541385443b443823cab66b367beb131a7a9567f3c0b4/llvm_nanobind-21.1.6.1-cp312-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:fc45990f16dda0dadbae13e013c069232780ad0131736ccaaaf1f8eb9f66d83e", size = 61753025, upload-time = "2026-05-12T16:22:26.942Z" },
-    { url = "https://files.pythonhosted.org/packages/36/2f/26d6839abd9a112132abc0dfb790b2c76e9d169959ff1580b9ee8c355264/llvm_nanobind-21.1.6.1-cp312-abi3-win_amd64.whl", hash = "sha256:3c18170be22f41b0552c0032ecd863b926e1132ba2f89af8c34a118df2abdc78", size = 44011407, upload-time = "2026-05-12T16:22:31.123Z" },
+    { url = "https://files.pythonhosted.org/packages/b7/03/3ffd26b70fe0db6f6be411b57f396c81a1ee01fcefea6392707192bcdba4/llvm_nanobind-21.1.6.3-cp312-abi3-macosx_13_0_arm64.whl", hash = "sha256:b153bfe4f2008f96dd0395f72e1e4413f8e5d2d9dc254cd51ae49cc48e7f84e9", size = 43962841, upload-time = "2026-05-14T01:05:32.964Z" },
+    { url = "https://files.pythonhosted.org/packages/9e/af/77b62c2dbe20ba8425e1d358da40436599b5540357145a76349098bdb33a/llvm_nanobind-21.1.6.3-cp312-abi3-macosx_13_0_x86_64.whl", hash = "sha256:fa46ac4e99698190eda9b4e6760d21618bac7c814997da63039e959ca7f9e8fa", size = 47013114, upload-time = "2026-05-14T01:05:38.565Z" },
+    { url = "https://files.pythonhosted.org/packages/4d/81/ca1b5ac53c272dcf1f4db98c0b2246c0a532f4eea42193a05e3ded90099b/llvm_nanobind-21.1.6.3-cp312-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:4e4e814a68c7a63c62badc845d0d1941b04a4cdcf11eafc9398e08c3e3a1bc40", size = 58501459, upload-time = "2026-05-14T01:05:44.558Z" },
+    { url = "https://files.pythonhosted.org/packages/1a/85/7830240b5d1bac8bc69c0b039478d2342f2809c0809825f0e3f157571f1d/llvm_nanobind-21.1.6.3-cp312-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:8be86943cd9c7a7e69ac07ee1f948d7971089377831b0c180de9c5ba714f97d7", size = 61796024, upload-time = "2026-05-14T01:05:52.125Z" },
+    { url = "https://files.pythonhosted.org/packages/e9/94/6bce36ab093677ba6643c6a8b9d51bb9893be3c4d9e4604bf797e93077a2/llvm_nanobind-21.1.6.3-cp312-abi3-win_amd64.whl", hash = "sha256:113a38f38201f7af137b3d7e19e1de9571ed7c7e648cb9f184e26471a1457193", size = 44061105, upload-time = "2026-05-14T01:05:58.011Z" },
 ]

 [[package]]
@@ -100,7 +101,7 @@ dev = [
 requires-dist = [
    { name = "capstone", specifier = ">=5.0.7" },
    { name = "icicle-emu", specifier = ">=0.0.11" },
-    { name = "llvm-nanobind" },
+    { name = "llvm-nanobind", specifier = "==21.1.6.3" },
    { name = "pefile", specifier = ">=2024.8.26" },
 ]