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