xref: /reactos/sdk/tools/gen_baseaddress.py (revision 177ae91b) 
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    'smdll.dll',
81    'userenv.dll',
82    'uxtheme.dll',
83    'cryptui.dll',
84    'csrsrv.dll',
85    'basesrv.dll',
86    'winsrv.dll',
87    'dplayx.dll',
88    'gdiplus.dll',
89    'msimg32.dll',
90    'mswsock.dll',
91    'oledlg.dll',
92    'rasapi32.dll',
93    'rsaenh.dll',
94    'samlib.dll',
95    'sensapi.dll',
96    'sfc_os.dll',
97    'snmpapi.dll',
98    'spoolss.dll',
99    'usp10.dll',
100)
101
102EXCLUDE = (
103    'bmfd.dll',
104    'bootvid.dll',
105    'framebuf.dll',
106    'framebuf_new.dll',
107    'ftfd.dll',
108    'fusion.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    'kbdbur.dll',
129    'kbdcan.dll',
130    'kbdcr.dll',
131    'kbdcz.dll',
132    'kbdcz1.dll',
133    'kbdda.dll',
134    'kbddv.dll',
135    'kbdes.dll',
136    'kbdest.dll',
137    'kbdfc.dll',
138    'kbdfi.dll',
139    'kbdfr.dll',
140    'kbdgeo.dll',
141    'kbdgerg.dll',
142    'kbdgneo.dll',
143    'kbdgr.dll',
144    'kbdgrist.dll',
145    'kbdhe.dll',
146    'kbdheb.dll',
147    'kbdhu.dll',
148    'kbdic.dll',
149    'kbdinasa.dll',
150    'kbdinben.dll',
151    'kbdindev.dll',
152    'kbdinguj.dll',
153    'kbdinmal.dll',
154    'kbdir.dll',
155    'kbdit.dll',
156    'kbdja.dll',
157    'kbdkaz.dll',
158    'kbdko.dll',
159    'kbdla.dll',
160    'kbdlt1.dll',
161    'kbdlv.dll',
162    'kbdmac.dll',
163    'kbdne.dll',
164    'kbdno.dll',
165    'kbdpl.dll',
166    'kbdpl1.dll',
167    'kbdpo.dll',
168    'kbdro.dll',
169    'kbdru.dll',
170    'kbdru1.dll',
171    'kbdsg.dll',
172    'kbdsk.dll',
173    'kbdsk1.dll',
174    'kbdsw.dll',
175    'kbdtat.dll',
176    'kbdth0.dll',
177    'kbdth1.dll',
178    'kbdth2.dll',
179    'kbdth3.dll',
180    'kbdtuf.dll',
181    'kbdtuq.dll',
182    'kbduk.dll',
183    'kbdur.dll',
184    'kbdurs.dll',
185    'kbdus.dll',
186    'kbdusa.dll',
187    'kbdusl.dll',
188    'kbdusr.dll',
189    'kbdusx.dll',
190    'kbduzb.dll',
191    'kbdvntc.dll',
192    'kbdycc.dll',
193    'kbdycl.dll',
194    'kdcom.dll',
195    'kdvbox.dll',
196    'vgaddi.dll',
197    'dllexport_test_dll1.dll',
198    'dllexport_test_dll2.dll',
199    'dllimport_test.dll',
200    'MyEventProvider.dll',
201    'w32kdll_2k3sp2.dll',
202    'w32kdll_ros.dll',
203    'w32kdll_xpsp2.dll',
204)
205
206IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x10b
207IMAGE_NT_OPTIONAL_HDR64_MAGIC = 0x20b
208
209IMAGE_TYPES = {
210    IMAGE_NT_OPTIONAL_HDR32_MAGIC: 0,
211    IMAGE_NT_OPTIONAL_HDR64_MAGIC: 0
212}
213
214def is_x64():
215    return IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] > IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC]
216
217def size_of_image(filename):
218    with open(filename, 'rb') as fin:
219        if fin.read(2) != b'MZ':
220            print(filename, 'No dos header found!')
221            return 0
222        fin.seek(0x3C)
223        e_lfanew = struct.unpack('i', fin.read(4))[0]
224        fin.seek(e_lfanew)
225        if fin.read(4) != b'PE\0\0':
226            print(filename, 'No PE header found!')
227            return 0
228        fin.seek(e_lfanew + 0x18)
229        pe_magic = struct.unpack('h', fin.read(2))[0]
230        if pe_magic in IMAGE_TYPES.keys():
231            IMAGE_TYPES[pe_magic] += 1
232            fin.seek(e_lfanew + 0x50)
233            pe_size_of_image = struct.unpack('i', fin.read(4))[0]
234            return pe_size_of_image
235        print(filename, 'Unknown executable format!')
236        return 0
237
238
239class Module(object):
240    def __init__(self, name, address, size, filename):
241        self._name = name
242        self.address = address
243        self.size = size
244        self._reserved = address != 0
245        self.filename = filename
246
247    def gen_baseaddress(self, output_file):
248        name, ext = os.path.splitext(self._name)
249        postfix = ''
250        if ext in('.acm', '.drv') and self._name != 'winspool.drv':
251            name = self._name
252        if name == 'ntdll':
253            postfix = ' # should be above 0x%08x' % self.address
254        elif self._reserved:
255            postfix = ' # reserved'
256        output_file.write('set(baseaddress_%-30s 0x%08x)%s\n' % (name, self.address, postfix))
257
258    def end(self):
259        return self.address + self.size
260
261    def __repr__(self):
262        return '%s (0x%08x - 0x%08x)' % (self._name, self.address, self.end())
263
264class MemoryLayout(object):
265    def __init__(self, startaddress):
266        self.addresses = []
267        self.found = {}
268        self.reserved = {}
269        self.initial = startaddress
270        self.start_at = 0
271        self.module_padding = 0x2000
272
273    def add_reserved(self, name, address):
274        self.reserved[name] = (address, 0)
275
276    def add(self, filename, name):
277        size = size_of_image(filename)
278        addr = 0
279        if name in self.found:
280            return  # Assume duplicate files (rshell, ...) are 1:1 copies
281        if name in self.reserved:
282            addr = self.reserved[name][0]
283            self.reserved[name] = (addr, size)
284        self.found[name] = Module(name, addr, size, filename)
285
286    def _next_address(self, size):
287        if self.start_at:
288            addr = (self.start_at - size - self.module_padding - 0xffff) & 0xffff0000
289            self.start_at = addr
290        else:
291            addr = self.start_at = self.initial
292        return addr
293
294    def next_address(self, size):
295        while True:
296            current_start = self._next_address(size)
297            current_end = current_start + size + self.module_padding
298            # Is there overlap with reserved modules?
299            for key, reserved in self.reserved.items():
300                res_start = reserved[0]
301                res_end = res_start + reserved[1] + self.module_padding
302                if (res_start <= current_start <= res_end) or \
303                   (res_start <= current_end <= res_end) or \
304                   (current_start < res_start and current_end > res_end):
305                    # We passed this reserved item, we can remove it now
306                    self.start_at = min(res_start, current_start)
307                    del self.reserved[key]
308                    current_start = 0
309                    break
310            # No overlap with a reserved module?
311            if current_start:
312                return current_start
313
314    def update(self, priorities):
315        # sort addresses, should only contain reserved modules at this point!
316        for key, reserved in self.reserved.items():
317            assert reserved[1] != 0, key
318        for curr in priorities:
319            if not curr in self.found:
320                print('# Did not find', curr, '!')
321            else:
322                obj = self.found[curr]
323                del self.found[curr]
324                if not obj.address:
325                    obj.address = self.next_address(obj.size)
326                self.addresses.append(obj)
327        # We handled all known modules now, run over the rest we found
328        for key in sorted(self.found):
329            obj = self.found[key]
330            obj.address = self.next_address(obj.size)
331            self.addresses.append(obj)
332
333    def gen_baseaddress(self, output_file):
334        for obj in self.addresses:
335            obj.gen_baseaddress(output_file)
336
337def get_target_file(ntdll_path):
338    if 'pefile' in globals():
339        ntdll_pe = pefile.PE(ntdll_path, fast_load=True)
340        names = [sect.Name.strip(b'\0') for sect in ntdll_pe.sections]
341        count = b'|'.join(names).count(b'/')
342        if b'.rossym' in names:
343            return 'baseaddress.cmake'
344        elif is_x64():
345            return 'baseaddress_msvc_x64.cmake'
346        elif count == 0:
347            return 'baseaddress_msvc.cmake'
348        elif count > 3:
349            return 'baseaddress_dwarf.cmake'
350        else:
351            assert False, "Unknown"
352    return None
353
354def run_dir(target):
355    layout = MemoryLayout(0x7c920000)
356    layout.add_reserved('user32.dll', 0x77a20000)
357    IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] = 0
358    IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC] = 0
359    for root, _, files in os.walk(target):
360        for dll in [filename for filename in files if filename.endswith(ALL_EXTENSIONS)]:
361            if not dll in EXCLUDE and not dll.startswith('api-ms-win-'):
362                layout.add(os.path.join(root, dll), dll)
363    ntdll_path = layout.found['ntdll.dll'].filename
364    target_file = get_target_file(ntdll_path)
365    if target_file:
366        target_dir = os.path.realpath(os.path.dirname(os.path.dirname(__file__)))
367        target_path = os.path.join(target_dir, 'cmake', target_file)
368        output_file = open(target_path, "w")
369    else:
370        output_file = sys.stdout
371    with output_file:
372        output_file.write('# Generated from {}\n'.format(target))
373        output_file.write('# Generated by sdk/tools/gen_baseaddress.py\n\n')
374        layout.update(PRIORITIES)
375        layout.gen_baseaddress(output_file)
376
377def main():
378    dirs = sys.argv[1:]
379    if len(dirs) < 1:
380        trydir = os.getcwd()
381        print(USAGE)
382        print('No path specified, trying the working directory: ', trydir)
383        dirs = [trydir]
384    for onedir in dirs:
385        if onedir.lower() in ['-help', '/help', '/h', '-h', '/?', '-?']:
386            print(USAGE)
387        else:
388            run_dir(onedir)
389
390
391if __name__ == '__main__':
392    main()
393