[MKHIVE] Minor improvements.
[reactos.git] / sdk / tools / gen_baseaddress.py
1 '''
2 PROJECT: ReactOS baseaddress updater
3 LICENSE: MIT (https://spdx.org/licenses/MIT)
4 PURPOSE: Update baseaddresses of all modules
5 COPYRIGHT: Copyright 2017,2018 Mark Jansen (mark.jansen@reactos.org)
6 '''
7 from __future__ import print_function, absolute_import, division
8 import os
9 import struct
10 import sys
11
12 try:
13 import pefile
14 except ImportError:
15 print('# Please install pefile from pip or https://github.com/erocarrera/pefile')
16 print('# Using fallback')
17 print()
18
19 ALL_EXTENSIONS = (
20 '.dll', '.acm', '.ax', '.cpl', '.drv', '.ocx'
21 )
22
23 PRIORITIES = (
24 'ntdll.dll',
25 'kernel32.dll',
26 'msvcrt.dll',
27 'advapi32.dll',
28 'gdi32.dll',
29 'user32.dll',
30 'dhcpcsvc.dll',
31 'dnsapi.dll',
32 'icmp.dll',
33 'iphlpapi.dll',
34 'ws2_32.dll',
35 'ws2help.dll',
36 'shlwapi.dll',
37 'rpcrt4.dll',
38 'comctl32.dll',
39 'ole32.dll',
40 'winspool.drv',
41 'winmm.dll',
42 'comdlg32.dll',
43 'shell32.dll',
44 'lz32.dll',
45 'version.dll',
46 'oleaut32.dll',
47 'setupapi.dll',
48 'mpr.dll',
49 'crypt32.dll',
50 'wininet.dll',
51 'urlmon.dll',
52 'psapi.dll',
53 'imm32.dll',
54 'msvfw32.dll',
55 'dbghelp.dll',
56 'devmgr.dll',
57 'msacm32.dll',
58 'netapi32.dll',
59 'powrprof.dll',
60 'secur32.dll',
61 'wintrust.dll',
62 'avicap32.dll',
63 'cabinet.dll',
64 'dsound.dll',
65 'glu32.dll',
66 'opengl32.dll',
67 'riched20.dll',
68 'smdll.dll',
69 'userenv.dll',
70 'uxtheme.dll',
71 'cryptui.dll',
72 'csrsrv.dll',
73 'basesrv.dll',
74 'winsrv.dll',
75 'dplayx.dll',
76 'gdiplus.dll',
77 'msimg32.dll',
78 'mswsock.dll',
79 'oledlg.dll',
80 'rasapi32.dll',
81 'rsaenh.dll',
82 'samlib.dll',
83 'sensapi.dll',
84 'sfc_os.dll',
85 'snmpapi.dll',
86 'spoolss.dll',
87 'usp10.dll',
88 )
89
90 EXCLUDE = (
91 'bmfd.dll',
92 'bootvid.dll',
93 'freeldr_pe.dll',
94 'ftfd.dll',
95 'fusion.dll',
96 'hal.dll',
97 'halaacpi.dll',
98 'halacpi.dll',
99 'halapic.dll',
100 'kbda1.dll',
101 'kbda2.dll',
102 'kbda3.dll',
103 'kbdal.dll',
104 'kbdarme.dll',
105 'kbdarmw.dll',
106 'kbdaze.dll',
107 'kbdazel.dll',
108 'kbdbe.dll',
109 'kbdbga.dll',
110 'kbdbgm.dll',
111 'kbdbgt.dll',
112 'kbdblr.dll',
113 'kbdbr.dll',
114 'kbdbur.dll',
115 'kbdcan.dll',
116 'kbdcr.dll',
117 'kbdcz.dll',
118 'kbdcz1.dll',
119 'kbdda.dll',
120 'kbddv.dll',
121 'kbdes.dll',
122 'kbdest.dll',
123 'kbdfc.dll',
124 'kbdfi.dll',
125 'kbdfr.dll',
126 'kbdgeo.dll',
127 'kbdgerg.dll',
128 'kbdgneo.dll',
129 'kbdgr.dll',
130 'kbdgrist.dll',
131 'kbdhe.dll',
132 'kbdheb.dll',
133 'kbdhu.dll',
134 'kbdic.dll',
135 'kbdinasa.dll',
136 'kbdinben.dll',
137 'kbdindev.dll',
138 'kbdinguj.dll',
139 'kbdinmal.dll',
140 'kbdir.dll',
141 'kbdit.dll',
142 'kbdja.dll',
143 'kbdkaz.dll',
144 'kbdko.dll',
145 'kbdla.dll',
146 'kbdlt1.dll',
147 'kbdlv.dll',
148 'kbdmac.dll',
149 'kbdne.dll',
150 'kbdno.dll',
151 'kbdpl.dll',
152 'kbdpl1.dll',
153 'kbdpo.dll',
154 'kbdro.dll',
155 'kbdru.dll',
156 'kbdru1.dll',
157 'kbdsg.dll',
158 'kbdsk.dll',
159 'kbdsk1.dll',
160 'kbdsw.dll',
161 'kbdtat.dll',
162 'kbdth0.dll',
163 'kbdth1.dll',
164 'kbdth2.dll',
165 'kbdth3.dll',
166 'kbdtuf.dll',
167 'kbdtuq.dll',
168 'kbduk.dll',
169 'kbdur.dll',
170 'kbdurs.dll',
171 'kbdus.dll',
172 'kbdusa.dll',
173 'kbdusl.dll',
174 'kbdusr.dll',
175 'kbdusx.dll',
176 'kbduzb.dll',
177 'kbdvntc.dll',
178 'kbdycc.dll',
179 'kbdycl.dll',
180 'kdcom.dll',
181 'kdvbox.dll',
182 'setupldr_pe.dll',
183 'vgaddi.dll',
184 'dllexport_test_dll1.dll',
185 'dllexport_test_dll2.dll',
186 'dllimport_test.dll',
187 'MyEventProvider.dll',
188 'w32kdll_2k3sp2.dll',
189 'w32kdll_ros.dll',
190 'w32kdll_xpsp2.dll',
191 )
192
193
194 def size_of_image_fallback(filename):
195 with open(filename, 'rb') as fin:
196 if fin.read(2) != 'MZ':
197 print(filename, 'No dos header found!')
198 return 0
199 fin.seek(0x3C)
200 e_lfanew = struct.unpack('i', fin.read(4))[0]
201 fin.seek(e_lfanew)
202 if fin.read(4) != 'PE\0\0':
203 print(filename, 'No PE header found!')
204 return 0
205 fin.seek(e_lfanew + 0x18)
206 pe_magic = struct.unpack('h', fin.read(2))[0]
207 if pe_magic != 0x10b:
208 print(filename, 'is not a 32 bit exe!')
209 return 0
210 fin.seek(e_lfanew + 0x50)
211 pe_size_of_image = struct.unpack('i', fin.read(4))[0]
212 return pe_size_of_image
213
214 def size_of_image(filename):
215 if 'pefile' in globals():
216 return pefile.PE(filename, fast_load=True).OPTIONAL_HEADER.SizeOfImage
217 return size_of_image_fallback(filename)
218
219
220 class Module(object):
221 def __init__(self, name, address, size):
222 self._name = name
223 self.address = address
224 self.size = size
225 self._reserved = address != 0
226
227 def gen_baseaddress(self):
228 name, ext = os.path.splitext(self._name)
229 postfix = ''
230 if ext in('.acm', '.drv') and self._name != 'winspool.drv':
231 name = self._name
232 if name == 'ntdll':
233 postfix = ' # should be above 0x%08x' % self.address
234 elif self._reserved:
235 postfix = ' # reserved'
236 print('set(baseaddress_%-30s 0x%08x)%s' % (name, self.address, postfix))
237
238 def end(self):
239 return self.address + self.size
240
241 def __repr__(self):
242 return '%s (0x%08x - 0x%08x)' % (self._name, self.address, self.end())
243
244 class MemoryLayout(object):
245 def __init__(self, startaddress):
246 self.addresses = []
247 self.found = {}
248 self.reserved = {}
249 self.initial = startaddress
250 self.start_at = 0
251 self.module_padding = 0x2000
252
253 def add_reserved(self, name, address):
254 self.reserved[name] = (address, 0)
255
256 def add(self, filename, name):
257 size = size_of_image(filename)
258 addr = 0
259 if name in self.found:
260 return # Assume duplicate files (rshell, ...) are 1:1 copies
261 if name in self.reserved:
262 addr = self.reserved[name][0]
263 self.reserved[name] = (addr, size)
264 self.found[name] = Module(name, addr, size)
265
266 def _next_address(self, size):
267 if self.start_at:
268 addr = (self.start_at - size - self.module_padding - 0xffff) & 0xffff0000
269 self.start_at = addr
270 else:
271 addr = self.start_at = self.initial
272 return addr
273
274 def next_address(self, size):
275 while True:
276 current_start = self._next_address(size)
277 current_end = current_start + size + self.module_padding
278 # Is there overlap with reserved modules?
279 for key, reserved in self.reserved.items():
280 res_start = reserved[0]
281 res_end = res_start + reserved[1] + self.module_padding
282 if (res_start <= current_start <= res_end) or \
283 (res_start <= current_end <= res_end) or \
284 (current_start < res_start and current_end > res_end):
285 # We passed this reserved item, we can remove it now
286 self.start_at = min(res_start, current_start)
287 del self.reserved[key]
288 current_start = 0
289 break
290 # No overlap with a reserved module?
291 if current_start:
292 return current_start
293
294 def update(self, priorities):
295 # sort addresses, should only contain reserved modules at this point!
296 for key, reserved in self.reserved.items():
297 assert reserved[1] != 0, key
298 for curr in priorities:
299 if not curr in self.found:
300 print('# Did not find', curr, '!')
301 else:
302 obj = self.found[curr]
303 del self.found[curr]
304 if not obj.address:
305 obj.address = self.next_address(obj.size)
306 self.addresses.append(obj)
307 # We handled all known modules now, run over the rest we found
308 for key in sorted(self.found):
309 obj = self.found[key]
310 obj.address = self.next_address(obj.size)
311 self.addresses.append(obj)
312
313 def gen_baseaddress(self):
314 for obj in self.addresses:
315 obj.gen_baseaddress()
316
317 def run_dir(target):
318 print('# Generated from', target)
319 print('# Generated by sdk/tools/gen_baseaddress.py')
320 layout = MemoryLayout(0x7c920000)
321 layout.add_reserved('user32.dll', 0x77a20000)
322 for root, _, files in os.walk(target):
323 for dll in [filename for filename in files if filename.endswith(ALL_EXTENSIONS)]:
324 if not dll in EXCLUDE and not dll.startswith('api-ms-win-'):
325 layout.add(os.path.join(root, dll), dll)
326 layout.update(PRIORITIES)
327 layout.gen_baseaddress()
328
329 def main(dirs):
330 if len(dirs) < 1:
331 trydir = os.getcwd()
332 print('# No path specified, trying', trydir)
333 dirs = [trydir]
334 for onedir in dirs:
335 run_dir(onedir)
336
337 if __name__ == '__main__':
338 main(sys.argv[1:])