added registry initiation and patch to jump addr

iat.py

@@ -11,6 +11,20 @@ iat_json_path = "rsc/upx-hostname.exe.bin_iat_wave1.json"
 # iat_json_path = "rsc/000155f2e0360f6ff6cd.exe_iat_wave2.json"
 
 
+def hex_address_to_memory_representation(
+    hex_addr: str, is_32b: bool, is_little_endian: bool
+) -> list[int]:
+    adress_size = 4 if is_32b else 8
+    mem_value = [0x00] * adress_size
+    hex_addr = hex_addr[::-1][:-2]  # reversing order and stripping zero
+    for i in range(0, adress_size):
+        byte_str = hex_addr[i * 2 : (i + 1) * 2][::-1]
+        mem_value[i] += int(byte_str, 16)
+    if not is_little_endian:
+        mem_value = mem_value[::-1]  # reverse byte order for big endian
+    return mem_value
+
+
 # Retrives all unique DLL names being imported
 def get_used_dlls(calls: list[dict[str, str]]) -> set[str]:
     res = set()
@@ -30,28 +44,36 @@ def get_used_functions_from_dll(dllname, calls):
 
 def patch_direct_adress_call(pe: lief.PE.Binary, rva: int, instruction_offset: int):
     # We can manually patch the instruction here: FF 15 08 10 00 01 represents `call [0x01001080]`
-    adress_size = 4 if pe.abstract.header.is_32 else 8
-    is_little_endian = pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE
-    new_value = [0x00] * adress_size
-    hex_adress = hex(rva + pe.imagebase)[::-1][
-        :-2
-    ]  # reversing order and stripping zero
-    for i in range(0, adress_size):
-        byte_str = hex_adress[i * 2 : (i + 1) * 2][::-1]
-        new_value[i] += int(byte_str, 16)
-    if not is_little_endian:
-        new_value = new_value[::-1]  # reverse byte order for big endian
+    new_value = hex_address_to_memory_representation(
+        hex(rva + pe.imagebase),
+        pe.abstract.header.is_32,
+        pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE,
+    )
     pe.patch_address(
         instruction_offset, [0xFF, 0x15] + new_value, lief.Binary.VA_TYPES.RVA
     )
 
 
-def patch_call_to_new_IAT_entry(pe: lief.PE.Binary, call: dict[str, str], rva: int):
+def patch_direct_adress_jump(pe: lief.PE.Binary, rva: int, instruction_offset: int):
+    # We can manually patch the instruction here: FF 15 08 10 00 01 represents `call [0x01001080]`
+    new_value = hex_address_to_memory_representation(
+        hex(rva + pe.imagebase),
+        pe.abstract.header.is_32,
+        pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE,
+    )
+    pe.patch_address(
+        instruction_offset, [0xFF, 0x25] + new_value, lief.Binary.VA_TYPES.RVA
+    )
+
+
+def patch_instr_to_new_IAT_entry(pe: lief.PE.Binary, call: dict[str, str], rva: int):
     base = pe.imagebase
     instruction_offset = int(call["adress"], 16) - base
-    memview = pe.get_content_from_virtual_address(instruction_offset,2)
-    if([memview[0],memview[1]] == [0xFF,0x15]):
-        patch_direct_adress_call(pe,rva, instruction_offset)
+    memview = pe.get_content_from_virtual_address(instruction_offset, 2)
+    if [memview[0], memview[1]] == [0xFF, 0x15]:
+        patch_direct_adress_call(pe, rva, instruction_offset)
+    elif [memview[0], memview[1]] == [0xFF, 0x25]:
+        patch_direct_adress_jump(pe, rva, instruction_offset)
 
 
 def patch_calls_to_new_IAT(
@@ -61,7 +83,7 @@ def patch_calls_to_new_IAT(
     for call in filter(
         lambda x: x["name"] == f"{imp.name.upper()}!{entry.name}", calls
     ):
-        patch_call_to_new_IAT_entry(pe, call, rva)
+        patch_instr_to_new_IAT_entry(pe, call, rva)
 
 
 # wave dump file to patch
@@ -75,13 +97,54 @@ with open(iat_json_path, "r") as iat_json_input:
 calls: list[dict[str, str]] = iat_data["calls"]
 wave_entry = int(iat_data["entry"], 16)
 
-# Define all sections as writeable, to help with some weird stuff we're seeing
-for section in pe.sections:
-    section.characteristics_lists.append(lief.PE.Section.CHARACTERISTICS.MEM_WRITE)
-    section.characteristics_lists.append(lief.PE.Section.CHARACTERISTICS.MEM_EXECUTE)
+# create new section
+patch_section = lief.PE.Section(".iatpatch")
+content = []
+
+# initiate registry values
+reg_to_inst_code = {
+    "EAX": 0xC0,
+    "EBX": 0xC3,
+    "ECX": 0xC1,
+    "EDX": 0xC2,
+    "ESI": 0xC6,
+    "EDI": 0xC7,
+    "EBP": 0xC5,
+    # "ESP": 0xC4,
+}
+for reg in iat_data["entry_reg_values"].keys():
+    if reg not in reg_to_inst_code:
+        continue
+    new_instruction = [
+        0xC7,
+        reg_to_inst_code[reg],
+    ] + hex_address_to_memory_representation(
+        iat_data["entry_reg_values"][reg].strip(),
+        pe.abstract.header.is_32,
+        pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE,
+    )
+    for byte in new_instruction:
+        content.append(byte)
+
+
+# add ret to actual OEP
+
+content += [0x68] + hex_address_to_memory_representation(
+    hex(wave_entry),
+    pe.abstract.header.is_32,
+    pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE,
+)
+
+content += [0xC3]
+
+patch_section.content = content
+
+# add new section to PE
+pe.add_section(patch_section)
 
 # patch entrypoint
-entrypoint_format = int(hex(wave_entry)[-4:], 16)
+# entrypoint_format = int(hex(pe.get_section(".iatpatch").virtual_address)[-4:], 16)
+entrypoint_format = int(hex(pe.get_section(".iatpatch").virtual_address)[-4:], 16)
 pe.optional_header.addressof_entrypoint = entrypoint_format
 
 # remove all current imports
@@ -96,6 +159,15 @@ for dll in get_used_dlls(calls):
 
 # At this point, the new IAT will only be constructed when the PE is written. We therefore need to make a callback function to patch calls afterwards.
 
+# Define all sections as writeable, to help with some weird stuff we're seeing
+for section in pe.sections:
+    section.characteristics = (
+        lief.PE.Section.CHARACTERISTICS.MEM_WRITE.value
+        + lief.PE.Section.CHARACTERISTICS.MEM_READ.value
+        + lief.PE.Section.CHARACTERISTICS.MEM_EXECUTE.value
+        + lief.PE.Section.CHARACTERISTICS.CNT_INITIALIZED_DATA.value
+    )
+
 # write result
 config = lief.PE.Builder.config_t()
 config.imports = True  # allows the config of the writer to write a new IAT