From 3d44b1f071f31d8ffe2c056a002292c880a24568 Mon Sep 17 00:00:00 2001 From: Joey Payne Date: Sat, 18 Oct 2014 22:07:10 +1300 Subject: [PATCH] Finished off making pe.py process the .rsrc section and added a method to replace the image. --- pe.py | 751 +++++++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 693 insertions(+), 58 deletions(-) diff --git a/pe.py b/pe.py index ab27dc7..21a315b 100644 --- a/pe.py +++ b/pe.py @@ -1,5 +1,10 @@ import os import struct +try: + import PythonMagick as pm +except ImportError: + print 'PythonMagick is required to replace pe icon file. This functionality is disabled because you need to install it.' + pm = None struct_symbols = {1:'B',#byte 2:'H',#word @@ -151,6 +156,91 @@ name_dictionary = {'PEHeader_Machine': { }, } +DEFAULT_ENDIAN = 'little' + +def read_data(file_data, offset, number_of_bytes, string_data=None): + """Just reads the straight data with no endianness.""" + + if number_of_bytes > 0: + data = file_data[offset:offset+number_of_bytes] + #if len(data) != number_of_bytes: + #print 'data out of bounds:', 'offset', hex(offset), 'data', data, 'data_len', len(data), 'num_bytes', number_of_bytes, 'total', hex(len(file_data)) + return data + else: + return bytearray('') + +def read_bytes(file_data, offset, number_of_bytes, endian=None, string_data=None): + """Returns a tuple of the data value and string representation. + Will read 1,2,4,8 bytes with little endian as the default + (value, string) + """ + + if number_of_bytes > 0: + + endian = endian or DEFAULT_ENDIAN + endian = endian_symbols[endian] + + data = str(file_data[offset:offset+number_of_bytes]) + if len(data) != number_of_bytes: + #print 'data out of bounds:', 'offset', hex(offset), 'data', data, 'data_len', len(data), 'num_bytes', number_of_bytes, 'total', hex(len(file_data)) + return 0, u'' + + return struct.unpack(endian+struct_symbols[number_of_bytes], data)[0], data + else: + return 0, u'' + +def value_to_byte_string(value, number_of_bytes, endian=None): + endian = endian or DEFAULT_ENDIAN + endian = endian_symbols[endian] + + return struct.pack(endian+struct_symbols[number_of_bytes], value) + +class ResourceTypes(object): + Cursor = 1 + Bitmap = 2 + Icon = 3 + Menu = 4 + Dialog = 5 + String = 6 + Font_Directory = 7 + Font = 8 + Accelerator = 9 + RC_Data = 10 + Message_Table = 11 + Group_Cursor = 12 + Group_Icon = 14 + Version_Info = 16 + DLG_Include = 17 + Plug_Play = 19 + VXD = 20 + Animated_Cursor = 21 + Animated_Icon = 22 + HTML = 23 + Manifest = 24 + + +resource_types = {1: 'Cursor', + 2: 'Bitmap', + 3: 'Icon', + 4: 'Menu', + 5: 'Dialog', + 6: 'String', + 7: 'Font Directory', + 8: 'Font', + 9: 'Accelerator', + 10: 'RC Data', + 11: 'Message Table', + 12: 'Group Cursor', + 14: 'Group Icon', + 16: 'Version Info', + 17: 'DLG Include', + 19: 'Plug and Play', + 20: 'VXD', + 21: 'Animated Cursor', + 22: 'Animated Icon', + 23: 'HTML', + 24: 'Manifest'} + _32BIT_PLUS_MAGIC = 0x20b _32BIT_MAGIC = 0x10b _ROM_MAGIC = 0x107 @@ -163,6 +253,18 @@ def test_bit(value, index): mask = 1 << index return (value & mask) +def set_bit(value, index): + mask = 1 << index + return (value | mask) + +def clear_bit(value, index): + mask = ~(1 << index) + return (value & mask) + +def toggle_bit(value, index): + mask = 1 << index + return (value ^ mask) + class PEFormatError(Exception): pass @@ -201,29 +303,51 @@ class Structure(Printable): _fields = {} - def __init__(self, offset=0, size=0, value=None, data=None, absolute_offset=0, + def __init__(self, size=0, value=None, data=None, absolute_offset=0, name='', friendly_name='', *args, **kwargs): - super(Structure, self).__init__() - self.offset = offset + self._value = value self.size = size - self.value = value self.data = data self.name = name self.friendly_name = friendly_name - self.absolute_offset = absolute_offset + self._absolute_offset = absolute_offset + self._file_data = None for k, v in kwargs.items(): setattr(self, k, v) - def process_field(self, pe_file, field_name, field_info): + @property + def absolute_offset(self): + return self._absolute_offset + + @absolute_offset.setter + def absolute_offset(self, abs_offset): + self._absolute_offset = abs_offset + for k, v in self._fields.items(): + field = getattr(self, k) + field.absolute_offset = self.absolute_offset + field.offset + + @property + def value(self): + return self._value + + @value.setter + def value(self, value): + if self._file_data is not None: + self.data = value_to_byte_string(value, self.size) + self._file_data[self.absolute_offset:self.absolute_offset+self.size] = bytearray(self.data) + self._value = value + + + def process_field(self, file_data, field_name, field_info): if hasattr(self, 'process_'+field_name) and callable(getattr(self, 'process_'+field_name)): - getattr(self, 'process_'+field_name)(pe_file, field_name, field_info) + getattr(self, 'process_'+field_name)(file_data, field_name, field_info) else: - absolute_offset = field_info['offset'] + self.offset + absolute_offset = field_info['offset'] + self.absolute_offset size = field_info['size'] self.size += size - int_value, data = pe_file.read_bytes(absolute_offset, size) + int_value, data = read_bytes(file_data, absolute_offset, size) field_name_dict = read_from_name_dict(self, field_name) name = field_name_dict.get(int_value, '') friendly_name = name.replace('_', ' ').capitalize() @@ -233,13 +357,14 @@ class Structure(Printable): value=int_value, data=data, absolute_offset=absolute_offset, name=name, friendly_name=friendly_name)) + getattr(self, field_name)._file_data = file_data - def process_Characteristics(self, pe_file, field_name, field_info): - absolute_offset = field_info['offset'] + self.offset + def process_Characteristics(self, file_data, field_name, field_info): + absolute_offset = field_info['offset'] + self.absolute_offset size = field_info['size'] self.size += size - int_value, data = pe_file.read_bytes(absolute_offset, size) + int_value, data = read_bytes(file_data, absolute_offset, size) field_name_dict = read_from_name_dict(self, field_name) bit_length = len(bin(int_value))-2 @@ -247,8 +372,9 @@ class Structure(Printable): characteristics = {} for i in xrange(bit_length): set_bit = test_bit(int_value, i) - if set_bit != 0: - characteristics[field_name_dict[set_bit]] = set_bit + char_name = field_name_dict.get(set_bit, '') + if set_bit != 0 and char_name: + characteristics[char_name] = set_bit setattr(self, field_name, Structure(offset=field_info['offset'], @@ -257,13 +383,15 @@ class Structure(Printable): absolute_offset=absolute_offset, values=characteristics, )) + getattr(self, field_name)._file_data = file_data @classmethod - def parse_from_data(cls, pe_file, **cls_args): + def parse_from_data(cls, file_data, **cls_args): """Parses the Structure from the file data.""" self = cls(**cls_args) + self._file_data = file_data for field_name, field_info in self._fields.items(): - self.process_field(pe_file, field_name, field_info) + self.process_field(file_data, field_name, field_info) return self @@ -436,7 +564,7 @@ class OptionalHeader(Structure): 'size':1}, 'SizeOfCode':{'offset':4, 'size':4}, - 'SizeOfInitializedData':{'offset':8, + 'SizeOfInitializedData':{'offset':8,# 'size':4}, 'SizeOfUninitializedData':{'offset':12, 'size':4}, @@ -466,12 +594,10 @@ class OptionalHeader(Structure): 'size':2}, 'Reserved':{'offset':52, 'size':4}, - 'SizeOfImage':{'offset':56, + 'SizeOfImage':{'offset':56,# 'size':4}, 'SizeOfHeaders':{'offset':60, 'size':4}, - 'SizeOfHeaders':{'offset':60, - 'size':4}, 'CheckSum': {'offset': 64, 'size':4}, 'Subsystem':{'offset':68, @@ -500,7 +626,7 @@ class OptionalHeader(Structure): 'size':4}, 'ResourceTableAddress':{'offset':112, 'size':4}, - 'ResourceTableSize':{'offset':116, + 'ResourceTableSize':{'offset':116,# 'size':4}, 'ExceptionTableAddress':{'offset':120, 'size':4}, @@ -510,7 +636,7 @@ class OptionalHeader(Structure): 'size':4}, 'CertificateTableSize':{'offset':132, 'size':4}, - 'BaseRelocationTableAddress':{'offset':136, + 'BaseRelocationTableAddress':{'offset':136,# 'size':4}, 'BaseRelocationTableSize':{'offset':140, 'size':4}, @@ -556,15 +682,36 @@ class OptionalHeader(Structure): } - def process_DLL_Characteristics(self, pe_file, field_name, field_info): - self.process_Characteristics(pe_file, field_name, field_info) + def process_DLL_Characteristics(self, file_data, field_name, field_info): + self.process_Characteristics(file_data, field_name, field_info) + def process_field(self, file_data, field_name, field_info): + + if hasattr(self, 'process_'+field_name) and callable(getattr(self, 'process_'+field_name)): + getattr(self, 'process_'+field_name)(file_data, field_name, field_info) + else: + absolute_offset = field_info['offset'] + self.absolute_offset + size = field_info['size'] + self.size += size + int_value, data = read_bytes(file_data, absolute_offset, size) + field_name_dict = read_from_name_dict(self, field_name) + name = field_name_dict.get(int_value, '') + friendly_name = name.replace('_', ' ').capitalize() + + setattr(self, field_name, Structure(offset=field_info['offset'], + size=size, + value=int_value, data=data, + absolute_offset=absolute_offset, + name=name, friendly_name=friendly_name)) + getattr(self, field_name)._file_data = file_data + @classmethod - def parse_from_data(cls, pe_file, **cls_args): + def parse_from_data(cls, file_data, **cls_args): """Parses the Structure from the file data.""" self = cls(**cls_args) - magic, _ = int_value, data = pe_file.read_bytes(self.offset, 2) + self._file_data = file_data + magic, _ = read_bytes(file_data, self.absolute_offset, 2) if magic == _32BIT_MAGIC: self._fields = self._fields_32 @@ -574,7 +721,7 @@ class OptionalHeader(Structure): raise PEFormatError('Magic for Optional Header is invalid.') for field_name, field_info in self._fields.items(): - self.process_field(pe_file, field_name, field_info) + self.process_field(file_data, field_name, field_info) return self @@ -583,13 +730,13 @@ class SectionHeader(Structure): _fields = {'Name':{'offset':0, 'size':8}, - 'VirtualSize':{'offset':8, + 'VirtualSize':{'offset':8, #.rsrc 'size':4}, - 'VirtualAddress':{'offset':12, + 'VirtualAddress':{'offset':12,#.reloc 'size':4}, - 'SizeOfRawData':{'offset':16, + 'SizeOfRawData':{'offset':16,#.rsrc 'size':4}, - 'PointerToRawData':{'offset':20, + 'PointerToRawData':{'offset':20,#.reloc 'size':4}, 'PointerToRelocations':{'offset':24, 'size':4}, @@ -603,10 +750,381 @@ class SectionHeader(Structure): 'size':4} } +class ResourceDirectoryTable(Structure): + _fields = {'Characteristics':{'offset':0, + 'size':4}, + 'TimeDateStamp':{'offset':4, + 'size':4}, + 'MajorVersion':{'offset':8, + 'size':2}, + 'MinorVersion':{'offset':10, + 'size':2}, + 'NumberOfNameEntries':{'offset':12, + 'size':2}, + 'NumberOfIDEntries':{'offset':14, + 'size':2} + } + + def __init__(self, *args, **kwargs): + self.name_entries = [] + self.id_entries = [] + self.subdirectory_tables = [] + self.data_entries = [] + + super(ResourceDirectoryTable, self).__init__(*args, **kwargs) + +class ResourceDirectoryEntryName(Structure): + _fields = {'NameRVA':{'offset':0, + 'size':4}, + 'DataOrSubdirectoryEntryRVA':{'offset':4, #high bit 1 for subdir RVA + 'size':4} + } + + + directory_string = None + + def is_data_entry(self): + return (test_bit(self.DataOrSubdirectoryEntryRVA.value, 31) == 0) + + def data_rva_empty(self): + return self.get_data_or_subdirectory_rva() == 0 + + def get_data_or_subdirectory_rva(self, virtual_to_physical=0): + return clear_bit(self.DataOrSubdirectoryEntryRVA.value-virtual_to_physical, 31) + + def get_data_or_subdirectory_absolute_offset(self): + return self.get_data_or_subdirectory_rva() + self._section_header.PointerToRawData.value + + def get_name_absolute_offset(self): + return clear_bit(self.NameRVA.value, 31) + self._section_header.PointerToRawData.value + + +class ResourceDirectoryEntryID(Structure): + _fields = {'IntegerID':{'offset':0, + 'size':4}, + 'DataOrSubdirectoryEntryRVA':{'offset':4,#high bit 1 for Subdir RVA + 'size':4} + } + + + def is_data_entry(self): + return (test_bit(self.DataOrSubdirectoryEntryRVA.value, 31) == 0) + + def data_rva_empty(self): + return self.get_data_or_subdirectory_rva() == 0 + + def get_data_or_subdirectory_rva(self, virtual_to_physical=0): + return clear_bit(self.DataOrSubdirectoryEntryRVA.value-virtual_to_physical, 31) + + def get_data_or_subdirectory_absolute_offset(self, vtp=0): + return self.get_data_or_subdirectory_rva(vtp) + self._section_header.PointerToRawData.value + +class ResourceDirectoryString(Structure): + _fields = {'Length':{'offset':0, + 'size':2}, + #String : offset=2, len=Length + } + + @classmethod + def parse_from_data(cls, file_data, **cls_args): + """Parses the Structure from the file data.""" + self = cls(**cls_args) + self._file_data = file_data + str_len, _ = read_bytes(file_data, self.absolute_offset, 2) + + self._fields['String'] = {'offset':2, 'size':str_len} + + for field_name, field_info in self._fields.items(): + self.process_field(file_data, field_name, field_info) + + return self + + def process_String(self, file_data, field_name, field_info): + absolute_offset = field_info['offset'] + self.absolute_offset + size = field_info['size'] + self.size += size + data = u'' + for i in xrange(size): + val, dat = read_bytes(file_data, absolute_offset+i*2,2) + data += dat + + setattr(self, field_name, Structure(offset=field_info['offset'], + size=size, + data=data, + absolute_offset=absolute_offset, + )) +class ResourceDataEntry(Structure): + _fields = {'DataRVA':{'offset':0, + 'size':4}, + 'Size':{'offset':4, + 'size':4}, + 'Codepage':{'offset':8, + 'size':4}, + 'Reserved':{'offset':12, + 'size':4}, + } + + + def get_data_absolute_offset(self): + return self._section_header.PointerToRawData.value - self._section_header.VirtualAddress.value + self.DataRVA.value + + def process_field(self, file_data, field_name, field_info): + + if hasattr(self, 'process_'+field_name) and callable(getattr(self, 'process_'+field_name)): + getattr(self, 'process_'+field_name)(file_data, field_name, field_info) + else: + absolute_offset = field_info['offset'] + self.absolute_offset + size = field_info['size'] + self.size += size + int_value, data = read_bytes(file_data, absolute_offset, size) + field_name_dict = read_from_name_dict(self, field_name) + name = field_name_dict.get(int_value, '') + friendly_name = name.replace('_', ' ').capitalize() + + setattr(self, field_name, Structure(offset=field_info['offset'], + size=size, + value=int_value, data=data, + absolute_offset=absolute_offset, + name=name, friendly_name=friendly_name)) + getattr(self, field_name)._file_data = file_data + + @classmethod + def parse_from_data(cls, file_data, **cls_args): + """Parses the Structure from the file data.""" + self = cls(**cls_args) + self._file_data = file_data + for field_name, field_info in self._fields.items(): + self.process_field(file_data, field_name, field_info) + + self.data = read_data(file_data, self.get_data_absolute_offset(), self.Size.value) + + return self + +class ResourceHeader(Structure): + _fields = {'DataSize':{'offset':0, + 'size':4}, + 'HeaderSize':{'offset':4, + 'size':4}, + 'Type':{'offset':8, + 'size':4}, + 'Name':{'offset':12, + 'size':4}, + 'DataVersion':{'offset':16, + 'size':4}, + 'MemoryFlags':{'offset':20, + 'size':2}, + 'LanguageID':{'offset':22, + 'size':2}, + 'Version':{'offset':24, + 'size':4}, + 'Characteristics':{'offset':28, + 'size':4}, + } + def get_name(self): + return resource_types[self.Type.value] + + def set_name(self, value): + for k,v in resource_types.items(): + if v == value: + self.Type.value = k + return + + +class IconHeader(Structure): + _fields = {'Reserved':{'offset':0, + 'size':2}, + 'ImageType':{'offset':2,#1 for ICO, 2 for CUR, others invalid + 'size':2}, + 'ImageCount':{'offset':4, + 'size':2}, + } + + def copy_from(self, group_header): + self.Reserved.value = group_header.Reserved.value + self.ImageType.value = group_header.ResourceType.value + self.ImageCount.value = group_header.ResourceCount.value + + self.entries = [] + entry_offset = 0 + self.total_size = self.size + for group_entry in group_header.entries: + icon_entry = IconEntry.parse_from_data(bytearray(''), absolute_offset=self.absolute_offset+self.size+entry_offset, offset=entry_offset) + icon_entry._file_data = self._file_data + icon_entry.copy_from(group_entry) + icon_entry.number = group_entry.number + self.entries.append(icon_entry) + entry_offset += icon_entry.size + self.total_size += icon_entry.size + + @classmethod + def parse_from_data(cls, file_data, **cls_args): + """Parses the Structure from the file data.""" + self = cls(**cls_args) + self._file_data = file_data + + for field_name, field_info in self._fields.items(): + self.process_field(file_data, field_name, field_info) + + self.entries = [] + entry_offset = 0 + self.total_size = self.size + for i in xrange(self.ImageCount.value): + entry = IconEntry.parse_from_data(file_data, absolute_offset=self.absolute_offset+self.size+entry_offset, offset=entry_offset) + entry.number = i + 1 + self.entries.append(entry) + entry_offset += entry.size + self.total_size += entry.size + + return self + + +class GroupHeader(Structure): + _fields = {'Reserved':{'offset':0, + 'size':2}, + 'ResourceType':{'offset':2,#1 for ICO, 2 for CUR, others invalid + 'size':2}, + 'ResourceCount':{'offset':4, + 'size':2}, + } + + def copy_from(self, icon_header): + self.Reserved._file_data = self._file_data + self.ResourceType._file_data = self._file_data + self.ResourceCount._file_data = self._file_data + + self.Reserved.value = icon_header.Reserved.value + self.ResourceType.value = icon_header.ImageType.value + self.ResourceCount.value = icon_header.ImageCount.value + + self.entries = [] + entry_offset = 0 + self.total_size = self.size + for icon_entry in icon_header.entries: + group_entry = GroupEntry.parse_from_data(bytearray(''), absolute_offset=self.absolute_offset+self.size+entry_offset, offset=entry_offset) + group_entry._file_data = self._file_data + group_entry.copy_from(icon_entry) + group_entry.number = icon_entry.number + self.entries.append(group_entry) + entry_offset += group_entry.size + self.total_size += group_entry.size + + + @classmethod + def parse_from_data(cls, file_data, **cls_args): + """Parses the Structure from the file data.""" + self = cls(**cls_args) + self._file_data = file_data + + for field_name, field_info in self._fields.items(): + self.process_field(file_data, field_name, field_info) + + self.entries = [] + entry_offset = 0 + self.total_size = self.size + for i in xrange(self.ResourceCount.value): + entry = GroupEntry.parse_from_data(file_data, absolute_offset=self.absolute_offset+self.size+entry_offset, offset=entry_offset) + entry.number = i + 1 + self.entries.append(entry) + entry_offset += entry.size + self.total_size += entry.size + + return self + +class IconEntry(Structure): + _fields = {'Width':{'offset':0, + 'size':1}, + 'Height':{'offset':1, + 'size':1}, + 'ColorCount':{'offset':2, + 'size':1}, + 'Reserved':{'offset':3, + 'size':1}, + 'ColorPlanes':{'offset':4, + 'size':2}, + 'BitCount':{'offset':6, #bits per pixel + 'size':2}, + 'DataSize':{'offset':8, + 'size': 4}, + 'OffsetToData':{'offset':12, #from start of file + 'size':4}, + } + + def copy_from(self, group_entry, entries): + self.Width.value = group_entry.Width.value + self.Height.value = group_entry.Height.value + self.ColorCount.value = group_entry.ColorCount.value + self.Reserved.value = group_entry.Reserved.value + self.ColorPlanes.value = group_entry.ColorPlanes.value + self.BitCount.value = group_entry.BitCount.value + self.DataSize.value = group_entry.DataSize.value + self.OffsetToData.value = self._get_entry_offset(group_entry, entries) + + def _get_entry_offset(self, group_entry, group_entries): + offset = 6 #Default icon header size + offset += self.size * len(group_entries) + + for i in xrange(group_entry.number-1): + offset += group_entries[i].DataSize.value + + return offset + + + @classmethod + def parse_from_data(cls, file_data, **cls_args): + """Parses the Structure from the file data.""" + self = cls(**cls_args) + self._file_data = file_data + for field_name, field_info in self._fields.items(): + self.process_field(file_data, field_name, field_info) + + self.data = read_data(file_data, self.OffsetToData.value, self.DataSize.value) + + return self + + +class GroupEntry(Structure): + _fields = {'Width':{'offset':0, + 'size':1}, + 'Height':{'offset':1, + 'size':1}, + 'ColorCount':{'offset':2, + 'size':1}, + 'Reserved':{'offset':3, + 'size':1}, + 'ColorPlanes':{'offset':4, + 'size':2}, + 'BitCount':{'offset':6, + 'size':2}, + 'DataSize':{'offset':8, + 'size': 4}, + 'IconCursorId':{'offset':12, + 'size':2}, + } + + def copy_from(self, icon_entry): + self.Width._file_data = self._file_data + self.Height._file_data = self._file_data + self.ColorCount._file_data = self._file_data + self.Reserved._file_data = self._file_data + self.ColorPlanes._file_data = self._file_data + self.BitCount._file_data = self._file_data + self.DataSize._file_data = self._file_data + self.IconCursorId._file_data = self._file_data + + self.Width.value = icon_entry.Width.value + self.Height.value = icon_entry.Height.value + self.ColorCount.value = icon_entry.ColorCount.value + self.Reserved.value = icon_entry.Reserved.value + self.ColorPlanes.value = icon_entry.ColorPlanes.value + self.BitCount.value = icon_entry.BitCount.value + self.DataSize.value = icon_entry.DataSize.value + self.IconCursorId.value = icon_entry.number class PEFile(Printable): - """Reads a portable exe file in either big or little endian.""" + """Reads a portable exe file in either big or little endian. + Right now this only reads the .rsrc section. + """ signature = 'MZ' dos_header = None @@ -618,22 +1136,157 @@ class PEFile(Printable): if not self.is_PEFile(): raise PEFormatError('File is not a proper portable executable formatted file!') - self.pe_file_data = open(self.file_path,'rb').read() + self.pe_file_data = bytearray(open(self.file_path,'rb').read()) - self.dos_header = DOSHeader.parse_from_data(self) - self.pe_header = PEHeader.parse_from_data(self, offset=self.dos_header.PEHeaderOffset.value) - self.optional_header = OptionalHeader.parse_from_data(self, offset=self.pe_header.size+self.pe_header.offset) + self.dos_header = DOSHeader.parse_from_data(self.pe_file_data) + self.pe_header = PEHeader.parse_from_data(self.pe_file_data, absolute_offset=self.dos_header.PEHeaderOffset.value) + self.optional_header = OptionalHeader.parse_from_data(self.pe_file_data, absolute_offset=self.pe_header.size+self.pe_header.absolute_offset) number_of_sections = self.pe_header.NumberOfSections.value section_size = 40 - section_offset = self.pe_header.size+self.pe_header.offset+self.pe_header.SizeOfOptionalHeader.value - self.section_headers = [] + section_offset = self.pe_header.size+self.pe_header.absolute_offset+self.pe_header.SizeOfOptionalHeader.value + self.sections = {} for section_number in xrange(number_of_sections): - section = SectionHeader.parse_from_data(self, offset=section_offset) + section_header = SectionHeader.parse_from_data(self.pe_file_data, absolute_offset=section_offset) section_offset += section_size - self.section_headers.append(section) + self.sections[section_header.Name.data.strip('\x00')] = section_header + if section_header.PointerToLineNumbers.value != 0: + print '{} section contains line number COFF table, which is not implemented yet.'.format(section_header.Name) + + if section_header.PointerToRelocations.value != 0: + print '{} section contains relocation table, which is not implemented yet.'.format(section_header.Name) + + if section_header.Name.data == '.rsrc\x00\x00\x00': + current_table_pointer = section_header.PointerToRawData.value + current_resource_directory_table = ResourceDirectoryTable.parse_from_data(self.pe_file_data, absolute_offset=current_table_pointer, _section_header=section_header, type=None) + self.resource_directory_table = current_resource_directory_table + cur_level = 0 + stack = [(current_resource_directory_table, cur_level)] + + delta = section_header.VirtualAddress.value - section_header.PointerToRawData.value + + while stack: + resource_directory_table, level = stack.pop() + num_name_entries = resource_directory_table.NumberOfNameEntries.value + num_id_entries = resource_directory_table.NumberOfIDEntries.value + current_offset = resource_directory_table.absolute_offset + resource_directory_table.size + + for i in xrange(num_name_entries): + name_entry = ResourceDirectoryEntryName.parse_from_data(self.pe_file_data, absolute_offset=current_offset, _section_header=section_header) + current_offset += name_entry.size + + string_offset = name_entry.get_name_absolute_offset() + name_entry.directory_string = ResourceDirectoryString.parse_from_data(self.pe_file_data, absolute_offset=string_offset, _section_header=section_header) + + offset = name_entry.get_data_or_subdirectory_absolute_offset() + + if not name_entry.data_rva_empty(): + if name_entry.is_data_entry(): + rd = ResourceDataEntry.parse_from_data(self.pe_file_data, absolute_offset=offset, _section_header=section_header) + resource_directory_table.data_entries.append(rd) + else: + rd = ResourceDirectoryTable.parse_from_data(self.pe_file_data, absolute_offset=offset, _section_header=section_header, type=None) + resource_directory_table.subdirectory_tables.append(rd) + + stack.append((rd, level+1)) + + resource_directory_table.name_entries.append(name_entry) + + for i in xrange(num_id_entries): + id_entry = ResourceDirectoryEntryID.parse_from_data(self.pe_file_data, absolute_offset=current_offset, _section_header=section_header) + current_offset += id_entry.size + + offset = id_entry.get_data_or_subdirectory_absolute_offset() + + if id_entry.is_data_entry(): + + rd = ResourceDataEntry.parse_from_data(self.pe_file_data, absolute_offset=offset, _section_header=section_header) + resource_directory_table.data_entries.append(rd) + else: + id_entry.name = str(id_entry.IntegerID.value) + if level+1 == 1: + id_entry.name = resource_types[id_entry.IntegerID.value] + rd = ResourceDirectoryTable.parse_from_data(self.pe_file_data, absolute_offset=offset, _section_header=section_header, type=id_entry.IntegerID.value) + resource_directory_table.subdirectory_tables.append(rd) + stack.append((rd, level+1)) + resource_directory_table.id_entries.append(id_entry) + + + def replace_icon(self, icon_path): + """Replaces an icon in the pe file with the one specified. + This only replaces the largest icon and resizes the input + image to match so that the data is undisturbed. I tried to + update the pointers automatically by moving the data to the end + of the file, but that did not work. Comments were left as history + to what I attempted. + """ + + if pm is None: + raise Exception('PythonMagick is required to run this function.') + + resource_section = self.sections['.rsrc'] + + g_icon_dir = self.get_directory_by_type(ResourceTypes.Group_Icon) + g_icon_data_entry = g_icon_dir.subdirectory_tables[0].data_entries[0] + icon_dir = self.get_directory_by_type(ResourceTypes.Icon) + icon_data_entry = icon_dir.subdirectory_tables[0].data_entries[0] + + group_header = GroupHeader.parse_from_data(self.pe_file_data, absolute_offset=g_icon_data_entry.get_data_absolute_offset()) + g_entry = group_header.entries[0] + + icon = pm.Image(icon_path) + icon.sample('!{}x{}'.format(g_entry.Width.value, g_entry.Height.value)) #resize. Force aspect ratio to change with ! + icon.magick('ICO') # convert to ico + b = pm.Blob() + icon.write(b) + + icon_data = bytearray(b.data) + + icon_header = IconHeader.parse_from_data(icon_data, absolute_offset=0) + + #group_header.absolute_offset = len(self.pe_file_data) + #g_icon_data_entry.DataRVA.value = len(self.pe_file_data) - resource_section.PointerToRawData.value + resource_section.VirtualAddress.value + #padding = 6+14*len(icon_header.entries) + #g_icon_data_entry.Size.value = padding + + #self.pe_file_data += bytearray(padding) + group_header._file_data = self.pe_file_data + group_header.copy_from(icon_header) + + #icon_data_entry.DataRVA.value = len(self.pe_file_data) - resource_section.PointerToRawData.value + resource_section.VirtualAddress.value + #print hex(icon_data_entry.DataRVA.value), hex(len(self.pe_file_data)), hex(icon_data_entry.get_data_absolute_offset()) + + #print hex(read_bytes(self.pe_file_data[icon_data_entry.DataRVA.absolute_offset:icon_data_entry.DataRVA.absolute_offset+icon_data_entry.DataRVA.size],0, icon_data_entry.DataRVA.size)[0]) + + #data = bytearray() + #for entry in icon_header.entries: + #data += entry.data + #self.pe_file_data += entry.data + + #icon_data_entry.Size.value = len(data) + + #self.optional_header.SizeOfImage.value = self.optional_header.SizeOfImage.value + len(data) + padding + #self.optional_header.ResourceTableSize.value = self.optional_header.ResourceTableSize.value + len(data) + padding + #self.optional_header.SizeOfInitializedData.value = self.optional_header.SizeOfInitializedData.value + len(data) + padding + #resource_section.SizeOfRawData.value = resource_section.SizeOfRawData.value + len(data) + padding + #resource_section.VirtualSize.value = resource_section.VirtualSize.value + len(data) + padding + #print icon_header.total_size + data_address = icon_data_entry.get_data_absolute_offset() + data_size = icon_data_entry.Size.value + self.pe_file_data = self.pe_file_data[:data_address] + icon_data[icon_header.total_size:] + self.pe_file_data[data_address+data_size:] + + + def write(self, file_name): + with open(file_name, 'wb+') as f: + f.write(self.pe_file_data) + + def get_directory_by_type(self, type): + """Gets the directory by resource type.""" + for d in self.resource_directory_table.subdirectory_tables: + if d.type == type: + return d def is_PEFile(self): """Checks if the file is a proper PE file""" @@ -645,21 +1298,3 @@ class PEFile(Printable): raise e finally: return signature == self.signature - - def read_bytes(self, offset, number_of_bytes, endian=None): - """Returns a tuple of the data value and string representation. - (value, string) - """ - - if number_of_bytes > 0: - - if endian: - endian = endian_symbols[endian] - else: - endian = endian_symbols[self.endian] - - data = self.pe_file_data[offset:offset+number_of_bytes] - - return struct.unpack(endian+struct_symbols[number_of_bytes], data)[0], data - else: - return 0,''