[SDK] Update gen_baseaddress.

* Speed up the script by using the fallback by default
* Suggest what cmake file to edit
* Add support for x64 modules (addresses not updated yet)
CORE-14923

diff --git a/sdk/tools/gen_baseaddress.py b/sdk/tools/gen_baseaddress.py
index 5219a4e..3a4c1e6 100644 (file)
--- a/sdk/tools/gen_baseaddress.py
+++ b/sdk/tools/gen_baseaddress.py
@@ -190,6 +190,16 @@ EXCLUDE = (
     'w32kdll_xpsp2.dll',
 )
 
+IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x10b
+IMAGE_NT_OPTIONAL_HDR64_MAGIC = 0x20b
+
+IMAGE_TYPES = {
+    IMAGE_NT_OPTIONAL_HDR32_MAGIC: 0,
+    IMAGE_NT_OPTIONAL_HDR64_MAGIC: 0
+}
+
+def is_x64():
+    return IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] > IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC]
 
 def size_of_image_fallback(filename):
     with open(filename, 'rb') as fin:
@@ -204,25 +214,30 @@ def size_of_image_fallback(filename):
             return 0
         fin.seek(e_lfanew + 0x18)
         pe_magic = struct.unpack('h', fin.read(2))[0]
-        if pe_magic != 0x10b:
-            print(filename, 'is not a 32 bit exe!')
-            return 0
-        fin.seek(e_lfanew + 0x50)
-        pe_size_of_image = struct.unpack('i', fin.read(4))[0]
-        return pe_size_of_image
+        if pe_magic in IMAGE_TYPES.keys():
+            IMAGE_TYPES[pe_magic] += 1
+            fin.seek(e_lfanew + 0x50)
+            pe_size_of_image = struct.unpack('i', fin.read(4))[0]
+            return pe_size_of_image
+        print(filename, 'Unknown executable format!')
+        return 0
 
-def size_of_image(filename):
-    if 'pefile' in globals():
-        return pefile.PE(filename, fast_load=True).OPTIONAL_HEADER.SizeOfImage
-    return size_of_image_fallback(filename)
 
+def size_of_image_verify(filename):
+    pefile_size = pefile.PE(filename, fast_load=True).OPTIONAL_HEADER.SizeOfImage
+    custom_size = size_of_image_fallback(filename)
+    assert custom_size == pefile_size, filename
+    return custom_size
+
+SIZE_OF_IMAGE_FN = size_of_image_fallback
 
 class Module(object):
-    def __init__(self, name, address, size):
+    def __init__(self, name, address, size, filename):
         self._name = name
         self.address = address
         self.size = size
         self._reserved = address != 0
+        self.filename = filename
 
     def gen_baseaddress(self):
         name, ext = os.path.splitext(self._name)
@@ -254,14 +269,14 @@ class MemoryLayout(object):
         self.reserved[name] = (address, 0)
 
     def add(self, filename, name):
-        size = size_of_image(filename)
+        size = SIZE_OF_IMAGE_FN(filename)
         addr = 0
         if name in self.found:
             return  # Assume duplicate files (rshell, ...) are 1:1 copies
         if name in self.reserved:
             addr = self.reserved[name][0]
             self.reserved[name] = (addr, size)
-        self.found[name] = Module(name, addr, size)
+        self.found[name] = Module(name, addr, size, filename)
 
     def _next_address(self, size):
         if self.start_at:
@@ -314,6 +329,22 @@ class MemoryLayout(object):
         for obj in self.addresses:
             obj.gen_baseaddress()
 
+def guess_version(ntdll_path):
+    if 'pefile' in globals():
+        ntdll_pe = pefile.PE(ntdll_path, fast_load=True)
+        names = [sect.Name.strip('\0') for sect in ntdll_pe.sections]
+        count = '|'.join(names).count('/')
+        if '.rossym' in names:
+            print('# This should probably go in baseaddress.cmake')
+        elif is_x64():
+            print('# This should probably go in baseaddress_msvc_x64.cmake')
+        elif count == 0:
+            print('# This should probably go in baseaddress_msvc.cmake')
+        elif count > 3:
+            print('# This should probably go in baseaddress_dwarf.cmake')
+        else:
+            print('# No clue where to put this')
+
 def run_dir(target):
     print('# Generated from', target)
     print('# Generated by sdk/tools/gen_baseaddress.py')
@@ -323,10 +354,13 @@ def run_dir(target):
         for dll in [filename for filename in files if filename.endswith(ALL_EXTENSIONS)]:
             if not dll in EXCLUDE and not dll.startswith('api-ms-win-'):
                 layout.add(os.path.join(root, dll), dll)
+    ntdll_path = layout.found['ntdll.dll'].filename
+    guess_version(ntdll_path)
     layout.update(PRIORITIES)
     layout.gen_baseaddress()
 
-def main(dirs):
+def main():
+    dirs = sys.argv[1:]
     if len(dirs) < 1:
         trydir = os.getcwd()
         print('# No path specified, trying', trydir)
@@ -334,5 +368,13 @@ def main(dirs):
     for onedir in dirs:
         run_dir(onedir)
 
+
+def profile():
+    import cProfile
+    # pyprof2calltree -k -i test.cprof
+    cProfile.run('main()', filename='test.cprof')
+
 if __name__ == '__main__':
-    main(sys.argv[1:])
+    #profile()
+    #SIZE_OF_IMAGE_FN = size_of_image_verify
+    main()