sdk/tools/gen_baseaddress.py

   1 '''
   2 PROJECT:     ReactOS baseaddress updater
   3 LICENSE:     MIT (https://spdx.org/licenses/MIT)
   4 PURPOSE:     Update baseaddresses of all modules
   5 COPYRIGHT:   Copyright 2017,2018 Mark Jansen (mark.jansen@reactos.org)
   6 '''
   7 from __future__ import print_function, absolute_import, division
   8 import os
   9 import struct
  10 import sys
  11
  12 try:
  13     import pefile
  14 except ImportError:
  15     print('# Please install pefile from pip or https://github.com/erocarrera/pefile')
  16     print('# Using fallback')
  17     print()
  18
  19 ALL_EXTENSIONS = (
  20     '.dll', '.acm', '.ax', '.cpl', '.drv', '.ocx'
  21 )
  22
  23 PRIORITIES = (
  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
  90 EXCLUDE = (
  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 IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x10b
 194 IMAGE_NT_OPTIONAL_HDR64_MAGIC = 0x20b
 195
 196 IMAGE_TYPES = {
 197     IMAGE_NT_OPTIONAL_HDR32_MAGIC: 0,
 198     IMAGE_NT_OPTIONAL_HDR64_MAGIC: 0
 199 }
 200
 201 def is_x64():
 202     return IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] > IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC]
 203
 204 def 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
 226 def 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
 232 SIZE_OF_IMAGE_FN = size_of_image_fallback
 233
 234 class 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
 259 class 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
 332 def 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
 348 def 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
 362 def 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
 372 def profile():
 373     import cProfile
 374     # pyprof2calltree -k -i test.cprof
 375     cProfile.run('main()', filename='test.cprof')
 376
 377 if __name__ == '__main__':
 378     #profile()
 379     #SIZE_OF_IMAGE_FN = size_of_image_verify
 380     main()