Mercurial > touhou
diff pytouhou/utils/pe.py @ 229:5afc75f71fed
Add “SHT” support to EoSD, and do a little cleanup.
author | Emmanuel Gil Peyrot <linkmauve@linkmauve.fr> |
---|---|
date | Fri, 30 Dec 2011 18:37:06 +0100 |
parents | |
children | e15672733c93 |
line wrap: on
line diff
new file mode 100644 --- /dev/null +++ b/pytouhou/utils/pe.py @@ -0,0 +1,137 @@ +# -*- encoding: utf-8 -*- +## +## Copyright (C) 2011 Thibaut Girka <thib@sitedethib.com> +## +## This program is free software; you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation; version 3 only. +## +## This program is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## + +from struct import Struct, unpack +from collections import namedtuple + + +class PEStructs: + _IMAGE_FILE_HEADER = namedtuple('_IMAGE_FILE_HEADER', + ('Machine', + 'NumberOfSections', + 'TimeDateStamp', + 'PointerToSymbolTable', + 'NumberOfSymbols', + 'SizeOfOptionalHeader', + 'Characteristics')) + @classmethod + def read_image_file_header(cls, file): + format = Struct('<HHIIIHH') + return cls._IMAGE_FILE_HEADER(*format.unpack(file.read(format.size))) + + _IMAGE_OPTIONAL_HEADER = namedtuple('_IMAGE_OPTIONAL_HEADER', + ('Magic', + 'MajorLinkerVersion', 'MinorLinkerVersion', + 'SizeOfCode', 'SizeOfInitializedData', + 'SizeOfUninitializedData', + 'AddressOfEntryPoint', 'BaseOfCode', + 'BaseOfData', 'ImageBase', + 'SectionAlignement', 'FileAlignement', + 'MajorOperatingSystemVersion', + 'MinorOperatingSystemVersion', + 'MajorImageVersion', + 'MinorImageVersion', + 'MajorSubsystemVersion', + 'MinorSubsystemVersion', + 'Win32VersionValue', + 'SizeOfImage', + 'SizeOfHeaders', + 'CheckSum', + 'Subsystem', + 'DllCharacteristics', + 'SizeOfStackReserve', + 'SizeOfStackCommit', + 'SizeOfHeapReserve', + 'SizeOfHeapCommit', + 'LoaderFlags', + 'NumberOfRvaAndSizes', + 'DataDirectory')) + _IMAGE_DATA_DIRECTORY = namedtuple('_IMAGE_DATA_DIRECTORY', + ('VirtualAddress', 'Size')) + @classmethod + def read_image_optional_header(cls, file): + format = Struct('<HBBIIIIIIIIIHHHHHHIIIIHHIIIIII') + directory_format = Struct('<II') + directory = [] + partial_header = format.unpack(file.read(format.size)) + directory = [cls._IMAGE_DATA_DIRECTORY(*directory_format.unpack(file.read(directory_format.size))) for i in xrange(16)] + return cls._IMAGE_OPTIONAL_HEADER(*(partial_header + (directory,))) + + _IMAGE_SECTION_HEADER = namedtuple('_IMAGE_SECTION_HEADER', + ('Name', 'VirtualSize', + 'VirtualAddress', + 'SizeOfRawData', 'PointerToRawData', + 'PointerToRelocations', + 'PointerToLinenumbers', + 'NumberOfRelocations', + 'NumberOfLinenumbers', + 'Characteristics')) + @classmethod + def read_image_section_header(cls, file): + format = Struct('<8sIIIIIIHHI') + return cls._IMAGE_SECTION_HEADER(*format.unpack(file.read(format.size))) + + + +class PEFile(object): + def __init__(self, file): + self.file = file + + self.image_base = 0 + self.sections = [] + + file.seek(0x3c) + pe_offset, = unpack('<I', file.read(4)) + + file.seek(pe_offset) + pe_sig = file.read(4) + assert pe_sig == b'PE\0\0' + + pe_file_header = PEStructs.read_image_file_header(file) + pe_optional_header = PEStructs.read_image_optional_header(file) + + # Read image base + self.image_base = pe_optional_header.ImageBase + + self.sections = [PEStructs.read_image_section_header(file) + for i in xrange(pe_file_header.NumberOfSections)] + + + def seek_to_va(self, va): + self.file.seek(self.va_to_offset(va)) + + + def offset_to_rva(self, offset): + for section in self.sections: + if 0 <= (offset - section.PointerToRawData) < section.SizeOfRawData: + #TODO: is that okay? + return offset - section.PointerToRawData + section.VirtualAddress + raise IndexError #TODO + + + def offset_to_va(self, offset): + return self.offset_to_rva(offset) + self.image_base + + + def rva_to_offset(self, rva): + for section in self.sections: + if 0 <= (rva - section.VirtualAddress) < section.SizeOfRawData: + #TODO: is that okay? + return rva - section.VirtualAddress + section.PointerToRawData + raise IndexError #TODO + + + def va_to_offset(self, va): + return self.rva_to_offset(va - self.image_base) +