Commit 5b530414 authored by reloxx13's avatar reloxx13

win autopatch.py:

- reuse temp dir
- fix wrong filemode xb to wb, also fixes to override existing patch on re-patching
- avoid double download and double extracation for 32bit and 64bit
- some cleanups and import ordering
parent b89a5097
#!/usr/bin/env python3 #!/usr/bin/env python3
import argparse import argparse
import sys import functools
import itertools
import os.path
import subprocess import subprocess
import sys
import tempfile import tempfile
import os.path
from binascii import unhexlify
import xml.etree.ElementTree as ET
import itertools
import functools
import urllib.request import urllib.request
import xml.etree.ElementTree as ET
from binascii import unhexlify
CRLF = b"\x0d\x0a" CRLF = b"\x0d\x0a"
HEADER_FORMAT = b">%s" HEADER_FORMAT = b">%s"
...@@ -19,7 +18,6 @@ OFFSET_ADJUSTMENT = 0xC00 # shift specific to x64dbg .1337 format ...@@ -19,7 +18,6 @@ OFFSET_ADJUSTMENT = 0xC00 # shift specific to x64dbg .1337 format
def parse_args(): def parse_args():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Generates .1337 patch for Nvidia drivers for Windows", description="Generates .1337 patch for Nvidia drivers for Windows",
formatter_class=argparse.ArgumentDefaultsHelpFormatter) formatter_class=argparse.ArgumentDefaultsHelpFormatter)
...@@ -84,9 +82,11 @@ class ExtractException(Exception): ...@@ -84,9 +82,11 @@ class ExtractException(Exception):
class PatternNotFoundException(Exception): class PatternNotFoundException(Exception):
pass pass
class MultipleOccurencesException(Exception): class MultipleOccurencesException(Exception):
pass pass
class UnknownPlatformException(Exception): class UnknownPlatformException(Exception):
pass pass
...@@ -145,13 +145,13 @@ def make_patch(archive, *, ...@@ -145,13 +145,13 @@ def make_patch(archive, *,
arch_tgt, arch_tgt,
search, search,
replacement, replacement,
tmpdir,
sevenzip="7z", sevenzip="7z",
direct=False): direct=False):
if direct: if direct:
with open(archive, 'rb') as fo: with open(archive, 'rb') as fo:
f = fo.read() f = fo.read()
else: else:
with tempfile.TemporaryDirectory() as tmpdir:
with ExtractedTarget(archive, with ExtractedTarget(archive,
tmpdir, tmpdir,
arch_tgt, arch_tgt,
...@@ -164,7 +164,7 @@ def make_patch(archive, *, ...@@ -164,7 +164,7 @@ def make_patch(archive, *,
offset = f.find(search) offset = f.find(search)
if offset == -1: if offset == -1:
raise PatternNotFoundException("Pattern not found.") raise PatternNotFoundException("Pattern not found.")
if f[offset+len(search):].find(search) != -1: if f[offset + len(search):].find(search) != -1:
raise MultipleOccurencesException("Multiple occurences of pattern found!") raise MultipleOccurencesException("Multiple occurences of pattern found!")
del f del f
print("Pattern found @ %016X" % (offset,), file=sys.stderr) print("Pattern found @ %016X" % (offset,), file=sys.stderr)
...@@ -181,7 +181,7 @@ def identify_driver(archive, *, sevenzip="7z"): ...@@ -181,7 +181,7 @@ def identify_driver(archive, *, sevenzip="7z"):
manifest = extract_single_file(archive, "setup.cfg", sevenzip=sevenzip) manifest = extract_single_file(archive, "setup.cfg", sevenzip=sevenzip)
root = ET.fromstring(manifest) root = ET.fromstring(manifest)
version = root.attrib['version'] version = root.attrib['version']
product_type = root.find('./properties/string[@name="ProductType"]')\ product_type = root.find('./properties/string[@name="ProductType"]') \
.attrib['value'] .attrib['value']
return version, product_type return version, product_type
...@@ -192,49 +192,69 @@ def format_patch(diff, filename): ...@@ -192,49 +192,69 @@ def format_patch(diff, filename):
res += LINE_FORMAT % (offset + OFFSET_ADJUSTMENT, left, right) res += LINE_FORMAT % (offset + OFFSET_ADJUSTMENT, left, right)
return res return res
def patch_flow(installer_file, search, replacement, target, target_name, patch_name, *, def patch_flow(installer_file, search, replacement, target, target_name, patch_name, *,
direct=False, stdout=False, sevenzip="7z"): tempdir, direct=False, stdout=False, sevenzip="7z"):
search = unhexlify(search) search = unhexlify(search)
replacement = unhexlify(replacement) replacement = unhexlify(replacement)
assert len(search) == len(replacement), "len() of search and replacement"\ assert len(search) == len(replacement), "len() of search and replacement is not equal"
" is not equal"
# check if installer file exists or try to download
if not os.path.isfile(installer_file): #installer file does not exists, get url for download
if not installer_file.startswith("http"): #installer_file is a version, parse to url
filename = installer_file+"-desktop-win10-win11-64bit-international-dch-whql.exe"
installer_file = "https://international.download.nvidia.com/Windows/"+installer_file+"/"+filename
else: # installer_file is an url
filename = os.path.basename(installer_file)
# download installer and save in .temp
if not os.path.isfile(os.path.join('temp', filename)): # check if file already downloaded
print(f"Downloading... ( {installer_file} TO {os.path.join('temp', filename)} )")
print("This may take a while (~800MB)")
urllib.request.urlretrieve(installer_file, os.path.join('temp', filename))
installer_file = os.path.join('temp', filename)
else:
installer_file = os.path.join('temp', filename)
print(f"Use downloaded file in `{installer_file}`")
# Check if installer file exists or try to download
print(f"Search for installer file `{installer_file}`...")
if not os.path.isfile(installer_file):
print("Installer file is not a file...")
if not installer_file.startswith("http"):
print("Installer file is not a URL...")
# Construct URL from version
print("Installer file is a version!")
filename = installer_file + "-desktop-win10-win11-64bit-international-dch-whql.exe"
installer_url = f"https://international.download.nvidia.com/Windows/{installer_file}/{filename}"
else:
print("Installer file is a URL!")
installer_url = installer_file
if installer_url:
try:
file_path = os.path.join(tempdir, os.path.basename(installer_url))
if not os.path.isfile(file_path):
with urllib.request.urlopen(installer_url) as response, open(file_path, 'wb') as out_file:
print(f"Downloading... ({installer_url} TO {file_path})")
print("This may take a while (~800MB)")
out_file.write(response.read())
print("Download completed successfully!")
installer_file = file_path
else:
print(f"Using downloaded file in '{file_path}'")
installer_file = file_path
except (urllib.error.URLError, Exception) as e:
print(f"Failed to download the file: {e}")
return
except Exception as e:
print(f"An error occurred during download: {str(e)}")
return
else:
print(f"Invalid installer file or version: {installer_file}")
return
# Rest of the code remains the same...
patch = make_patch(installer_file, patch = make_patch(installer_file,
arch_tgt=target, arch_tgt=target,
search=search, search=search,
replacement=replacement, replacement=replacement,
tmpdir=tempdir,
sevenzip=sevenzip, sevenzip=sevenzip,
direct=direct) direct=direct)
patch_content = format_patch(patch, target_name) patch_content = format_patch(patch, target_name)
if stdout: if stdout:
with open(sys.stdout.fileno(), mode='wb', closefd=False) as out: sys.stdout.buffer.write(patch_content)
out.write(patch_content)
elif direct: elif direct:
with open(patch_name, mode='wb') as out: with open(patch_name, mode='wb') as out:
out.write(patch_content) out.write(patch_content)
else: else:
version, product_type = identify_driver(installer_file, version, product_type = identify_driver(installer_file, sevenzip=sevenzip)
sevenzip=sevenzip)
drv_prefix = { drv_prefix = {
"100": "quadro_", "100": "quadro_",
"103": "quadro_", "103": "quadro_",
...@@ -244,40 +264,37 @@ def patch_flow(installer_file, search, replacement, target, target_name, patch_n ...@@ -244,40 +264,37 @@ def patch_flow(installer_file, search, replacement, target, target_name, patch_n
"304": "nsd_", "304": "nsd_",
} }
installer_name = os.path.basename(installer_file).lower() installer_name = os.path.basename(installer_file).lower()
if 'winserv2008' in installer_name: if 'winserv2008' in installer_name or 'winserv-2012' in installer_name:
os_prefix = 'ws2012_x64'
elif 'winserv-2012' in installer_name:
os_prefix = 'ws2012_x64' os_prefix = 'ws2012_x64'
elif 'winserv-2016' in installer_name: elif 'winserv-2016' in installer_name or 'win10' in installer_name:
os_prefix = 'ws2016_x64'
elif 'win10' in installer_name:
os_prefix = 'win10_x64' os_prefix = 'win10_x64'
elif 'win7' in installer_name: elif 'win7' in installer_name:
os_prefix = 'win7_x64' os_prefix = 'win7_x64'
else: else:
raise UnknownPlatformException("Can't infer platform from filename %s" raise UnknownPlatformException(f"Can't infer platform from filename {installer_name}")
% (repr(installer_name),))
driver_name = drv_prefix[product_type] + version driver_name = drv_prefix.get(product_type, "") + version
out_dir = os.path.join( out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', os_prefix, driver_name)
os.path.dirname( os.makedirs(out_dir, 0o755, exist_ok=True)
os.path.abspath(__file__)), '..', '..', os_prefix, driver_name) out_filename = os.path.join(out_dir, patch_name)
os.makedirs(out_dir, 0o755, True) with open(out_filename, 'wb') as out:
out_filename = os.path.join(out_dir,
patch_name)
with open(out_filename, 'xb') as out:
out.write(patch_content) out.write(patch_content)
def main(): def main():
args = parse_args() args = parse_args()
if args.direct: if args.direct:
combinations = zip(args.installer_file, args.search, args.replacement, combinations = zip(args.installer_file, args.search, args.replacement,
args.target, args.target_name, args.patch_name) args.target, args.target_name, args.patch_name)
else: else:
base_params = zip(args.search, args.replacement, args.target, args.target_name, args.patch_name) 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)) combinations = ((l,) + r for l, r in itertools.product(args.installer_file, base_params))
with tempfile.TemporaryDirectory() as tempdir:
print(f"Using tempdir `{tempdir}`")
for params in combinations: for params in combinations:
patch_flow(*params, direct=args.direct, stdout=args.stdout) patch_flow(*params, tempdir=tempdir, direct=args.direct, stdout=args.stdout)
if __name__ == '__main__': if __name__ == '__main__':
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment