Merge pull request #21 from purseclab/code-fixes

Code fixes

src/patcherex2/components/archinfo/aarch64.py

@@ -3,6 +3,8 @@ class Aarch64Info:
     nop_size = 4
     jmp_asm = "b {dst}"
     jmp_size = 4
+    alignment = 4
+    is_variable_length_isa = False
     call_asm = "bl {dst}"
     pc_reg_names = ["pc", "ip"]
     save_context_asm = """

src/patcherex2/components/archinfo/amd64.py

@@ -3,6 +3,8 @@ class Amd64Info:
     nop_size = 1
     jmp_asm = "jmp {dst}"
     jmp_size = 6
+    alignment = 4
+    is_variable_length_isa = True
     call_asm = "call {dst}"
     pc_reg_names = ["rip"]
     save_context_asm = """

src/patcherex2/components/archinfo/arm.py

@@ -3,6 +3,8 @@ class ArmInfo:
     nop_size = 4
     jmp_asm = "b {dst}"
     jmp_size = 4
+    alignment = 4
+    is_variable_length_isa = False
     call_asm = "bl {dst}"
     pc_reg_names = ["pc", "r15", "ip"]
     save_context_asm = """

src/patcherex2/components/archinfo/mips.py

@@ -4,6 +4,8 @@ class MipsInfo:
     jmp_asm = "j {dst}"
     # NOTE: keystone will always add nop for branch delay slot, so include it in size
     jmp_size = 8
+    alignment = 4
+    is_variable_length_isa = False
     call_asm = "jal {dst}"
     pc_reg_names = ["pc"]
     save_context_asm = """

src/patcherex2/components/archinfo/mips64.py

@@ -4,6 +4,8 @@ class Mips64Info:
     jmp_asm = "j {dst}"
     # NOTE: keystone will aldays add nop for branch delay slot, so include it in size
     jmp_size = 8
+    alignment = 4
+    is_variable_length_isa = False
     call_asm = "jal {dst}"
     pc_reg_names = ["pc"]
     save_context_asm = """

src/patcherex2/components/archinfo/ppc.py

@@ -3,6 +3,8 @@ class PpcInfo:
     nop_size = 4
     jmp_asm = "b {dst}"
     jmp_size = 4
+    alignment = 4
+    is_variable_length_isa = False
     call_asm = "bl {dst}"
     pc_reg_names = []
     save_context_asm = """

src/patcherex2/components/archinfo/ppc64.py

@@ -3,6 +3,8 @@ class Ppc64Info:
     nop_size = 4
     jmp_asm = "b {dst}"
     jmp_size = 4
+    alignment = 4
+    is_variable_length_isa = False
     call_asm = "bl {dst}"
     pc_reg_names = []
     save_context_asm = """

src/patcherex2/components/archinfo/ppc_vle.py

@@ -3,6 +3,8 @@ class PpcVleInfo:
     nop_size = 4
     jmp_asm = "b {dst}"
     jmp_size = 4
+    alignment = 4
+    is_variable_length_isa = True
     call_asm = "bl {dst}"
     pc_reg_names = []
     save_context_asm = ""  # TODO

src/patcherex2/components/archinfo/sparc.py

@@ -3,6 +3,8 @@ class SparcInfo:
     nop_size = 4
     jmp_asm = "b {dst}\nnop"  # nop due to delay slot
     jmp_size = 8
+    alignment = 4
+    is_variable_length_isa = False
     call_asm = "call {dst}"
     pc_reg_names = ["pc"]
     save_context_asm = ""  # TODO

src/patcherex2/components/archinfo/x86.py

@@ -3,6 +3,8 @@ class X86Info:
     nop_size = 1
     jmp_asm = "jmp {dst}"
     jmp_size = 5
+    alignment = 4
+    is_variable_length_isa = True
     call_asm = "call {dst}"
     pc_reg_names = ["eip"]
     save_context_asm = """

src/patcherex2/components/binary_analyzers/angr.py

@@ -86,10 +86,10 @@ def get_basic_block(self, addr: int) -> Dict[str, Union[int, List[int]]]:
             ],
         }
 
-    def get_instr_bytes_at(self, addr: int) -> angr.Block:
+    def get_instr_bytes_at(self, addr: int, num_instr=1) -> angr.Block:
         addr += 1 if self.is_thumb(addr) else 0
         addr = self.denormalize_addr(addr)
-        return self.p.factory.block(addr, num_inst=1).bytes
+        return self.p.factory.block(addr, num_inst=num_instr).bytes
 
     def get_unused_funcs(self) -> List[Dict[str, int]]:
         logger.info("Getting unused functions with angr")

src/patcherex2/components/utils/utils.py

@@ -63,8 +63,8 @@ def insert_trampoline_code(
         )
         if detour_pos == -1:
             trampoline_block = self.p.allocation_manager.allocate(
-                trampoline_size, align=0x4, flag=MemoryFlag.RX
-            )  # TODO: get alignment from arch info
+                trampoline_size, align=self.p.archinfo.alignment, flag=MemoryFlag.RX
+            )
             logger.debug(f"Allocated trampoline block: {trampoline_block}")
             mem_addr = trampoline_block.mem_addr
             file_addr = trampoline_block.file_addr

src/patcherex2/patcherex.py

@@ -1,6 +1,8 @@
+# ruff: noqa: F403, F405
 import logging
 
