From e20f2c539c98426bf64c213407106b5dacee656f Mon Sep 17 00:00:00 2001 From: Will Toohey Date: Sun, 25 Oct 2020 16:48:06 +1000 Subject: [PATCH] Super files: MD5 checks and balances --- ifstools/handlers/GenericFolder.py | 45 ++++++- ifstools/handlers/MD5Folder.py | 7 +- ifstools/handlers/TexFolder.py | 6 +- ifstools/ifs.py | 17 ++- ifstools/ifstools.py | 189 +++++++++++++++-------------- 5 files changed, 158 insertions(+), 106 deletions(-) diff --git a/ifstools/handlers/GenericFolder.py b/ifstools/handlers/GenericFolder.py index 939b416..acf88ef 100644 --- a/ifstools/handlers/GenericFolder.py +++ b/ifstools/handlers/GenericFolder.py @@ -3,13 +3,16 @@ from os.path import getmtime, basename, dirname, join, realpath, isfile from collections import OrderedDict import lxml.etree as etree +from tqdm import tqdm from . import GenericFile from .Node import Node class GenericFolder(Node): - def __init__(self, ifs_data, obj, parent = None, path = '', name = '', supers = None): + def __init__(self, ifs_data, obj, parent = None, path = '', name = '', + supers = None, super_disable = False, super_skip_bad = False, + super_abort_if_bad = False): # circular dependencies mean we import here from . import AfpFolder, TexFolder self.folder_handlers = { @@ -17,6 +20,9 @@ class GenericFolder(Node): 'tex' : TexFolder, } self.supers = supers if supers else [] + self.super_disable = super_disable + self.super_skip_bad = super_skip_bad + self.super_abort_if_bad = super_abort_if_bad Node.__init__(self, ifs_data, obj, parent, path, name) file_handler = GenericFile @@ -37,24 +43,51 @@ class GenericFolder(Node): if filename == '_info_': # metadata continue elif filename == '_super_': # sub-reference + if self.super_disable: + continue + super_file = join(my_path, child.text) if not isfile(super_file): - raise IOError('IFS references super-IFS {} but it does not exist'.format(child.text)) + raise IOError('IFS references super-IFS {} but it does not exist. Use --super-disable to ignore.'.format(child.text)) - self.supers.append(IFS(super_file)) + md5_expected = None + if list(child) and child[0].tag == 'md5': + md5_expected = bytearray.fromhex(child[0].text) + + super_ifs = IFS(super_file, super_skip_bad=self.super_skip_bad, + super_abort_if_bad=self.super_abort_if_bad) + super_ifs.md5_good = (super_ifs.manifest_md5 == md5_expected) # add our own sentinel + if not super_ifs.md5_good: + super_msg = 'IFS references super-IFS {} with MD5 {} but the actual MD5 is {}. One IFS may be corrupt.'.format( + child.text, md5_expected.hex(), super_ifs.manifest_md5.hex() + ) + if self.super_abort_if_bad: + raise IOError(super_msg + ' Aborting.') + elif self.super_skip_bad: + tqdm.write('WARNING: {} Skipping all files it contains.'.format(super_msg)) + else: + tqdm.write('WARNING: {}'.format(super_msg)) + + self.supers.append(super_ifs) # folder: has children or timestamp only, and isn't a reference elif (list(child) or len(child.text.split(' ')) == 1) and child[0].tag != 'i': handler = self.folder_handlers.get(filename, GenericFolder) - self.folders[filename] = handler(self.ifs_data, child, self, self.full_path, filename, self.supers) + self.folders[filename] = handler(self.ifs_data, child, self, self.full_path, filename, self.supers, + self.super_disable, self.super_skip_bad, self.super_abort_if_bad) else: # file - self.files[filename] = self.file_handler(self.ifs_data, child, self, self.full_path, filename) if list(child) and child[0].tag == 'i': + if self.super_disable: + continue + # backref super_ref = int(child[0].text) if super_ref > len(self.supers): raise IOError('IFS references super-IFS {} but we only have {}'.format(super_ref, len(self.supers))) super_ifs = self.supers[super_ref - 1] + if not super_ifs.md5_good and self.super_skip_bad: + continue + super_files = super_ifs.tree.all_files try: super_file = next(x for x in super_files if ( @@ -65,6 +98,8 @@ class GenericFolder(Node): raise IOError('IFS references super-IFS entry {} in {} but it does not exist'.format(filename, super_ifs.ifs_out)) self.files[filename] = super_file + else: + self.files[filename] = self.file_handler(self.ifs_data, child, self, self.full_path, filename) if not self.full_path: # root self.tree_complete() diff --git a/ifstools/handlers/MD5Folder.py b/ifstools/handlers/MD5Folder.py index d950a59..0dc8e53 100644 --- a/ifstools/handlers/MD5Folder.py +++ b/ifstools/handlers/MD5Folder.py @@ -6,8 +6,11 @@ from . import GenericFolder class MD5Folder(GenericFolder): - def __init__(self, ifs_data, parent, obj, path = '', name = '', supers = None, md5_tag = None, extension = None): - GenericFolder.__init__(self, ifs_data, parent, obj, path, name, supers) + def __init__(self, ifs_data, parent, obj, path = '', name = '', supers = None, + super_disable = False, super_skip_bad = False, + super_abort_if_bad = False, md5_tag = None, extension = None): + GenericFolder.__init__(self, ifs_data, parent, obj, path, name, supers, + super_disable, super_skip_bad, super_abort_if_bad) self.md5_tag = md5_tag if md5_tag else self.name self.extension = extension diff --git a/ifstools/handlers/TexFolder.py b/ifstools/handlers/TexFolder.py index d139080..fed2216 100644 --- a/ifstools/handlers/TexFolder.py +++ b/ifstools/handlers/TexFolder.py @@ -60,8 +60,10 @@ class ImageCanvas(GenericFile): return class TexFolder(MD5Folder): - def __init__(self, ifs_data, obj, parent = None, path = '', name = '', supers = None): - MD5Folder.__init__(self, ifs_data, obj, parent, path, name, supers, 'image', '.png') + def __init__(self, ifs_data, obj, parent = None, path = '', name = '', supers = None, + super_disable = False, super_skip_bad = False, super_abort_if_bad = False): + MD5Folder.__init__(self, ifs_data, obj, parent, path, name, supers, + super_disable, super_skip_bad, super_abort_if_bad, 'image', '.png') def tree_complete(self): MD5Folder.tree_complete(self) diff --git a/ifstools/ifs.py b/ifstools/ifs.py index fca7743..fa2cb6c 100644 --- a/ifstools/ifs.py +++ b/ifstools/ifs.py @@ -41,15 +41,17 @@ class FileBlob(object): return self.file.read(size) class IFS: - def __init__(self, path): + def __init__(self, path, super_disable = False, super_skip_bad = False, + super_abort_if_bad = False): if isfile(path): - self.load_ifs(path) + self.load_ifs(path, super_disable, super_skip_bad, super_abort_if_bad) elif isdir(path): self.load_dir(path) else: raise IOError('Input path {} does not exist'.format(path)) - def load_ifs(self, path): + def load_ifs(self, path, super_disable = False, super_skip_bad = False, + super_abort_if_bad = False): self.is_file = True name = basename(path) @@ -71,13 +73,16 @@ class IFS: manifest_end = header.get_u32() self.data_blob = FileBlob(self.file, manifest_end) + self.manifest_md5 = None if self.file_version > 1: - # md5 of manifest, unchecked - header.offset += 16 + self.manifest_md5 = header.get_bytes(16) self.file.seek(header.offset) self.manifest = KBinXML(self.file.read(manifest_end-header.offset)) - self.tree = GenericFolder(self.data_blob, self.manifest.xml_doc) + self.tree = GenericFolder(self.data_blob, self.manifest.xml_doc, + super_disable=super_disable, super_skip_bad=super_skip_bad, + super_abort_if_bad=super_abort_if_bad + ) # IFS files repacked with other tools usually have wrong values - don't validate this #assert ifs_tree_size == self.manifest.mem_size diff --git a/ifstools/ifstools.py b/ifstools/ifstools.py index 1cf0129..98b4d97 100644 --- a/ifstools/ifstools.py +++ b/ifstools/ifstools.py @@ -1,92 +1,99 @@ -import argparse -import os -import multiprocessing # for pyinstaller fixes -from sys import exit # exe freeze -try: - # py 2 - input = raw_input -except NameError: - # py 3 - pass - -from .ifs import IFS - -def get_choice(prompt): - while True: - q = input(prompt + ' [Y/n] ').lower() - if not q: - return True # default to yes - elif q == 'y': - return True - elif q == 'n': - return False - else: - print('Please answer y/n') - -def extract(i, args, path): - if args.progress: - print('Extracting...') - i.extract(path=path, **vars(args)) - -def repack(i, args, path): - if args.progress: - print('Repacking...') - i.repack(path=path, **vars(args)) - -def main(): - multiprocessing.freeze_support() # pyinstaller - parser = argparse.ArgumentParser(description='Unpack/pack IFS files and textures') - parser.add_argument('files', metavar='file_to_unpack.ifs|folder_to_repack_ifs', type=str, nargs='+', - help='files/folders to process. Files will be unpacked, folders will be repacked') - parser.add_argument('-e', '--extract-folders', action='store_true', help='do not repack folders, instead unpack any IFS files inside them', dest='extract_folders') - parser.add_argument('-y', action='store_true', help='don\'t prompt for file/folder overwrite', dest='overwrite') - parser.add_argument('-o', default='.', help='output directory', dest='out_dir') - parser.add_argument('--tex-only', action='store_true', help='only extract textures', dest='tex_only') - parser.add_argument('-c', '--canvas', action='store_true', help='dump the image canvas as defined by the texturelist.xml in _canvas.png', dest='dump_canvas') - parser.add_argument('--bounds', action='store_true', help='draw image bounds on the exported canvas in red', dest='draw_bbox') - parser.add_argument('--uv', action='store_true', help='crop images to uvrect (usually 1px smaller than imgrect). Forces --tex-only', dest='crop_to_uvrect') - parser.add_argument('--no-cache', action='store_false', help='ignore texture cache, recompress all', dest='use_cache') - parser.add_argument('--rename-dupes', action='store_true', dest='rename_dupes', +import argparse +import os +import multiprocessing # for pyinstaller fixes +from sys import exit # exe freeze +try: + # py 2 + input = raw_input +except NameError: + # py 3 + pass + +from .ifs import IFS + +def get_choice(prompt): + while True: + q = input(prompt + ' [Y/n] ').lower() + if not q: + return True # default to yes + elif q == 'y': + return True + elif q == 'n': + return False + else: + print('Please answer y/n') + +def extract(i, args, path): + if args.progress: + print('Extracting...') + i.extract(path=path, **vars(args)) + +def repack(i, args, path): + if args.progress: + print('Repacking...') + i.repack(path=path, **vars(args)) + +def main(): + multiprocessing.freeze_support() # pyinstaller + parser = argparse.ArgumentParser(description='Unpack/pack IFS files and textures') + parser.add_argument('files', metavar='file_to_unpack.ifs|folder_to_repack_ifs', type=str, nargs='+', + help='files/folders to process. Files will be unpacked, folders will be repacked') + parser.add_argument('-e', '--extract-folders', action='store_true', help='do not repack folders, instead unpack any IFS files inside them', dest='extract_folders') + parser.add_argument('-y', action='store_true', help='don\'t prompt for file/folder overwrite', dest='overwrite') + parser.add_argument('-o', default='.', help='output directory', dest='out_dir') + parser.add_argument('--tex-only', action='store_true', help='only extract textures') + parser.add_argument('-c', '--canvas', action='store_true', help='dump the image canvas as defined by the texturelist.xml in _canvas.png', dest='dump_canvas') + parser.add_argument('--bounds', action='store_true', help='draw image bounds on the exported canvas in red', dest='draw_bbox') + parser.add_argument('--uv', action='store_true', help='crop images to uvrect (usually 1px smaller than imgrect). Forces --tex-only', dest='crop_to_uvrect') + parser.add_argument('--no-cache', action='store_false', help='ignore texture cache, recompress all', dest='use_cache') + parser.add_argument('--rename-dupes', action='store_true', help='if two files have the same name but differing case (A.png vs a.png) rename the second as "a (1).png" to allow both to be extracted on Windows') - parser.add_argument('-m', '--extract-manifest', action='store_true', help='extract the IFS manifest for inspection', dest='extract_manifest') - parser.add_argument('-s', '--silent', action='store_false', dest='progress', - help='don\'t display files as they are processed') - parser.add_argument('-r', '--norecurse', action='store_false', dest='recurse', - help='if file contains another IFS, don\'t extract its contents') - - args = parser.parse_args() - - if args.crop_to_uvrect: - args.tex_only = True - - if args.extract_folders: - dirs = [f for f in args.files if os.path.isdir(f)] - # prune - args.files = [f for f in args.files if not os.path.isdir(f)] - # add the extras - for d in dirs: - args.files.extend((os.path.join(d,f) for f in os.listdir(d) if f.lower().endswith('.ifs'))) - - for f in args.files: - if args.progress: - print(f) - try: - i = IFS(f) - except IOError as e: - # human friendly - print('{}: {}'.format(os.path.basename(f), str(e))) - exit(1) - - path = os.path.join(args.out_dir, i.default_out) - if os.path.exists(path) and not args.overwrite: - if not get_choice('{} exists. Overwrite?'.format(path)): - continue - - if i.is_file: - extract(i, args, path) - else: - repack(i, args, path) - - -if __name__ == '__main__': - main() + parser.add_argument('-m', '--extract-manifest', action='store_true', help='extract the IFS manifest for inspection', dest='extract_manifest') + parser.add_argument('--super-disable', action='store_true', + help='only extract files unique to this IFS, do not follow "super" parent references at all') + parser.add_argument('--super-skip-bad', action='store_true', + help='if a "super" IFS reference has a checksum mismatch, do not extract it') + parser.add_argument('--super-abort-if-bad', action='store_true', + help='if a "super" IFS reference has a checksum mismatch, cancel and display an error') + parser.add_argument('-s', '--silent', action='store_false', dest='progress', + help='don\'t display files as they are processed') + parser.add_argument('-r', '--norecurse', action='store_false', dest='recurse', + help='if file contains another IFS, don\'t extract its contents') + + args = parser.parse_args() + + if args.crop_to_uvrect: + args.tex_only = True + + if args.extract_folders: + dirs = [f for f in args.files if os.path.isdir(f)] + # prune + args.files = [f for f in args.files if not os.path.isdir(f)] + # add the extras + for d in dirs: + args.files.extend((os.path.join(d,f) for f in os.listdir(d) if f.lower().endswith('.ifs'))) + + for f in args.files: + if args.progress: + print(f) + try: + i = IFS(f, super_disable=args.super_disable, super_skip_bad=args.super_skip_bad, + super_abort_if_bad=args.super_abort_if_bad) + except IOError as e: + # human friendly + print('{}: {}'.format(os.path.basename(f), str(e))) + exit(1) + + path = os.path.join(args.out_dir, i.default_out) + if os.path.exists(path) and not args.overwrite: + if not get_choice('{} exists. Overwrite?'.format(path)): + continue + + if i.is_file: + extract(i, args, path) + else: + repack(i, args, path) + + +if __name__ == '__main__': + main()