Mercurial > touhou

comparison pytouhou/utils/pe.py @ 229:5afc75f71fed

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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
comparison
equal deleted inserted replaced
228:8f4cd1c01d22 229:5afc75f71fed
1 # -*- encoding: utf-8 -*-
2 ##
3 ## Copyright (C) 2011 Thibaut Girka <thib@sitedethib.com>
4 ##
5 ## This program is free software; you can redistribute it and/or modify
6 ## it under the terms of the GNU General Public License as published
7 ## by the Free Software Foundation; version 3 only.
8 ##
9 ## This program is distributed in the hope that it will be useful,
10 ## but WITHOUT ANY WARRANTY; without even the implied warranty of
11 ## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 ## GNU General Public License for more details.
13 ##
14
15 from struct import Struct, unpack
16 from collections import namedtuple
17
18
19 class PEStructs:
20 _IMAGE_FILE_HEADER = namedtuple('_IMAGE_FILE_HEADER',
21 ('Machine',
22 'NumberOfSections',
23 'TimeDateStamp',
24 'PointerToSymbolTable',
25 'NumberOfSymbols',
26 'SizeOfOptionalHeader',
27 'Characteristics'))
28 @classmethod
29 def read_image_file_header(cls, file):
30 format = Struct('<HHIIIHH')
31 return cls._IMAGE_FILE_HEADER(*format.unpack(file.read(format.size)))
32
33 _IMAGE_OPTIONAL_HEADER = namedtuple('_IMAGE_OPTIONAL_HEADER',
34 ('Magic',
35 'MajorLinkerVersion', 'MinorLinkerVersion',
36 'SizeOfCode', 'SizeOfInitializedData',
37 'SizeOfUninitializedData',
38 'AddressOfEntryPoint', 'BaseOfCode',
39 'BaseOfData', 'ImageBase',
40 'SectionAlignement', 'FileAlignement',
41 'MajorOperatingSystemVersion',
42 'MinorOperatingSystemVersion',
43 'MajorImageVersion',
44 'MinorImageVersion',
45 'MajorSubsystemVersion',
46 'MinorSubsystemVersion',
47 'Win32VersionValue',
48 'SizeOfImage',
49 'SizeOfHeaders',
50 'CheckSum',
51 'Subsystem',
52 'DllCharacteristics',
53 'SizeOfStackReserve',
54 'SizeOfStackCommit',
55 'SizeOfHeapReserve',
56 'SizeOfHeapCommit',
57 'LoaderFlags',
58 'NumberOfRvaAndSizes',
59 'DataDirectory'))
60 _IMAGE_DATA_DIRECTORY = namedtuple('_IMAGE_DATA_DIRECTORY',
61 ('VirtualAddress', 'Size'))
62 @classmethod
63 def read_image_optional_header(cls, file):
64 format = Struct('<HBBIIIIIIIIIHHHHHHIIIIHHIIIIII')
65 directory_format = Struct('<II')
66 directory = []
67 partial_header = format.unpack(file.read(format.size))
68 directory = [cls._IMAGE_DATA_DIRECTORY(*directory_format.unpack(file.read(directory_format.size))) for i in xrange(16)]
69 return cls._IMAGE_OPTIONAL_HEADER(*(partial_header + (directory,)))
70
71 _IMAGE_SECTION_HEADER = namedtuple('_IMAGE_SECTION_HEADER',
72 ('Name', 'VirtualSize',
73 'VirtualAddress',
74 'SizeOfRawData', 'PointerToRawData',
75 'PointerToRelocations',
76 'PointerToLinenumbers',
77 'NumberOfRelocations',
78 'NumberOfLinenumbers',
79 'Characteristics'))
80 @classmethod
81 def read_image_section_header(cls, file):
82 format = Struct('<8sIIIIIIHHI')
83 return cls._IMAGE_SECTION_HEADER(*format.unpack(file.read(format.size)))
84
85
86
87 class PEFile(object):
88 def __init__(self, file):
89 self.file = file
90
91 self.image_base = 0
92 self.sections = []
93
94 file.seek(0x3c)
95 pe_offset, = unpack('<I', file.read(4))
96
97 file.seek(pe_offset)
98 pe_sig = file.read(4)
99 assert pe_sig == b'PE\0\0'
100
101 pe_file_header = PEStructs.read_image_file_header(file)
102 pe_optional_header = PEStructs.read_image_optional_header(file)
103
104 # Read image base
105 self.image_base = pe_optional_header.ImageBase
106
107 self.sections = [PEStructs.read_image_section_header(file)
108 for i in xrange(pe_file_header.NumberOfSections)]
109
110
111 def seek_to_va(self, va):
112 self.file.seek(self.va_to_offset(va))
113
114
115 def offset_to_rva(self, offset):
116 for section in self.sections:
117 if 0 <= (offset - section.PointerToRawData) < section.SizeOfRawData:
118 #TODO: is that okay?
119 return offset - section.PointerToRawData + section.VirtualAddress
120 raise IndexError #TODO
121
122
123 def offset_to_va(self, offset):
124 return self.offset_to_rva(offset) + self.image_base
125
126
127 def rva_to_offset(self, rva):
128 for section in self.sections:
129 if 0 <= (rva - section.VirtualAddress) < section.SizeOfRawData:
130 #TODO: is that okay?
131 return rva - section.VirtualAddress + section.PointerToRawData
132 raise IndexError #TODO
133
134
135 def va_to_offset(self, va):
136 return self.rva_to_offset(va - self.image_base)
137