xref: /reactos/sdk/tools/gen_baseaddress.py (revision ebaf247c)
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'''
7from __future__ import print_function, absolute_import, division
8import os
9import struct
10import sys
11
12try:
13    import pefile
14except ImportError:
15    print('# Please install pefile from pip or https://github.com/erocarrera/pefile')
16    print('# Using fallback')
17    print()
18
19ALL_EXTENSIONS = (
20    '.dll', '.acm', '.ax', '.cpl', '.drv', '.ocx'
21)
22
23PRIORITIES = (
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
90EXCLUDE = (
91    'bmfd.dll',
92    'bootvid.dll',
93    'framebuf.dll',
94    'framebuf_new.dll',
95    'freeldr_pe.dll',
96    'ftfd.dll',
97    'fusion.dll',
98    'genincdata.dll',
99    'hal.dll',
100    'halaacpi.dll',
101    'halacpi.dll',
102    'halapic.dll',
103    'kbda1.dll',
104    'kbda2.dll',
105    'kbda3.dll',
106    'kbdal.dll',
107    'kbdarme.dll',
108    'kbdarmw.dll',
109    'kbdaze.dll',
110    'kbdazel.dll',
111    'kbdbe.dll',
112    'kbdbga.dll',
113    'kbdbgm.dll',
114    'kbdbgt.dll',
115    'kbdblr.dll',
116    'kbdbr.dll',
117    'kbdbur.dll',
118    'kbdcan.dll',
119    'kbdcr.dll',
120    'kbdcz.dll',
121    'kbdcz1.dll',
122    'kbdda.dll',
123    'kbddv.dll',
124    'kbdes.dll',
125    'kbdest.dll',
126    'kbdfc.dll',
127    'kbdfi.dll',
128    'kbdfr.dll',
129    'kbdgeo.dll',
130    'kbdgerg.dll',
131    'kbdgneo.dll',
132    'kbdgr.dll',
133    'kbdgrist.dll',
134    'kbdhe.dll',
135    'kbdheb.dll',
136    'kbdhu.dll',
137    'kbdic.dll',
138    'kbdinasa.dll',
139    'kbdinben.dll',
140    'kbdindev.dll',
141    'kbdinguj.dll',
142    'kbdinmal.dll',
143    'kbdir.dll',
144    'kbdit.dll',
145    'kbdja.dll',
146    'kbdkaz.dll',
147    'kbdko.dll',
148    'kbdla.dll',
149    'kbdlt1.dll',
150    'kbdlv.dll',
151    'kbdmac.dll',
152    'kbdne.dll',
153    'kbdno.dll',
154    'kbdpl.dll',
155    'kbdpl1.dll',
156    'kbdpo.dll',
157    'kbdro.dll',
158    'kbdru.dll',
159    'kbdru1.dll',
160    'kbdsg.dll',
161    'kbdsk.dll',
162    'kbdsk1.dll',
163    'kbdsw.dll',
164    'kbdtat.dll',
165    'kbdth0.dll',
166    'kbdth1.dll',
167    'kbdth2.dll',
168    'kbdth3.dll',
169    'kbdtuf.dll',
170    'kbdtuq.dll',
171    'kbduk.dll',
172    'kbdur.dll',
173    'kbdurs.dll',
174    'kbdus.dll',
175    'kbdusa.dll',
176    'kbdusl.dll',
177    'kbdusr.dll',
178    'kbdusx.dll',
179    'kbduzb.dll',
180    'kbdvntc.dll',
181    'kbdycc.dll',
182    'kbdycl.dll',
183    'kdcom.dll',
184    'kdvbox.dll',
185    'setupldr_pe.dll',
186    'vgaddi.dll',
187    'dllexport_test_dll1.dll',
188    'dllexport_test_dll2.dll',
189    'dllimport_test.dll',
190    'MyEventProvider.dll',
191    'w32kdll_2k3sp2.dll',
192    'w32kdll_ros.dll',
193    'w32kdll_xpsp2.dll',
194)
195
196IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x10b
197IMAGE_NT_OPTIONAL_HDR64_MAGIC = 0x20b
198
199IMAGE_TYPES = {
200    IMAGE_NT_OPTIONAL_HDR32_MAGIC: 0,
201    IMAGE_NT_OPTIONAL_HDR64_MAGIC: 0
202}
203
204def is_x64():
205    return IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] > IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC]
206
207def size_of_image_fallback(filename):
208    with open(filename, 'rb') as fin:
209        if fin.read(2) != b'MZ':
210            print(filename, 'No dos header found!')
211            return 0
212        fin.seek(0x3C)
213        e_lfanew = struct.unpack('i', fin.read(4))[0]
214        fin.seek(e_lfanew)
215        if fin.read(4) != b'PE\0\0':
216            print(filename, 'No PE header found!')
217            return 0
218        fin.seek(e_lfanew + 0x18)
219        pe_magic = struct.unpack('h', fin.read(2))[0]
220        if pe_magic in IMAGE_TYPES.keys():
221            IMAGE_TYPES[pe_magic] += 1
222            fin.seek(e_lfanew + 0x50)
223            pe_size_of_image = struct.unpack('i', fin.read(4))[0]
224            return pe_size_of_image
225        print(filename, 'Unknown executable format!')
226        return 0
227
228
229def size_of_image_verify(filename):
230    pefile_size = pefile.PE(filename, fast_load=True).OPTIONAL_HEADER.SizeOfImage
231    custom_size = size_of_image_fallback(filename)
232    assert custom_size == pefile_size, filename
233    return custom_size
234
235SIZE_OF_IMAGE_FN = size_of_image_fallback
236
237class Module(object):
238    def __init__(self, name, address, size, filename):
239        self._name = name
240        self.address = address
241        self.size = size
242        self._reserved = address != 0
243        self.filename = filename
244
245    def gen_baseaddress(self):
246        name, ext = os.path.splitext(self._name)
247        postfix = ''
248        if ext in('.acm', '.drv') and self._name != 'winspool.drv':
249            name = self._name
250        if name == 'ntdll':
251            postfix = ' # should be above 0x%08x' % self.address
252        elif self._reserved:
253            postfix = ' # reserved'
254        print('set(baseaddress_%-30s 0x%08x)%s' % (name, self.address, postfix))
255
256    def end(self):
257        return self.address + self.size
258
259    def __repr__(self):
260        return '%s (0x%08x - 0x%08x)' % (self._name, self.address, self.end())
261
262class MemoryLayout(object):
263    def __init__(self, startaddress):
264        self.addresses = []
265        self.found = {}
266        self.reserved = {}
267        self.initial = startaddress
268        self.start_at = 0
269        self.module_padding = 0x2000
270
271    def add_reserved(self, name, address):
272        self.reserved[name] = (address, 0)
273
274    def add(self, filename, name):
275        size = SIZE_OF_IMAGE_FN(filename)
276        addr = 0
277        if name in self.found:
278            return  # Assume duplicate files (rshell, ...) are 1:1 copies
279        if name in self.reserved:
280            addr = self.reserved[name][0]
281            self.reserved[name] = (addr, size)
282        self.found[name] = Module(name, addr, size, filename)
283
284    def _next_address(self, size):
285        if self.start_at:
286            addr = (self.start_at - size - self.module_padding - 0xffff) & 0xffff0000
287            self.start_at = addr
288        else:
289            addr = self.start_at = self.initial
290        return addr
291
292    def next_address(self, size):
293        while True:
294            current_start = self._next_address(size)
295            current_end = current_start + size + self.module_padding
296            # Is there overlap with reserved modules?
297            for key, reserved in self.reserved.items():
298                res_start = reserved[0]
299                res_end = res_start + reserved[1] + self.module_padding
300                if (res_start <= current_start <= res_end) or \
301                   (res_start <= current_end <= res_end) or \
302                   (current_start < res_start and current_end > res_end):
303                    # We passed this reserved item, we can remove it now
304                    self.start_at = min(res_start, current_start)
305                    del self.reserved[key]
306                    current_start = 0
307                    break
308            # No overlap with a reserved module?
309            if current_start:
310                return current_start
311
312    def update(self, priorities):
313        # sort addresses, should only contain reserved modules at this point!
314        for key, reserved in self.reserved.items():
315            assert reserved[1] != 0, key
316        for curr in priorities:
317            if not curr in self.found:
318                print('# Did not find', curr, '!')
319            else:
320                obj = self.found[curr]
321                del self.found[curr]
322                if not obj.address:
323                    obj.address = self.next_address(obj.size)
324                self.addresses.append(obj)
325        # We handled all known modules now, run over the rest we found
326        for key in sorted(self.found):
327            obj = self.found[key]
328            obj.address = self.next_address(obj.size)
329            self.addresses.append(obj)
330
331    def gen_baseaddress(self):
332        for obj in self.addresses:
333            obj.gen_baseaddress()
334
335def guess_version(ntdll_path):
336    if 'pefile' in globals():
337        ntdll_pe = pefile.PE(ntdll_path, fast_load=True)
338        names = [sect.Name.strip(b'\0') for sect in ntdll_pe.sections]
339        count = b'|'.join(names).count(b'/')
340        if '.rossym' in names:
341            print('# This should probably go in baseaddress.cmake')
342        elif is_x64():
343            print('# This should probably go in baseaddress_msvc_x64.cmake')
344        elif count == 0:
345            print('# This should probably go in baseaddress_msvc.cmake')
346        elif count > 3:
347            print('# This should probably go in baseaddress_dwarf.cmake')
348        else:
349            print('# No clue where to put this')
350
351def run_dir(target):
352    print('# Generated from', target)
353    print('# Generated by sdk/tools/gen_baseaddress.py')
354    layout = MemoryLayout(0x7c920000)
355    layout.add_reserved('user32.dll', 0x77a20000)
356    for root, _, files in os.walk(target):
357        for dll in [filename for filename in files if filename.endswith(ALL_EXTENSIONS)]:
358            if not dll in EXCLUDE and not dll.startswith('api-ms-win-'):
359                layout.add(os.path.join(root, dll), dll)
360    ntdll_path = layout.found['ntdll.dll'].filename
361    guess_version(ntdll_path)
362    layout.update(PRIORITIES)
363    layout.gen_baseaddress()
364
365def main():
366    dirs = sys.argv[1:]
367    if len(dirs) < 1:
368        trydir = os.getcwd()
369        print('# No path specified, trying', trydir)
370        dirs = [trydir]
371    for onedir in dirs:
372        run_dir(onedir)
373
374
375def profile():
376    import cProfile
377    # pyprof2calltree -k -i test.cprof
378    cProfile.run('main()', filename='test.cprof')
379
380if __name__ == '__main__':
381    #profile()
382    #SIZE_OF_IMAGE_FN = size_of_image_verify
383    main()
384