[SDK] Update the gen_baseaddress script

- Remove some debug leftovers
- Automatically write the correct file

diff --git a/sdk/tools/gen_baseaddress.py b/sdk/tools/gen_baseaddress.py
index 2c53c6c..f214012 100644 (file)
--- a/sdk/tools/gen_baseaddress.py
+++ b/sdk/tools/gen_baseaddress.py
@@ -13,8 +13,7 @@ try:
     import pefile
 except ImportError:
     print('# Please install pefile from pip or https://github.com/erocarrera/pefile')
-    print('# Using fallback')
-    print()
+    sys.exit(-1)
 
 ALL_EXTENSIONS = (
     '.dll', '.acm', '.ax', '.cpl', '.drv', '.ocx'
@@ -202,7 +201,7 @@ IMAGE_TYPES = {
 def is_x64():
     return IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] > IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC]
 
-def size_of_image_fallback(filename):
+def size_of_image(filename):
     with open(filename, 'rb') as fin:
         if fin.read(2) != b'MZ':
             print(filename, 'No dos header found!')
@@ -224,14 +223,6 @@ def size_of_image_fallback(filename):
         return 0
 
 
-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, filename):
         self._name = name
@@ -240,7 +231,7 @@ class Module(object):
         self._reserved = address != 0
         self.filename = filename
 
-    def gen_baseaddress(self):
+    def gen_baseaddress(self, output_file):
         name, ext = os.path.splitext(self._name)
         postfix = ''
         if ext in('.acm', '.drv') and self._name != 'winspool.drv':
@@ -249,7 +240,7 @@ class Module(object):
             postfix = ' # should be above 0x%08x' % self.address
         elif self._reserved:
             postfix = ' # reserved'
-        print('set(baseaddress_%-30s 0x%08x)%s' % (name, self.address, postfix))
+        output_file.write('set(baseaddress_%-30s 0x%08x)%s\n' % (name, self.address, postfix))
 
     def end(self):
         return self.address + self.size
@@ -270,7 +261,7 @@ class MemoryLayout(object):
         self.reserved[name] = (address, 0)
 
     def add(self, filename, name):
-        size = SIZE_OF_IMAGE_FN(filename)
+        size = size_of_image(filename)
         addr = 0
         if name in self.found:
             return  # Assume duplicate files (rshell, ...) are 1:1 copies
@@ -326,39 +317,49 @@ class MemoryLayout(object):
             obj.address = self.next_address(obj.size)
             self.addresses.append(obj)
 
-    def gen_baseaddress(self):
+    def gen_baseaddress(self, output_file):
         for obj in self.addresses:
-            obj.gen_baseaddress()
+            obj.gen_baseaddress(output_file)
 
-def guess_version(ntdll_path):
+def get_target_file(ntdll_path):
     if 'pefile' in globals():
         ntdll_pe = pefile.PE(ntdll_path, fast_load=True)
         names = [sect.Name.strip(b'\0') for sect in ntdll_pe.sections]
         count = b'|'.join(names).count(b'/')
         if b'.rossym' in names:
-            print('# This should probably go in sdk/cmake/baseaddress.cmake')
+            return 'baseaddress.cmake'
         elif is_x64():
-            print('# This should probably go in sdk/cmake/baseaddress_msvc_x64.cmake')
+            return 'baseaddress_msvc_x64.cmake'
         elif count == 0:
-            print('# This should probably go in sdk/cmake/baseaddress_msvc.cmake')
+            return 'baseaddress_msvc.cmake'
         elif count > 3:
-            print('# This should probably go in sdk/cmake/baseaddress_dwarf.cmake')
+            return 'baseaddress_dwarf.cmake'
         else:
-            print('# No clue where to put this')
+            assert False, "Unknown"
+    return None
 
 def run_dir(target):
-    print('# Generated from', target)
-    print('# Generated by sdk/tools/gen_baseaddress.py')
     layout = MemoryLayout(0x7c920000)
     layout.add_reserved('user32.dll', 0x77a20000)
+    IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR64_MAGIC] = 0
+    IMAGE_TYPES[IMAGE_NT_OPTIONAL_HDR32_MAGIC] = 0
     for root, _, files in os.walk(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()
+    target_file = get_target_file(ntdll_path)
+    if target_file:
+        target_dir = os.path.realpath(os.path.dirname(os.path.dirname(__file__)))
+        target_path = os.path.join(target_dir, 'cmake', target_file)
+        output_file = open(target_path, "w")
+    else:
+        output_file = sys.stdout
+    with output_file:
+        output_file.write('# Generated from {}\n'.format(target))
+        output_file.write('# Generated by sdk/tools/gen_baseaddress.py\n\n')
+        layout.update(PRIORITIES)
+        layout.gen_baseaddress(output_file)
 
 def main():
     dirs = sys.argv[1:]
@@ -370,12 +371,5 @@ def main():
         run_dir(onedir)
 
 
-def profile():
-    import cProfile
-    # pyprof2calltree -k -i test.cprof
-    cProfile.run('main()', filename='test.cprof')
-
 if __name__ == '__main__':
-    #profile()
-    #SIZE_OF_IMAGE_FN = size_of_image_verify
     main()