Fix in place binary rebuild

themida_unmutate/miasm_utils.py

@@ -2,10 +2,11 @@
 from typing import Optional, Self
 
 import miasm.expression.expression as m2_expr
+import miasm.core.asmblock as m2_asmblock
 from miasm.analysis.binary import Container
 from miasm.analysis.machine import Machine
 from miasm.core import parse_asm
-from miasm.core.asmblock import disasmEngine, asm_resolve_final
+from miasm.core.asmblock import disasmEngine, AsmCFG, asm_resolve_final
 from miasm.core.interval import interval
 from miasm.core.locationdb import LocationDB
 from miasm.ir.ir import Lifter
@@ -78,3 +79,53 @@ def generate_code_redirect_patch(miasm_ctx: MiasmContext, src_addr: int, dst_add
     assert jmp_patch is not None
 
     return jmp_patch
+
+
+# Custom version of miasm's `asm_resolve_final` which works better for our
+# in-place rewriting use case
+def asm_resolve_final_in_place(loc_db: LocationDB, mnemo, asmcfg: AsmCFG, dst_interval: Optional[interval] = None):
+    """Resolve and assemble @asmcfg into interval
+    @dst_interval"""
+
+    asmcfg.sanity_check()
+
+    guess_blocks_size(loc_db, asmcfg, dst_interval)
+    blockChains = m2_asmblock.group_constrained_blocks(asmcfg)
+    resolved_blockChains = m2_asmblock.resolve_symbol(blockChains, asmcfg.loc_db, dst_interval)
+    m2_asmblock.asmblock_final(mnemo, asmcfg, resolved_blockChains)
+    patches = {}
+    output_interval = interval()
+
+    for block in asmcfg.blocks:
+        offset = asmcfg.loc_db.get_location_offset(block.loc_key)
+        for instr in block.lines:
+            if not instr.data:
+                # Empty line
+                continue
+            assert len(instr.data) == instr.l
+            patches[offset] = instr.data
+            instruction_interval = interval([(offset, offset + instr.l - 1)])
+            if not (instruction_interval & output_interval).empty:
+                raise RuntimeError("overlapping bytes %X" % int(offset))
+            output_interval = output_interval.union(instruction_interval)
+            instr.offset = offset
+            offset += instr.l
+    return patches
+
+
+# This is a custom version of miasm's `guess_blocks_size` which simply reuse
+# data from the original CFG's interval to compute true basic block sizes so
+# that it properly fits inside of the destination interval
+def guess_blocks_size(loc_db: LocationDB, asmcfg: AsmCFG, interval: interval):
+    # Compute bb sizes from interval
+    for block in asmcfg.blocks:
+        for start, end in interval.intervals:
+            if start == loc_db.get_location_offset(block.loc_key):
+                block.max_size = end - start
+                block.size = block.max_size
+
+        # Set up instructions
+        for instr in block.lines:
+            if instr.b is None:
+                instr.b = b""
+                instr.l = 0

themida_unmutate/rebuilding.py

@@ -4,7 +4,8 @@
 from miasm.core.asmblock import AsmCFG, asm_resolve_final, bbl_simplifier
 from miasm.core.interval import interval
 
-from themida_unmutate.miasm_utils import MiasmContext, MiasmFunctionInterval, generate_code_redirect_patch
+from themida_unmutate.miasm_utils import (MiasmContext, MiasmFunctionInterval, generate_code_redirect_patch,
+                                          asm_resolve_final_in_place)
 
 NEW_SECTION_NAME = ".unmut"
 NEW_SECTION_MAX_SIZE = 2**16
@@ -148,36 +149,23 @@ def __rebuild_simplified_binary_in_place(
             func_addr_to_simplified_cfg.items():
         simplified_asmcfg, orignal_asmcfg_interval = val
 
-        # Unpin blocks to be able to relocate the CFG
-        head = simplified_asmcfg.heads()[0]
-        for asm_block in simplified_asmcfg.blocks:
-            miasm_ctx.loc_db.unset_location_offset(asm_block.loc_key)
-
-        # Start rewriting at the first part of the interval (i.e., at the start
-        # of the mutated code)
-        target_addr: int = orignal_asmcfg_interval.intervals[0][0]
-        # Unpin loc_key if it's pinned
-        original_loc = miasm_ctx.loc_db.get_offset_location(target_addr)
-        if original_loc is not None:
-            miasm_ctx.loc_db.unset_location_offset(original_loc)
-
-        # Relocate the function's entry block
-        miasm_ctx.loc_db.set_location_offset(head, target_addr)
-
         # Generate the simplified machine code
-        new_section_patches = asm_resolve_final(miasm_ctx.mdis.arch,
-                                                simplified_asmcfg,
-                                                dst_interval=orignal_asmcfg_interval)
+        new_section_patches = asm_resolve_final_in_place(miasm_ctx.loc_db,
+                                                         miasm_ctx.mdis.arch,
+                                                         simplified_asmcfg,
+                                                         dst_interval=orignal_asmcfg_interval)
 
         # Merge patches into the patch list
         for patch in new_section_patches.items():
             unmut_patches.append(patch)
 
         # Associate original addr to simplified addr
-        original_to_simplified[protected_func_addr] = min(new_section_patches.keys())
+        head = simplified_asmcfg.heads()[0]
+        head_addr = miasm_ctx.loc_db.get_location_offset(head)
+        original_to_simplified[protected_func_addr] = head_addr
 
     # Find Themida's section
-    themida_section = __section_from_virtual_address(pe_obj, target_addr)
+    themida_section = __section_from_virtual_address(pe_obj, head_addr)
     assert themida_section is not None
 
     # Overwrite Themida's section content