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)
+