mirror of
https://github.com/idapython/src
synced 2026-06-08 14:47:00 +00:00
Arnaud Diederen
138 lines
4.0 KiB
Python
138 lines
4.0 KiB
Python
"""
|
|
summary: assign DOS/PE headers structures to a PE binary
|
|
|
|
description:
|
|
The goal of this script is to demonstrate some usage of the type API.
|
|
|
|
In this script, we:
|
|
|
|
* load a PE64 file in binary mode
|
|
* import some types from the mssdk64 til
|
|
* apply these types at the correct ofsset in the DB
|
|
* finally, rebase the program based on the information stored
|
|
in the ImageBase field of the IMAGE_OPTIONAL_HEADER64.
|
|
|
|
level: intermediate
|
|
"""
|
|
|
|
import ida_typeinf
|
|
import ida_bytes
|
|
import ida_name
|
|
import ida_segment
|
|
import ida_netnode
|
|
import ida_hexrays
|
|
import idc
|
|
|
|
|
|
def create_struct_at(id, ea, var_name = None):
|
|
"""
|
|
Create the structure identified by id at
|
|
address ea.
|
|
"""
|
|
tif = ida_typeinf.tinfo_t()
|
|
ida_hexrays.get_type(id, tif, ida_typeinf.BTF_STRUCT)
|
|
ida_bytes.create_struct(ea, tif.get_size(), id, True)
|
|
if var_name:
|
|
ida_name.set_name(ea, var_name)
|
|
|
|
|
|
def get_struct_size(id):
|
|
"""
|
|
Return the size of the structure identified by id.
|
|
"""
|
|
tif = ida_typeinf.tinfo_t()
|
|
ida_hexrays.get_type(id, tif, ida_typeinf.BTF_STRUCT)
|
|
return tif.get_size()
|
|
|
|
|
|
def get_pe_ea(id):
|
|
"""
|
|
Return the address of the pe header.
|
|
"""
|
|
tif = ida_typeinf.tinfo_t()
|
|
ida_hexrays.get_type(id, tif, ida_typeinf.BTF_STRUCT)
|
|
udt = ida_typeinf.udt_type_data_t()
|
|
tif.get_udt_details(udt)
|
|
udm = udt[tif.get_udt_nmembers() - 1]
|
|
return ida_bytes.get_dword(udm.offset // 8)
|
|
|
|
|
|
|
|
def get_field_off(id, field):
|
|
"""
|
|
Return the offset in bytes of the member
|
|
'field' in the structure identified by 'id'.
|
|
"""
|
|
tif = ida_typeinf.tinfo_t()
|
|
ida_hexrays.get_type(id, tif, ida_typeinf.BTF_STRUCT)
|
|
udt = ida_typeinf.udt_type_data_t()
|
|
tif.get_udt_details(udt)
|
|
udm = udt[udt.find_member(field)]
|
|
return udm.offset // 8
|
|
|
|
|
|
def get_struct_field_off(struct_name, field):
|
|
"""
|
|
Return the offset in bytes of the member
|
|
'field' in the structure identified by 'struct_name'.
|
|
"""
|
|
tif = ida_typeinf.tinfo_t()
|
|
tif.get_named_type(None, struct_name, ida_typeinf.BTF_STRUCT, True, False)
|
|
udt = ida_typeinf.udt_type_data_t()
|
|
tif.get_udt_details(udt)
|
|
udm = udt[udt.find_member(field)]
|
|
return udm.offset // 8
|
|
|
|
|
|
def main():
|
|
|
|
# Add mssdk64_win10 til to the local types.
|
|
ida_typeinf.add_til("mssdk64_win10", ida_typeinf.ADDTIL_DEFAULT)
|
|
|
|
# Get _IMAGE_XXX struct ids. -1 Means that the type is appended
|
|
# to the local type list.
|
|
idh_id = idc.import_type(-1, "_IMAGE_DOS_HEADER")
|
|
inh_id = idc.import_type(-1, "_IMAGE_NT_HEADERS64")
|
|
ish_id = idc.import_type(-1, "_IMAGE_SECTION_HEADER")
|
|
if idh_id == ida_netnode.BADNODE or inh_id == ida_netnode.BADNODE \
|
|
or ish_id == ida_netnode.BADNODE:
|
|
print("At least one type has not been imported. Quit.")
|
|
return
|
|
|
|
# Create _IMAGE_DOS_HEADER at offset 0.
|
|
create_struct_at(idh_id, 0, "DOS_HEADER")
|
|
|
|
# Get the e_lfanew value.
|
|
inh_ea = get_pe_ea(idh_id)
|
|
|
|
# Create _IMAGE_NT_HEADERS at offset inh_ea.
|
|
create_struct_at(inh_id, inh_ea, "NT_HEADERS")
|
|
|
|
# Get the number of sections.
|
|
ifh_off = get_field_off(inh_id, "FileHeader")
|
|
nos_off = get_struct_field_off("_IMAGE_FILE_HEADER", "NumberOfSections")
|
|
sections = ida_bytes.get_word(inh_ea + ifh_off + nos_off)
|
|
|
|
# Create _IMAGE_SECTION_HEADERs.
|
|
section_ea = inh_ea + get_struct_size(inh_id)
|
|
ish_size = get_struct_size(ish_id)
|
|
section_eas = []
|
|
for section in range(sections):
|
|
section_eas.append(section_ea)
|
|
create_struct_at(ish_id, section_ea)
|
|
section_ea += ish_size
|
|
|
|
|
|
ioh_off = get_field_off(inh_id, "OptionalHeader")
|
|
ioh_ea = inh_ea + ioh_off
|
|
|
|
# Get Section alignment.
|
|
sec_align_off = get_struct_field_off("_IMAGE_OPTIONAL_HEADER64", "SectionAlignment")
|
|
|
|
# Get the image base ea.
|
|
image_base_member_off = get_struct_field_off("_IMAGE_OPTIONAL_HEADER64", "ImageBase")
|
|
image_base = ida_bytes.get_qword(ioh_ea + image_base_member_off)
|
|
ida_segment.rebase_program(image_base, ida_segment.MSF_FIXONCE)
|
|
|
|
main()
|