Merge pull request #176 from Snawoot/win_x86

Win x86

win/tools/autopatch/autopatch.py

 #!/usr/bin/env python3
 
+import argparse
 import sys
 import subprocess
 import tempfile
 import os.path
 from binascii import unhexlify
 import xml.etree.ElementTree as ET
+import itertools
+import functools
 
 
-PATCH_EXT = ".1337"
 CRLF = b"\x0d\x0a"
 HEADER_FORMAT = b">%s"
 LINE_FORMAT = CRLF + b"%016X:%02X->%02X"
@@ -16,7 +18,6 @@ OFFSET_ADJUSTMENT = 0xC00  # shift specific to x64dbg .1337 format
 
 
 def parse_args():
-    import argparse
 
     parser = argparse.ArgumentParser(
         description="Generates .1337 patch for Nvidia drivers for Windows",
@@ -29,18 +30,41 @@ def parse_args():
                         dest="sevenzip",
                         help="location of 7-zip `7z` executable")
     parser.add_argument("-T", "--target",
-                        default="Display.Driver/nvcuvid64.dl_",
-                        help="target location in archive")
+                        nargs="+",
+                        default=[
+                            "Display.Driver/nvcuvid64.dl_",
+                            "Display.Driver/nvcuvid32.dl_",
+                        ],
+                        help="target location(s) in archive")
     parser.add_argument("-N", "--target-name",
-                        default="nvcuvid.dll",
-                        help="name of installed target file. Used for patch "
+                        nargs="+",
+                        default=[
+                            "nvcuvid.dll",
+                            "nvcuvid.dll",
+                        ],
+                        help="name(s) of installed target file. Used for patch "
                         "header")
+    parser.add_argument("-P", "--patch-name",
+                        nargs="+",
+                        default=[
+                            "nvcuvid64.1337",
+                            "nvcuvid32.1337",
+                        ],
+                        help="relative filename(s) of generated patch(es)")
     parser.add_argument("-S", "--search",
-                        default="FF909800000084C075",
-                        help="representation of search pattern binary string")
+                        nargs="+",
+                        default=[
+                            "FF909800000084C075",
+                            "8B404CFFD084C075",
+                        ],
+                        help="representation of search pattern(s) binary string")
     parser.add_argument("-R", "--replacement",
-                        default="FF90980000000C0175",
-                        help="representation of replacement binary string")
+                        nargs="+",
+                        default=[
+                            "FF90980000000C0175",
+                            "8B404CFFD00C0175",
+                        ],
+                        help="representation of replacement(s) binary string")
     parser.add_argument("-o", "--stdout",
                         action="store_true",
                         help="output into stdout")
