xref: /reactos/sdk/tools/gen_baseaddress.py (revision dc4992ce) 
1'''
2PROJECT:     ReactOS baseaddress updater
3LICENSE:     MIT (https://spdx.org/licenses/MIT)
4PURPOSE:     Update baseaddresses of all modules
5COPYRIGHT:   Copyright 2017,2018 Mark Jansen (mark.jansen@reactos.org)
6'''
7
8from __future__ import print_function, absolute_import, division
9
10USAGE = """
11This script will update the baseaddresses of all modules, based on the build output.
12
13Specify the build output dir as commandline argument to the script:
14`python gen_baseaddress.py C:\\Users\\Mark\\reactos\\output-MinGW-i386`
15
16Multiple directories can be specified:
17`python gen_baseaddress r:/build/msvc r:/build/gcc`
18
19Specify -64 as first argument to use 64-bit addresses:
20`python gen_baseaddress -64 r:\\build\\msvc-x64`
21"""
22
23import os
24import struct
25import sys
26
27try:
28    import pefile
29except ImportError:
30    print('# Please install pefile from pip or https://github.com/erocarrera/pefile')
31    sys.exit(-1)
32
33
34ALL_EXTENSIONS = (
35    '.dll', '.acm', '.ax', '.cpl', '.drv', '.ocx', '.ime'
36)
37
38PRIORITIES = (
39    'ntdll.dll',
40    'kernel32.dll',
41    'msvcrt.dll',
42    'advapi32.dll',
43    'gdi32.dll',
44    'user32.dll',
45    'dhcpcsvc.dll',
46    'dnsapi.dll',
47    'icmp.dll',
48    'iphlpapi.dll',
49    'ws2_32.dll',
50    'ws2help.dll',
51    'shlwapi.dll',
52    'rpcrt4.dll',
53    'comctl32.dll',
54    'ole32.dll',
55    'winspool.drv',
56    'winmm.dll',
57    'comdlg32.dll',
58    'shell32.dll',
59    'lz32.dll',
60    'version.dll',
61    'oleaut32.dll',
62    'setupapi.dll',
63    'mpr.dll',
64    'crypt32.dll',
65    'wininet.dll',
66    'urlmon.dll',
67    'psapi.dll',
68    'imm32.dll',
69    'msvfw32.dll',
70    'dbghelp.dll',
71    'devmgr.dll',
72    'msacm32.dll',
73    'netapi32.dll',
74    'powrprof.dll',
75    'secur32.dll',
76    'wintrust.dll',
77    'avicap32.dll',
78    'cabinet.dll',
79    'dsound.dll',
80    'glu32.dll',
81    'opengl32.dll',
82    'riched20.dll',
83    'userenv.dll',
84    'uxtheme.dll',
85    'cryptui.dll',
86    'csrsrv.dll',
87    'basesrv.dll',
88    'winsrv.dll',
89    'dplayx.dll',
90    'gdiplus.dll',
91    'msimg32.dll',
92    'mswsock.dll',
93    'oledlg.dll',
94    'rasapi32.dll',
95    'rsaenh.dll',
96    'samlib.dll',
97    'sensapi.dll',
98    'sfc_os.dll',
99    'snmpapi.dll',
100    'spoolss.dll',
101    'usp10.dll',
102)
103
104EXCLUDE = (
105    'bmfd.dll',
106    'bootvid.dll',
107    'framebuf.dll',
108    'ftfd.dll',
109    'genincdata.dll',
110    'hal.dll',
111    'halaacpi.dll',
112    'halacpi.dll',
113    'halapic.dll',
114    'kbda1.dll',
115    'kbda2.dll',
116    'kbda3.dll',
117    'kbdal.dll',
118    'kbdarme.dll',
119    'kbdarmw.dll',
120    'kbdaze.dll',
121    'kbdazel.dll',
122    'kbdbe.dll',
123    'kbdbga.dll',
124    'kbdbgm.dll',
125    'kbdbgt.dll',
126    'kbdblr.dll',
127    'kbdbr.dll',
128    'kbdbu.dll',
129    'kbdbur.dll',
130    'kbdcan.dll',
131    'kbdcr.dll',
132    'kbdcz.dll',
133    'kbdcz1.dll',
134    'kbdda.dll',
135    'kbddv.dll',
136    'kbdeo.dll',
137    'kbdes.dll',
138    'kbdest.dll',
139    'kbdfc.dll',
140    'kbdfi.dll',
141    'kbdfr.dll',
142    'kbdgeo.dll',
143    'kbdgerg.dll',
144    'kbdgneo.dll',
145    'kbdgr.dll',
146    'kbdgr1.dll',
147    'kbdgrist.dll',
148    'kbdhe.dll',
149    'kbdheb.dll',
150    'kbdhu.dll',
151    'kbdic.dll',
152    'kbdinasa.dll',
153    'kbdinben.dll',
154    'kbdindev.dll',
155    'kbdinguj.dll',
156    'kbdinmal.dll',
157    'kbdir.dll',
158    'kbdit.dll',
159    'kbdja.dll',
160    'kbdjpn.dll',
161    'kbdkaz.dll',
162    'kbdko.dll',
163    'kbdkor.dll',
164    'kbdla.dll',
165    'kbdlt1.dll',
166    'kbdlv.dll',
167    'kbdmac.dll',
168    'kbdne.dll',
169    'kbdno.dll',
170    'kbdpl.dll',
171    'kbdpl1.dll',
172    'kbdpo.dll',
173    'kbdro.dll',
174    'kbdrost.dll',
175    'kbdru.dll',
176    'kbdru1.dll',
177    'kbdsf.dll',
178    'kbdsg.dll',
179    'kbdsk.dll',
180    'kbdsk1.dll',
181    'kbdsl.dll',
182    'kbdsl1.dll',
183    'kbdsp.dll',
184    'kbdsw.dll',
185    'kbdtat.dll',
186    'kbdth0.dll',
187    'kbdth1.dll',
188    'kbdth2.dll',
189    'kbdth3.dll',
190    'kbdtuf.dll',
191    'kbdtuq.dll',
192    'kbduk.dll',
193    'kbdur.dll',
194    'kbdurs.dll',
195    'kbdus.dll',
196    'kbdusa.dll',
197    'kbdusl.dll',
198    'kbdusr.dll',
199    'kbdusx.dll',
200    'kbduzb.dll',
201    'kbdvntc.dll',
202    'kbdycc.dll',
203    'kbdycl.dll',
204    'kdcom.dll',
205    'kdvbox.dll',
206    'vgaddi.dll',
207    'dllexport_test_dll1.dll',
208    'dllexport_test_dll2.dll',
209    'dllimport_test.dll',
210    'localspl_apitest.dll',
211    'MyEventProvider.dll',
212    'redirtest1.dll',
213    'redirtest2.dll',
214    'win32u_2k3sp2.dll',
215    'win32u_vista.dll',
216    'win32u_xpsp2.dll',
217    'testvdd.dll',
218)
219
220IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x10b
221IMAGE_NT_OPTIONAL_HDR64_MAGIC = 0x20b
222
223IMAGE_TYPES = {
224    IMAGE_NT_OPTIONAL_HDR32_MAGIC: 0,
225    IMAGE_NT_OPTIONAL_HDR64_MAGIC: 0
226}
227
228IS_64_BIT = False
229
230def size_of_image(filename):
231    if IS_64_BIT:
232        return 0xFE0000 # This results in 0x1000000 / 16MB space
233    with open(filename, 'rb') as fin:
234        if fin.read(2) != b'MZ':
235            print(filename, 'No dos header found!')
236            return 0
237        fin.seek(0x3C)
238        e_lfanew = struct.unpack('i', fin.read(4))[0]
239        fin.seek(e_lfanew)
240        if fin.read(4) != b'PE\0\0':
241            print(filename, 'No PE header found!')
242            return 0
243        fin.seek(e_lfanew + 0x18)
244        pe_magic = struct.unpack('h', fin.read(2))[0]
245        if pe_magic in IMAGE_TYPES.keys():
246            IMAGE_TYPES[pe_magic] += 1
247            fin.seek(e_lfanew + 0x50)
248            pe_size_of_image = struct.unpack('i', fin.read(4))[0]
249            return pe_size_of_image
250        print(filename, 'Unknown executable format!')
251        return 0
252
253
254class Module(object):
255    def __init__(self, name, address, size, filename):
256        self._name = name
257        self.address = address
258        self.size = size
259        self._reserved = address != 0
260        self.filename = filename
261
262    def gen_baseaddress(self, output_file):
263        name, ext = os.path.splitext(self._name)
264        postfix = ''
265        if ext in('.acm', '.drv') and self._name != 'winspool.drv':
266            name = self._name
267        if name == 'ntdll':
268            postfix = ' # should be above 0x%08x' % self.address
269        elif self._reserved:
270            postfix = ' # reserved'
271        output_file.write('set(baseaddress_%-30s 0x%08x)%s\n' % (name, self.address, postfix))
272
273    def end(self):
274        return self.address + self.size
275
276    def __repr__(self):
277        return '%s (0x%08x - 0x%08x)' % (self._name, self.address, self.end())
278
279class MemoryLayout(object):
280    def __init__(self, startaddress):
281        self.addresses = []
282        self.found = {}
283        self.reserved = {}
284        self.initial = startaddress
285        self.start_at = 0
286        self.module_padding = 0x2000
287
288    def add_reserved(self, name, address):
289        self.reserved[name] = (address, 0)
290
291    def add(self, filename, name):
292        size = size_of_image(filename)
293        addr = 0
294        if name in self.found:
295            return  # Assume duplicate files (rshell, ...) are 1:1 copies
296        if name in self.reserved:
297            addr = self.reserved[name][0]
298            self.reserved[name] = (addr, size)
299        self.found[name] = Module(name, addr, size, filename)
300
301    def _next_address(self, size):
302        if self.start_at:
303            addr = (self.start_at - size - self.module_padding - 0xffff) & 0xffffffffffff0000
304            self.start_at = addr
305        else:
306            addr = self.start_at = self.initial
307        return addr
308
309    def next_address(self, size):
310        while True:
311            current_start = self._next_address(size)
312            current_end = current_start + size + self.module_padding
313            # Is there overlap with reserved modules?
314            for key, reserved in self.reserved.items():
315                res_start = reserved[0]
316                res_end = res_start + reserved[1] + self.module_padding
317                if (res_start <= current_start <= res_end) or \
318                   (res_start <= current_end <= res_end) or \
319                   (current_start < res_start and current_end > res_end):
320                    # We passed this reserved item, we can remove it now
321                    self.start_at = min(res_start, current_start)
322                    del self.reserved[key]
323                    current_start = 0
324                    break
325            # No overlap with a reserved module?
326            if current_start:
327                return current_start
328
329    def update(self, priorities):
330        # sort addresses, should only contain reserved modules at this point!
331        for key, reserved in self.reserved.items():
332            assert reserved[1] != 0, key
333        for curr in priorities:
334            if not curr in self.found:
335                print('# Did not find', curr, '!')
336            else:
337                obj = self.found[curr]
338                del self.found[curr]
339                if not obj.address:
340                    obj.address = self.next_address(obj.size)
341                self.addresses.append(obj)
342        # We handled all known modules now, run over the rest we found
343        for key in sorted(self.found):
344            obj = self.found[key]
345            obj.address = self.next_address(obj.size)
346            self.addresses.append(obj)
347
348    def gen_baseaddress(self, output_file):
349        for obj in self.addresses:
350            obj.gen_baseaddress(output_file)
351
352def get_target_file(ntdll_path):
353    if 'pefile' in globals():
354        ntdll_pe = pefile.PE(ntdll_path, fast_load=True)
355        names = [sect.Name.strip(b'\0') for sect in ntdll_pe.sections]
356        count = b'|'.join(names).count(b'/')
357        if IS_64_BIT:
358            return 'baseaddress64.cmake'
359        elif b'.rossym' in names:
360            return 'baseaddress.cmake'
361        elif count == 0:
362            return 'baseaddress_msvc.cmake'
363        elif count > 3:
364            return 'baseaddress_dwarf.cmake'
365        else:
366            assert False, "Unknown"
367    return None
368
369def run_dir(target):
370    if IS_64_BIT:
371        layout = MemoryLayout(0x7FFB7000000)
372    else:
373        layout = MemoryLayout(0x7c920000)
374        layout.add_reserved('user32.dll', 0x77a20000)
375    IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] = 0
376    IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC] = 0
377    for root, _, files in os.walk(target):
378        for dll in [filename for filename in files if filename.endswith(ALL_EXTENSIONS)]:
379            if not dll in EXCLUDE and not dll.startswith('api-ms-win-'):
380                layout.add(os.path.join(root, dll), dll)
381    ntdll_path = layout.found['ntdll.dll'].filename
382    target_file = get_target_file(ntdll_path)
383    if target_file:
384        target_dir = os.path.realpath(os.path.dirname(os.path.dirname(__file__)))
385        target_path = os.path.join(target_dir, 'cmake', target_file)
386        output_file = open(target_path, "w")
387    else:
388        output_file = sys.stdout
389    with output_file:
390        output_file.write('# Generated from {}\n'.format(target))
391        output_file.write('# Generated by sdk/tools/gen_baseaddress.py\n\n')
392        layout.update(PRIORITIES)
393        layout.gen_baseaddress(output_file)
394
395def main():
396    dirs = sys.argv[1:]
397
398    if len(dirs) > 0 and dirs[0] == '-64':
399        print('Using 64-bit addresses')
400        global IS_64_BIT
401        IS_64_BIT = True
402        dirs = sys.argv[2:]
403    if len(dirs) < 1:
404        trydir = os.getcwd()
405        print(USAGE)
406        print('No path specified, trying the working directory: ', trydir)
407        dirs = [trydir]
408    for onedir in dirs:
409        if onedir.lower() in ['-help', '/help', '/h', '-h', '/?', '-?']:
410            print(USAGE)
411        else:
412            run_dir(onedir)
413
414
415if __name__ == '__main__':
416    main()
417