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