@@ -147,6 +171,7 @@ def make_patch(archive, *,
     return res
 
 
+@functools.lru_cache(maxsize=None)
 def identify_driver(archive, *, sevenzip="7z"):
     manifest = extract_single_file(archive, "setup.cfg", sevenzip=sevenzip)
     root = ET.fromstring(manifest)
@@ -162,53 +187,63 @@ def format_patch(diff, filename):
         res += LINE_FORMAT % (offset + OFFSET_ADJUSTMENT, left, right)
     return res
 
-
-def main():
-    args = parse_args()
-    search = unhexlify(args.search)
-    replacement = unhexlify(args.replacement)
+def patch_flow(installer_file, search, replacement, target, target_name, patch_name, *,
+               direct=False, stdout=False, sevenzip="7z"):
+    search = unhexlify(search)
+    replacement = unhexlify(replacement)
     assert len(search) == len(replacement), "len() of search and replacement"\
         " is not equal"
-    for installer_file in args.installer_file:
-        patch = make_patch(installer_file,
-                           arch_tgt=args.target,
-                           search=search,
-                           replacement=replacement,
-                           sevenzip=args.sevenzip,
-                           direct=args.direct)
-        patch_content = format_patch(patch, args.target_name)
-        if args.stdout or args.direct:
-            with open(sys.stdout.fileno(), mode='wb', closefd=False) as out:
-                out.write(patch_content)
+    patch = make_patch(installer_file,
+                       arch_tgt=target,
+                       search=search,
+                       replacement=replacement,
+                       sevenzip=sevenzip,
+                       direct=direct)
+    patch_content = format_patch(patch, target_name)
+    if stdout or direct:
+        with open(sys.stdout.fileno(), mode='wb', closefd=False) as out:
+            out.write(patch_content)
+    else:
+        version, product_type = identify_driver(installer_file,
+                                                sevenzip=sevenzip)
+        drv_prefix = {
+            "100": "quadro_",
+            "300": "",
+            "301": "nsd_",
+        }
+        installer_name = os.path.basename(installer_file).lower()
+        if 'winserv2008' in installer_name:
+            os_prefix = 'ws2012_x64'
+        elif 'winserv-2016' in installer_name:
+            os_prefix = 'ws2016_x64'
+        elif 'win10' in installer_name:
+            os_prefix = 'win10_x64'
+        elif 'win7' in installer_name:
+            os_prefix = 'win7_x64'
         else:
-            version, product_type = identify_driver(installer_file,
-                                                    sevenzip=args.sevenzip)
-            drv_prefix = {
-                "100": "quadro_",
-                "300": "",
-                "301": "nsd_",
-            }
-            installer_name = os.path.basename(installer_file).lower()
-            if 'winserv2008' in installer_name:
-                os_prefix = 'ws2012_x64'
-            elif 'winserv-2016' in installer_name:
-                os_prefix = 'ws2016_x64'
-            elif 'win10' in installer_name:
-                os_prefix = 'win10_x64'
-            elif 'win7' in installer_name:
-                os_prefix = 'win7_x64'
-            else:
-                raise UnknownPlatformException("Can't infer platform from filename %s"
-                                               % (repr(installer_name),))
-            driver_name = drv_prefix[product_type] + version
-            out_dir = os.path.join(
-                os.path.dirname(
-                    os.path.abspath(__file__)), '..', '..', os_prefix, driver_name)
-            os.makedirs(out_dir, 0o755, True)
-            out_filename = os.path.join(out_dir,
-                os.path.splitext(args.target_name)[0] + PATCH_EXT)
-            with open(out_filename, 'xb') as out:
-                out.write(patch_content)
+            raise UnknownPlatformException("Can't infer platform from filename %s"
+                                           % (repr(installer_name),))
+        driver_name = drv_prefix[product_type] + version
+        out_dir = os.path.join(
+            os.path.dirname(
+                os.path.abspath(__file__)), '..', '..', os_prefix, driver_name)
+        os.makedirs(out_dir, 0o755, True)
+        out_filename = os.path.join(out_dir,
+            patch_name)
+        with open(out_filename, 'xb') as out:
+            out.write(patch_content)
+
+
+def main():
+    args = parse_args()
+    if args.direct:
+        combinations = zip(args.installer_file, args.search, args.replacement,
+                           args.target, args.target_name, args.patch_name)
+    else:
+        base_params = zip(args.search, args.replacement, args.target, args.target_name, args.patch_name)
+        combinations = ((l,) + r for l, r in itertools.product(args.installer_file, base_params))
+    for params in combinations:
+        patch_flow(*params, direct=args.direct, stdout=args.stdout)
 
 
 if __name__ == '__main__':

win/win10_x64/440.97/nvcuvid32.1337

+>nvcuvid.dll
+00000000000E6109:84->0C
+00000000000E610A:C0->01
\ No newline at end of file

win/win10_x64/quadro_440.97/nvcuvid32.1337

+>nvcuvid.dll
+00000000000E6109:84->0C
+00000000000E610A:C0->01
\ No newline at end of file

win/win7_x64/440.97/nvcuvid32.1337

+>nvcuvid.dll
+00000000000D1ED9:84->0C
+00000000000D1EDA:C0->01
\ No newline at end of file

win/win7_x64/quadro_440.97/nvcuvid32.1337

+>nvcuvid.dll
+00000000000D1ED9:84->0C
+00000000000D1EDA:C0->01
\ No newline at end of file

win/ws2012_x64/quadro_440.97/nvcuvid32.1337

+>nvcuvid.dll
+00000000000D1ED9:84->0C
+00000000000D1EDA:C0->01
\ No newline at end of file

win/ws2016_x64/quadro_440.97/nvcuvid32.1337

+>nvcuvid.dll
+00000000000E6109:84->0C
+00000000000E610A:C0->01
\ No newline at end of file