xref: /reactos/sdk/tools/gen_baseaddress.py (revision 4567e13e)
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    '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
193IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x10b
194IMAGE_NT_OPTIONAL_HDR64_MAGIC = 0x20b
195
196IMAGE_TYPES = {
197    IMAGE_NT_OPTIONAL_HDR32_MAGIC: 0,
198    IMAGE_NT_OPTIONAL_HDR64_MAGIC: 0
199}
200
201def is_x64():
202    return IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] > IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC]
203
204def size_of_image_fallback(filename):
205    with open(filename, 'rb') as fin:
206        if fin.read(2) != 'MZ':
207            print(filename, 'No dos header found!')
208            return 0
209        fin.seek(0x3C)
210        e_lfanew = struct.unpack('i', fin.read(4))[0]
211        fin.seek(e_lfanew)
212        if fin.read(4) != 'PE\0\0':
213            print(filename, 'No PE header found!')
214            return 0
215        fin.seek(e_lfanew + 0x18)
216        pe_magic = struct.unpack('h', fin.read(2))[0]
217        if pe_magic in IMAGE_TYPES.keys():
218            IMAGE_TYPES[pe_magic] += 1
219            fin.seek(e_lfanew + 0x50)
220            pe_size_of_image = struct.unpack('i', fin.read(4))[0]
221            return pe_size_of_image
222        print(filename, 'Unknown executable format!')
223        return 0
224
225
226def size_of_image_verify(filename):
227    pefile_size = pefile.PE(filename, fast_load=True).OPTIONAL_HEADER.SizeOfImage
228    custom_size = size_of_image_fallback(filename)
229    assert custom_size == pefile_size, filename
230    return custom_size
231
232SIZE_OF_IMAGE_FN = size_of_image_fallback
233
234class Module(object):
235    def __init__(self, name, address, size, filename):
236        self._name = name
237        self.address = address
238        self.size = size
239        self._reserved = address != 0
240        self.filename = filename
241
242    def gen_baseaddress(self):
243        name, ext = os.path.splitext(self._name)
244        postfix = ''
245        if ext in('.acm', '.drv') and self._name != 'winspool.drv':
246            name = self._name
247        if name == 'ntdll':
248            postfix = ' # should be above 0x%08x' % self.address
249        elif self._reserved:
250            postfix = ' # reserved'
251        print('set(baseaddress_%-30s 0x%08x)%s' % (name, self.address, postfix))
252
253    def end(self):
254        return self.address + self.size
255
256    def __repr__(self):
257        return '%s (0x%08x - 0x%08x)' % (self._name, self.address, self.end())
258
259class MemoryLayout(object):
260    def __init__(self, startaddress):
261        self.addresses = []
262        self.found = {}
263        self.reserved = {}
264        self.initial = startaddress
265        self.start_at = 0
266        self.module_padding = 0x2000
267
268    def add_reserved(self, name, address):
269        self.reserved[name] = (address, 0)
270
271    def add(self, filename, name):
272        size = SIZE_OF_IMAGE_FN(filename)
273        addr = 0
274        if name in self.found:
275            return  # Assume duplicate files (rshell, ...) are 1:1 copies
276        if name in self.reserved:
277            addr = self.reserved[name][0]
278            self.reserved[name] = (addr, size)
279        self.found[name] = Module(name, addr, size, filename)
280
281    def _next_address(self, size):
282        if self.start_at:
283            addr = (self.start_at - size - self.module_padding - 0xffff) & 0xffff0000
284            self.start_at = addr
285        else:
286            addr = self.start_at = self.initial
287        return addr
288
289    def next_address(self, size):
290        while True:
291            current_start = self._next_address(size)
292            current_end = current_start + size + self.module_padding
293            # Is there overlap with reserved modules?
294            for key, reserved in self.reserved.items():
295                res_start = reserved[0]
296                res_end = res_start + reserved[1] + self.module_padding
297                if (res_start <= current_start <= res_end) or \
298                   (res_start <= current_end <= res_end) or \
299                   (current_start < res_start and current_end > res_end):
300                    # We passed this reserved item, we can remove it now
301                    self.start_at = min(res_start, current_start)
302                    del self.reserved[key]
303                    current_start = 0
304                    break
305            # No overlap with a reserved module?
306            if current_start:
307                return current_start
308
309    def update(self, priorities):
310        # sort addresses, should only contain reserved modules at this point!
311        for key, reserved in self.reserved.items():
312            assert reserved[1] != 0, key
313        for curr in priorities:
314            if not curr in self.found:
315                print('# Did not find', curr, '!')
316            else:
317                obj = self.found[curr]
318                del self.found[curr]
319                if not obj.address:
320                    obj.address = self.next_address(obj.size)
321                self.addresses.append(obj)
322        # We handled all known modules now, run over the rest we found
323        for key in sorted(self.found):
324            obj = self.found[key]
325            obj.address = self.next_address(obj.size)
326            self.addresses.append(obj)
327
328    def gen_baseaddress(self):
329        for obj in self.addresses:
330            obj.gen_baseaddress()
331
332def guess_version(ntdll_path):
333    if 'pefile' in globals():
334        ntdll_pe = pefile.PE(ntdll_path, fast_load=True)
335        names = [sect.Name.strip('\0') for sect in ntdll_pe.sections]
336        count = '|'.join(names).count('/')
337        if '.rossym' in names:
338            print('# This should probably go in baseaddress.cmake')
339        elif is_x64():
340            print('# This should probably go in baseaddress_msvc_x64.cmake')
341        elif count == 0:
342            print('# This should probably go in baseaddress_msvc.cmake')
343        elif count > 3:
344            print('# This should probably go in baseaddress_dwarf.cmake')
345        else:
346            print('# No clue where to put this')
347
348def run_dir(target):
349    print('# Generated from', target)
350    print('# Generated by sdk/tools/gen_baseaddress.py')
351    layout = MemoryLayout(0x7c920000)
352    layout.add_reserved('user32.dll', 0x77a20000)
353    for root, _, files in os.walk(target):
354        for dll in [filename for filename in files if filename.endswith(ALL_EXTENSIONS)]:
355            if not dll in EXCLUDE and not dll.startswith('api-ms-win-'):
356                layout.add(os.path.join(root, dll), dll)
357    ntdll_path = layout.found['ntdll.dll'].filename
358    guess_version(ntdll_path)
359    layout.update(PRIORITIES)
360    layout.gen_baseaddress()
361
362def main():
363    dirs = sys.argv[1:]
364    if len(dirs) < 1:
365        trydir = os.getcwd()
366        print('# No path specified, trying', trydir)
367        dirs = [trydir]
368    for onedir in dirs:
369        run_dir(onedir)
370
371
372def profile():
373    import cProfile
374    # pyprof2calltree -k -i test.cprof
375    cProfile.run('main()', filename='test.cprof')
376
377if __name__ == '__main__':
378    #profile()
379    #SIZE_OF_IMAGE_FN = size_of_image_verify
380    main()
381