nvidia-patch/win/tools/autopatch/autopatch.py
Jai Luthra 2f11563e5c autopatch: Ensure non-zero retval for missing installer
Signed-off-by: Jai Luthra <me@jailuthra.in>
2023-11-09 14:19:38 +05:30

301 lines
11 KiB
Python
Executable File

#!/usr/bin/env python3
import argparse
import functools
import itertools
import os.path
import subprocess
import sys
import tempfile
import urllib.request
import xml.etree.ElementTree as ET
from binascii import unhexlify
CRLF = b"\x0d\x0a"
HEADER_FORMAT = b">%s"
LINE_FORMAT = CRLF + b"%016X:%02X->%02X"
OFFSET_ADJUSTMENT = 0xC00 # shift specific to x64dbg .1337 format
def parse_args():
parser = argparse.ArgumentParser(
description="Generates .1337 patch for Nvidia drivers for Windows",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("installer_file",
nargs="+",
help="location of installer executable(s)")
parser.add_argument("-7", "--7zip",
default="7z",
dest="sevenzip",
help="location of 7-zip `7z` executable")
parser.add_argument("-T", "--target",
nargs="+",
default=[
"Display.Driver/nvencodeapi64.dll",
"Display.Driver/nvencodeapi.dll",
],
help="target location(s) in archive")
parser.add_argument("-N", "--target-name",
nargs="+",
default=[
"nvencodeapi64.dll",
"nvencodeapi.dll",
],
help="name(s) of installed target file. Used for patch "
"header")
parser.add_argument("-P", "--patch-name",
nargs="+",
default=[
"nvencodeapi64.1337",
"nvencodeapi.1337",
],
help="relative filename(s) of generated patch(es)")
parser.add_argument("-S", "--search",
nargs="+",
default=[
"8BF085C0750549892FEB",
"89450885C08B450C75048938EB",
],
help="representation of search pattern(s) binary string")
parser.add_argument("-R", "--replacement",
nargs="+",
default=[
"33C08BF0750549892FEB",
"33C08945088B450C75048938EB",
],
help="representation of replacement(s) binary string")
parser.add_argument("-o", "--stdout",
action="store_true",
help="output into stdout")
parser.add_argument("-D", "--direct",
action="store_true",
help="supply patched library directly instead of "
"installer file")
args = parser.parse_args()
return args
class ExtractException(Exception):
pass
class PatternNotFoundException(Exception):
pass
class MultipleOccurencesException(Exception):
pass
class UnknownPlatformException(Exception):
pass
class InstallerNotFoundException(Exception):
pass
class ExtractedTarget:
name = None
def __init__(self, archive, dst_dir, arch_tgt, *, sevenzip="7z"):
self._archive = archive
self._dst_dir = dst_dir
self._sevenzip = sevenzip
self._arch_tgt = arch_tgt
def __enter__(self):
ret = subprocess.call([self._sevenzip,
"e",
"-o" + self._dst_dir,
self._archive,
self._arch_tgt],
stdout=sys.stderr)
if ret != 0:
raise ExtractException("Subprocess returned non-zero exit code.")
name = os.path.join(self._dst_dir, os.path.basename(self._arch_tgt))
self.name = name
return name
def __exit__(self, exc_type, exc_value, traceback):
if self.name is not None:
os.remove(self.name)
def expand(filename, *, sevenzip="7z"):
proc = subprocess.Popen([sevenzip,
"x",
"-so",
filename], stdout=subprocess.PIPE)
result = proc.communicate()[0]
if proc.returncode != 0:
raise ExtractException("Subprocess returned non-zero exit code.")
return result
def extract_single_file(archive, filename, *, sevenzip="7z"):
proc = subprocess.Popen([sevenzip,
"e",
"-so",
archive,
filename], stdout=subprocess.PIPE)
result = proc.communicate()[0]
if proc.returncode != 0:
raise ExtractException("Subprocess returned non-zero exit code.")
return result
def make_patch(archive, *,
arch_tgt,
search,
replacement,
tmpdir,
sevenzip="7z",
direct=False):
if direct:
with open(archive, 'rb') as fo:
f = fo.read()
else:
with ExtractedTarget(archive,
tmpdir,
arch_tgt,
sevenzip=sevenzip) as tgt:
if tgt.endswith(".dll"):
with open(tgt, 'rb') as fo:
f = fo.read()
else:
f = expand(tgt, sevenzip=sevenzip)
offset = f.find(search)
if offset == -1:
raise PatternNotFoundException("Pattern not found.")
if f[offset + len(search):].find(search) != -1:
raise MultipleOccurencesException("Multiple occurences of pattern found!")
del f
print("Pattern found @ %016X" % (offset,), file=sys.stderr)
res = []
for (i, (left, right)) in enumerate(zip(search, replacement)):
if left != right:
res.append((offset + i, left, right))
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)
version = root.attrib['version']
product_type = root.find('./properties/string[@name="ProductType"]') \
.attrib['value']
return version, product_type
def format_patch(diff, filename):
res = HEADER_FORMAT % filename.encode('utf-8')
for offset, left, right in diff:
res += LINE_FORMAT % (offset + OFFSET_ADJUSTMENT, left, right)
return res
def patch_flow(installer_file, search, replacement, target, target_name, patch_name, *,
tempdir, 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"
# 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:
raise InstallerNotFoundException(f"Failed to download the file: {e}")
except Exception as e:
raise InstallerNotFoundException(f"An error occurred during download: {str(e)}")
else:
raise InstallerNotFoundException(f"Invalid installer file or version: {installer_file}")
# Rest of the code remains the same...
patch = make_patch(installer_file,
arch_tgt=target,
search=search,
replacement=replacement,
tmpdir=tempdir,
sevenzip=sevenzip,
direct=direct)
patch_content = format_patch(patch, target_name)
if stdout:
sys.stdout.buffer.write(patch_content)
elif direct:
with open(patch_name, mode='wb') as out:
out.write(patch_content)
else:
version, product_type = identify_driver(installer_file, sevenzip=sevenzip)
drv_prefix = {
"100": "quadro_",
"103": "quadro_",
"300": "",
"301": "nsd_",
"303": "", # DCH
"304": "nsd_",
}
installer_name = os.path.basename(installer_file).lower()
if 'winserv2008' in installer_name or 'winserv-2012' in installer_name:
os_prefix = 'ws2012_x64'
elif 'winserv-2016' in installer_name or 'win10' in installer_name:
os_prefix = 'win10_x64'
elif 'win7' in installer_name:
os_prefix = 'win7_x64'
else:
raise UnknownPlatformException(f"Can't infer platform from filename {installer_name}")
driver_name = drv_prefix.get(product_type, "") + version
out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', os_prefix, driver_name)
os.makedirs(out_dir, 0o755, exist_ok=True)
out_filename = os.path.join(out_dir, patch_name)
with open(out_filename, 'wb') 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))
with tempfile.TemporaryDirectory() as tempdir:
print(f"Using tempdir `{tempdir}`")
for params in combinations:
patch_flow(*params, tempdir=tempdir, direct=args.direct, stdout=args.stdout)
if __name__ == '__main__':
main()