-from .patches import InsertDataPatch, ModifyDataPatch, RemoveDataPatch
+from .patches import *
+from .patches import __all__
 from .targets import Target
 
 logger = logging.getLogger(__name__)
@@ -46,13 +48,26 @@ def __init__(
                 ),
             )
 
-    def apply_patches(self):
-        # TODO: sort patches properly
-        self.patches.sort(
-            key=lambda x: not isinstance(
-                x, (ModifyDataPatch, InsertDataPatch, RemoveDataPatch)
-            )
+        # Chosen patch order, making sure all are accounted for
+        self.patch_order = (
+            ModifyRawBytesPatch,
+            RemoveDataPatch,
+            InsertDataPatch,
+            ModifyDataPatch,
+            RemoveLabelPatch,
+            ModifyLabelPatch,
+            InsertLabelPatch,
+            RemoveInstructionPatch,
+            InsertInstructionPatch,
+            ModifyInstructionPatch,
+            RemoveFunctionPatch,
+            InsertFunctionPatch,
+            ModifyFunctionPatch,
         )
+        assert len(self.patch_order) == len(__all__)
+
+    def apply_patches(self):
+        self.patches.sort(key=lambda x: self.patch_order.index(type(x)))
         logger.debug(f"Applying patches: {self.patches}")
         for patch in self.patches:
             patch.apply(self)

src/patcherex2/patches/function_patches.py

@@ -153,8 +153,8 @@ def apply(self, p) -> None:
             )
             if self.detour_pos == -1:
                 block = p.allocation_manager.allocate(
-                    compiled_size + 0x20, align=0x4, flag=MemoryFlag.RX
-                )  # TODO: get alignment from arch info, TODO: adjust that 0x20 part
+                    compiled_size + 0x20, align=p.archinfo.alignment, flag=MemoryFlag.RX
+                )  # TODO: adjust that 0x20 part
                 mem_addr = block.mem_addr
                 file_addr = block.file_addr
             else:

src/patcherex2/patches/instruction_patches.py

@@ -16,13 +16,28 @@ def __init__(
         self.symbols = symbols if symbols else {}
 
     def apply(self, p) -> None:
-        # TODO: check size, insert jump if necessary
         asm_bytes = p.assembler.assemble(
             self.instr,
             self.addr,
             symbols=self.symbols,
             is_thumb=p.binary_analyzer.is_thumb(self.addr),
         )
+        if p.archinfo.is_variable_length_isa:
+            asm_size = len(asm_bytes)
+            overwritten_size = 0
+            num_instrs = 1
+            while overwritten_size < asm_size:
+                overwritten_size = len(
+                    p.binary_analyzer.get_instr_bytes_at(
+                        self.addr, num_instr=num_instrs
+                    )
+                )
+                num_instrs += 1
+            remaining_size = overwritten_size - asm_size
+            assert (
+                remaining_size % p.archinfo.nop_size == 0
+            ), f"Cannot fill in {remaining_size} bytes when modifying instruction, must be a multiple of {p.archinfo.nop_size}"
+            asm_bytes += p.archinfo.nop_bytes * (remaining_size // p.archinfo.nop_size)
         offset = p.binary_analyzer.mem_addr_to_file_offset(self.addr)
         p.binfmt_tool.update_binary_content(offset, asm_bytes)
 
@@ -80,8 +95,8 @@ def apply(self, p) -> None:
             )
             if self.detour_pos == -1:
                 block = p.allocation_manager.allocate(
-                    assembled_size, align=0x4, flag=MemoryFlag.RX
-                )  # TODO: get alignment from arch info
+                    assembled_size, align=p.archinfo.alignment, flag=MemoryFlag.RX
+                )
                 p.symbols[self.name] = block.mem_addr
                 p.binfmt_tool.update_binary_content(
                     block.file_addr,

tests/test_i386.py

@@ -77,6 +77,22 @@ def test_modify_instruction_patch_pie(self):
             expected_returnCode=0,
         )
 
+    def test_modify_instruction_patch2_nopie(self):
+        self.run_one(
+            "printf_nopie",
+            [ModifyInstructionPatch(0x80491A0, "mov eax, 0x0")],
+            expected_output="",
+            expected_returnCode=0,
+        )
+
+    def test_modify_instruction_patch2_pie(self):
+        self.run_one(
+            "printf_pie",
+            [ModifyInstructionPatch(0x11C9, "mov eax, 0x0")],
+            expected_output="",
+            expected_returnCode=0,
+        )
+
     def test_insert_instruction_patch_nopie(self):
         instrs = """
             mov eax, 0x4

tests/test_x86_64.py

@@ -76,6 +76,22 @@ def test_modify_instruction_patch_pie(self):
             expected_returnCode=0,
         )
 
+    def test_modify_instruction_patch2_nopie(self):
+        self.run_one(
+            "printf_nopie",
+            [ModifyInstructionPatch(0x401152, "lea rax, [rip + 0xeb8]")],
+            expected_output="",
+            expected_returnCode=0,
+        )
+
+    def test_modify_instruction_patch2_pie(self):
+        self.run_one(
+            "printf_pie",
+            [ModifyInstructionPatch(0x115D, "mov eax, 0x0")],
+            expected_output="",
+            expected_returnCode=0,
+        )
+
     def test_insert_instruction_patch_nopie(self):
         instrs = """
             mov rax, 0x1