def parse_disk_to_header(sample, query_header=None):
offsets = {}
# Read DOS Header
with open(sample, 'rb') as f:
dos_read = bytearray(f.read(header_sizes["_IMAGE_DOS_HEADER"]))
dos_header = _IMAGE_DOS_HEADER.from_buffer(dos_read)
if getattr(dos_header, "e_magic") != 0x5A4D:
print(f"e_magic = {getattr(dos_header, 'e_magic')}")
print("Wrong DOS Magic Value (MZ). Aborting...")
raise InvalidPEFile
e_lfanew = getattr(dos_header, "e_lfanew")
if query_header == "e_lfanew":
return e_lfanew
if query_header == "_IMAGE_DOS_HEADER":
return dos_header
offsets["_IMAGE_DOS_HEADER"] = 0
# Read PE Header
pe_hdr_offset = e_lfanew
f.seek(pe_hdr_offset)
offsets["_IMAGE_FILE_HEADER"] = pe_hdr_offset
pe_read = bytearray(f.read(header_sizes["_IMAGE_FILE_HEADER"]))
pe_header = _IMAGE_FILE_HEADER.from_buffer(pe_read)
if getattr(pe_header, "Signature") != 0x4550:
print(f"Signature: {getattr(pe_header, 'Signature')}")
print("Wrong PE Header Signature. Aborting...")
raise InvalidPEFile
number_of_sections = getattr(pe_header, "NumberOfSections")
if query_header == "NumberOfSections":
return number_of_sections
if query_header == "_IMAGE_FILE_HEADER":
return pe_header
# Read Optional Header
opt_hdr_offset = pe_hdr_offset + header_sizes["_IMAGE_FILE_HEADER"]
f.seek(opt_hdr_offset)
offsets["_IMAGE_OPTIONAL_HEADER"] = opt_hdr_offset
opt_read = bytearray(f.read(header_sizes["_IMAGE_OPTIONAL_HEADER"]))
opt_header = _IMAGE_OPTIONAL_HEADER.from_buffer(opt_read)
if getattr(opt_header, "Magic") != 0x10B:
print(f"OPT Magic: {getattr(opt_header, 'Magic')}")
print("Wrong Optional Header Magic. Aborting...")
raise InvalidPEFile
if query_header == "_IMAGE_OPTIONAL_HEADER":
return opt_header
if query_header == "_IMAGE_DATA_DIRECTORY":
return getattr(opt_header, "DataDirectory")
rva_to_IMAGE_IMPORT_DESCRIPTOR = getattr(
getattr(opt_header, "DataDirectory")[1], "VirtualAddress")
size_import_table = getattr(
getattr(opt_header, "DataDirectory")[1], "Size")
#import_table = get_imp(uc, rva_to_IMAGE_IMPORT_DESCRIPTOR, base_addr, size_import_table)
import_descriptor_table = []
# TODO Use this when custom loader finished
#for x in range(int((size_import_table / header_sizes["IMAGE_IMPORT_DESCRIPTOR"]))):
# f.seek(rva_to_IMAGE_IMPORT_DESCRIPTOR)
# imp_descriptor_read = bytearray(f.read(header_sizes["IMAGE_IMPORT_DESCRIPTOR"]))
# imp_descriptor = IMAGE_IMPORT_DESCRIPTOR.from_buffer(imp_descriptor_read)
# if getattr(imp_descriptor, "Characteristics") != 0 or getattr(imp_descriptor, "FirstThunk") != 0:
# import_descriptor_table.append(imp_descriptor)
# rva_to_IMAGE_IMPORT_DESCRIPTOR += header_sizes["IMAGE_IMPORT_DESCRIPTOR"]
# else:
# break
if query_header == "IMPORTS":
return import_descriptor_table
# Read Section Header
section_hdr_offset = opt_hdr_offset + header_sizes[
"_IMAGE_OPTIONAL_HEADER"]
offsets["IMAGE_SECTION_HEADER"] = section_hdr_offset
section_headers = []
for i in range(number_of_sections):
f.seek(section_hdr_offset)
sec_read = bytearray(f.read(header_sizes["IMAGE_SECTION_HEADER"]))
sec_header = IMAGE_SECTION_HEADER.from_buffer(sec_read)
section_headers.append(sec_header)
section_hdr_offset += header_sizes["IMAGE_SECTION_HEADER"]
if query_header == "IMAGE_SECTION_HEADER":
return section_headers
if query_header == "Offsets":
return offsets
headers = {
"_IMAGE_DOS_HEADER": dos_header,
"_IMAGE_FILE_HEADER": pe_header,
"_IMAGE_OPTIONAL_HEADER": opt_header,
"IMAGE_SECTION_HEADER": section_headers,
"IMPORTS": import_descriptor_table,
}
return headers
def parse_memory_to_header(uc, base_addr, query_header=None):
# Read DOS Header
uc_dos = uc.mem_read(base_addr, header_sizes["_IMAGE_DOS_HEADER"])
dos_header = _IMAGE_DOS_HEADER.from_buffer(uc_dos)
if getattr(dos_header, "e_magic") != 0x5A4D:
print(f"e_magic = {getattr(dos_header, 'e_magic')}")
print("Wrong DOS Magic Value (MZ). Aborting...")
raise InvalidPEFile
e_lfanew = getattr(dos_header, "e_lfanew")
if query_header == "e_lfanew":
return e_lfanew
if query_header == "_IMAGE_DOS_HEADER":
return dos_header
# Read PE Header
pe_hdr_offset = base_addr + e_lfanew
uc_pe = uc.mem_read(pe_hdr_offset, header_sizes["_IMAGE_FILE_HEADER"])
pe_header = _IMAGE_FILE_HEADER.from_buffer(uc_pe)
if getattr(pe_header, "Signature") != 0x4550:
print(f"Signature: {getattr(pe_header, 'Signature')}")
print("Wrong PE Header Signature. Aborting...")
raise InvalidPEFile
number_of_sections = getattr(pe_header, "NumberOfSections")
if query_header == "NumberOfSections":
return number_of_sections
if query_header == "_IMAGE_FILE_HEADER":
return pe_header
# Read Optional Header
opt_hdr_offset = pe_hdr_offset + header_sizes["_IMAGE_FILE_HEADER"]
uc_opt = uc.mem_read(opt_hdr_offset,
header_sizes["_IMAGE_OPTIONAL_HEADER"])
opt_header = _IMAGE_OPTIONAL_HEADER.from_buffer(uc_opt)
if getattr(opt_header, "Magic") != 0x10B:
print(f"OPT Magic: {getattr(opt_header, 'Magic')}")
print("Wrong Optional Header Magic. Aborting...")
raise InvalidPEFile
if query_header == "_IMAGE_OPTIONAL_HEADER":
return opt_header
if query_header == "_IMAGE_DATA_DIRECTORY":
return getattr(opt_header, "DataDirectory")
rva_to_IMAGE_IMPORT_DESCRIPTOR = getattr(
getattr(opt_header, "DataDirectory")[1], "VirtualAddress")
size_import_table = getattr(
getattr(opt_header, "DataDirectory")[1], "Size")
import_table = get_imp(uc, rva_to_IMAGE_IMPORT_DESCRIPTOR, base_addr,
size_import_table)
if query_header == "IMPORTS":
return import_table
# Read Section Header
section_hdr_offset = opt_hdr_offset + header_sizes["_IMAGE_OPTIONAL_HEADER"]
section_headers = []
for i in range(number_of_sections):
uc_sec = uc.mem_read(section_hdr_offset,
header_sizes["IMAGE_SECTION_HEADER"])
sec_header = IMAGE_SECTION_HEADER.from_buffer(uc_sec)
section_headers.append(sec_header)
section_hdr_offset += header_sizes["IMAGE_SECTION_HEADER"]
if query_header == "IMAGE_SECTION_HEADER":
return section_headers
headers = {
"_IMAGE_DOS_HEADER": dos_header,
"_IMAGE_FILE_HEADER": pe_header,
"_IMAGE_OPTIONAL_HEADER": opt_header,
"IMAGE_SECTION_HEADER": section_headers,
"IMPORTS": import_table,
}
return headers