81 lines
No EOL
3.2 KiB
Python
81 lines
No EOL
3.2 KiB
Python
import json
|
|
import lief
|
|
lief.disable_leak_warning() # warnings to disable for the callback
|
|
|
|
# Retrives all unique DLL names being imported
|
|
def get_used_dlls(calls: list[dict[str,str]]) -> set[str]:
|
|
res = set()
|
|
for call in calls:
|
|
res.add(call["name"].split('!')[0])
|
|
return res
|
|
|
|
# Retrieves all unique function names used for a single DLL name
|
|
def get_used_functions_from_dll(dllname,calls):
|
|
res = set()
|
|
for [dll,func] in map(lambda x: x["name"].split('!'),calls):
|
|
if dll == dllname:
|
|
res.add(func)
|
|
return res
|
|
|
|
def patch_call_to_new_IAT_entry(pe: lief.PE.Binary, call: dict[str,str], rva: int):
|
|
base = pe.imagebase
|
|
instruction_offset = int(call["adress"],16)-base
|
|
# 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
|
|
pe.patch_address(instruction_offset, [0xFF,0x15] + new_value, lief.Binary.VA_TYPES.RVA)
|
|
|
|
|
|
def patch_calls_to_new_IAT(pe: lief.PE.Binary, imp: lief.PE.Import, entry:lief.PE.ImportEntry, rva: int):
|
|
# print(f"{imp.name}!{entry.name}: 0x{rva:010x}")
|
|
for call in filter(lambda x : x["name"] == f"{imp.name.upper()}!{entry.name}" ,calls):
|
|
patch_call_to_new_IAT_entry(pe,call,rva)
|
|
|
|
|
|
|
|
# wave dump file to patch
|
|
with open("rsc/wave-0001.dump", "rb") as f:
|
|
pe = lief.parse(f)
|
|
assert isinstance(pe, lief.PE.Binary)
|
|
|
|
# JSON generated with the python reader files
|
|
with open("rsc/upx-hostname.exe.bin_iat_wave1.json", "r") as iat_json_input:
|
|
iat_data = json.load(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)
|
|
|
|
# patch entrypoint
|
|
entrypoint_format = int(hex(wave_entry)[-4:],16)
|
|
pe.optional_header.addressof_entrypoint = entrypoint_format
|
|
|
|
# remove all current imports
|
|
# pe.remove_all_imports()
|
|
|
|
# recreate all DLL imports
|
|
for dll in get_used_dlls(calls):
|
|
imported_dll = pe.add_import(dll.lower())
|
|
# recreate all function calls related to that dll import
|
|
for func in get_used_functions_from_dll(dll,calls):
|
|
imported_dll.add_entry(func)
|
|
|
|
# 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.
|
|
|
|
# write result
|
|
config = lief.PE.Builder.config_t()
|
|
config.imports = True # allows the config of the writer to write a new IAT
|
|
config.resolved_iat_cbk = patch_calls_to_new_IAT # callback after the IAT has been written
|
|
pe.write("patched.exe", config)
|
|
print("Wrote the patched executable as patched.exe") |