xref: /reactos/sdk/tools/gen_baseaddress.py (revision f4d29a74)
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
193
194def size_of_image_fallback(filename):
195    with open(filename, 'rb') as fin:
196        if fin.read(2) != 'MZ':
197            print(filename, 'No dos header found!')
198            return 0
199        fin.seek(0x3C)
200        e_lfanew = struct.unpack('i', fin.read(4))[0]
201        fin.seek(e_lfanew)
202        if fin.read(4) != 'PE\0\0':
203            print(filename, 'No PE header found!')
204            return 0
205        fin.seek(e_lfanew + 0x18)
206        pe_magic = struct.unpack('h', fin.read(2))[0]
207        if pe_magic != 0x10b:
208            print(filename, 'is not a 32 bit exe!')
209            return 0
210        fin.seek(e_lfanew + 0x50)
211        pe_size_of_image = struct.unpack('i', fin.read(4))[0]
212        return pe_size_of_image
213
214def size_of_image(filename):
215    if 'pefile' in globals():
216        return pefile.PE(filename, fast_load=True).OPTIONAL_HEADER.SizeOfImage
217    return size_of_image_fallback(filename)
218
219
220class Module(object):
221    def __init__(self, name, address, size):
222        self._name = name
223        self.address = address
224        self.size = size
225        self._reserved = address != 0
226
227    def gen_baseaddress(self):
228        name, ext = os.path.splitext(self._name)
229        postfix = ''
230        if ext in('.acm', '.drv') and self._name != 'winspool.drv':
231            name = self._name
232        if name == 'ntdll':
233            postfix = ' # should be above 0x%08x' % self.address
234        elif self._reserved:
235            postfix = ' # reserved'
236        print('set(baseaddress_%-30s 0x%08x)%s' % (name, self.address, postfix))
237
238    def end(self):
239        return self.address + self.size
240
241    def __repr__(self):
242        return '%s (0x%08x - 0x%08x)' % (self._name, self.address, self.end())
243
244class MemoryLayout(object):
245    def __init__(self, startaddress):
246        self.addresses = []
247        self.found = {}
248        self.reserved = {}
249        self.initial = startaddress
250        self.start_at = 0
251        self.module_padding = 0x2000
252
253    def add_reserved(self, name, address):
254        self.reserved[name] = (address, 0)
255
256    def add(self, filename, name):
257        size = size_of_image(filename)
258        addr = 0
259        if name in self.found:
260            return  # Assume duplicate files (rshell, ...) are 1:1 copies
261        if name in self.reserved:
262            addr = self.reserved[name][0]
263            self.reserved[name] = (addr, size)
264        self.found[name] = Module(name, addr, size)
265
266    def _next_address(self, size):
267        if self.start_at:
268            addr = (self.start_at - size - self.module_padding - 0xffff) & 0xffff0000
269            self.start_at = addr
270        else:
271            addr = self.start_at = self.initial
272        return addr
273
274    def next_address(self, size):
275        while True:
276            current_start = self._next_address(size)
277            current_end = current_start + size + self.module_padding
278            # Is there overlap with reserved modules?
279            for key, reserved in self.reserved.items():
280                res_start = reserved[0]
281                res_end = res_start + reserved[1] + self.module_padding
282                if (res_start <= current_start <= res_end) or \
283                   (res_start <= current_end <= res_end) or \
284                   (current_start < res_start and current_end > res_end):
285                    # We passed this reserved item, we can remove it now
286                    self.start_at = min(res_start, current_start)
287                    del self.reserved[key]
288                    current_start = 0
289                    break
290            # No overlap with a reserved module?
291            if current_start:
292                return current_start
293
294    def update(self, priorities):
295        # sort addresses, should only contain reserved modules at this point!
296        for key, reserved in self.reserved.items():
297            assert reserved[1] != 0, key
298        for curr in priorities:
299            if not curr in self.found:
300                print('# Did not find', curr, '!')
301            else:
302                obj = self.found[curr]
303                del self.found[curr]
304                if not obj.address:
305                    obj.address = self.next_address(obj.size)
306                self.addresses.append(obj)
307        # We handled all known modules now, run over the rest we found
308        for key in sorted(self.found):
309            obj = self.found[key]
310            obj.address = self.next_address(obj.size)
311            self.addresses.append(obj)
312
313    def gen_baseaddress(self):
314        for obj in self.addresses:
315            obj.gen_baseaddress()
316
317def run_dir(target):
318    print('# Generated from', target)
319    print('# Generated by sdk/tools/gen_baseaddress.py')
320    layout = MemoryLayout(0x7c920000)
321    layout.add_reserved('user32.dll', 0x77a20000)
322    for root, _, files in os.walk(target):
323        for dll in [filename for filename in files if filename.endswith(ALL_EXTENSIONS)]:
324            if not dll in EXCLUDE and not dll.startswith('api-ms-win-'):
325                layout.add(os.path.join(root, dll), dll)
326    layout.update(PRIORITIES)
327    layout.gen_baseaddress()
328
329def main(dirs):
330    if len(dirs) < 1:
331        trydir = os.getcwd()
332        print('# No path specified, trying', trydir)
333        dirs = [trydir]
334    for onedir in dirs:
335        run_dir(onedir)
336
337if __name__ == '__main__':
338    main(sys.argv[1:])
339