Super files: MD5 checks and balances

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()

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):
+    def __init__(self, ifs_data, parent, obj, path = '', name = '', supers = None,
-        GenericFolder.__init__(self, ifs_data, parent, obj, path, name, supers)
+            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
 

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):
+    def __init__(self, ifs_data, obj, parent = None, path = '', name = '', supers = None,
-        MD5Folder.__init__(self, ifs_data, obj, parent, path, name, supers, 'image', '.png')
+            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)

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
+            self.manifest_md5 = header.get_bytes(16)
-            header.offset += 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

ifstools/ifstools.py

@@ -1,92 +1,99 @@
-import argparse
+import argparse
-import os
+import os
-import multiprocessing # for pyinstaller fixes
+import multiprocessing # for pyinstaller fixes
-from sys import exit # exe freeze
+from sys import exit # exe freeze
-try:
+try:
-    # py 2
+    # py 2
-    input = raw_input
+    input = raw_input
-except NameError:
+except NameError:
-    # py 3
+    # py 3
-    pass
+    pass
-
+
-from .ifs import IFS
+from .ifs import IFS
-
+
-def get_choice(prompt):
+def get_choice(prompt):
-    while True:
+    while True:
-        q = input(prompt + ' [Y/n] ').lower()
+        q = input(prompt + ' [Y/n] ').lower()
-        if not q:
+        if not q:
-            return True # default to yes
+            return True # default to yes
-        elif q == 'y':
+        elif q == 'y':
-            return True
+            return True
-        elif q == 'n':
+        elif q == 'n':
-            return False
+            return False
-        else:
+        else:
-            print('Please answer y/n')
+            print('Please answer y/n')
-
+
-def extract(i, args, path):
+def extract(i, args, path):
-    if args.progress:
+    if args.progress:
-        print('Extracting...')
+        print('Extracting...')
-    i.extract(path=path, **vars(args))
+    i.extract(path=path, **vars(args))
-
+
-def repack(i, args, path):
+def repack(i, args, path):
-    if args.progress:
+    if args.progress:
-        print('Repacking...')
+        print('Repacking...')
-    i.repack(path=path, **vars(args))
+    i.repack(path=path, **vars(args))
-
+
-def main():
+def main():
-    multiprocessing.freeze_support() # pyinstaller
+    multiprocessing.freeze_support() # pyinstaller
-    parser = argparse.ArgumentParser(description='Unpack/pack IFS files and textures')
+    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='+',
+    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')
+                       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('-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('-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('-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('--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('-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('--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('--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('--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',
+    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('-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',
+    parser.add_argument('--super-disable', action='store_true',
-                       help='don\'t display files as they are processed')
+                       help='only extract files unique to this IFS, do not follow "super" parent references at all')
-    parser.add_argument('-r', '--norecurse', action='store_false', dest='recurse',
+    parser.add_argument('--super-skip-bad', action='store_true',
-                       help='if file contains another IFS, don\'t extract its contents')
+                       help='if a "super" IFS reference has a checksum mismatch, do not extract it')
-
+    parser.add_argument('--super-abort-if-bad', action='store_true',
-    args = parser.parse_args()
+                       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',
-    if args.crop_to_uvrect:
+                       help='don\'t display files as they are processed')
-        args.tex_only = True
+    parser.add_argument('-r', '--norecurse', action='store_false', dest='recurse',
-
+                       help='if file contains another IFS, don\'t extract its contents')
-    if args.extract_folders:
+
-        dirs = [f for f in args.files if os.path.isdir(f)]
+    args = parser.parse_args()
-        # prune
+
-        args.files = [f for f in args.files if not os.path.isdir(f)]
+    if args.crop_to_uvrect:
-        # add the extras
+        args.tex_only = True
-        for d in dirs:
+
-            args.files.extend((os.path.join(d,f) for f in os.listdir(d) if f.lower().endswith('.ifs')))
+    if args.extract_folders:
-
+        dirs = [f for f in args.files if os.path.isdir(f)]
-    for f in args.files:
+        # prune
-        if args.progress:
+        args.files = [f for f in args.files if not os.path.isdir(f)]
-            print(f)
+        # add the extras
-        try:
+        for d in dirs:
-            i = IFS(f)
+            args.files.extend((os.path.join(d,f) for f in os.listdir(d) if f.lower().endswith('.ifs')))
-        except IOError as e:
+
-            # human friendly
+    for f in args.files:
-            print('{}: {}'.format(os.path.basename(f), str(e)))
+        if args.progress:
-            exit(1)
+            print(f)
-
+        try:
-        path = os.path.join(args.out_dir, i.default_out)
+            i = IFS(f, super_disable=args.super_disable, super_skip_bad=args.super_skip_bad,
-        if os.path.exists(path) and not args.overwrite:
+                super_abort_if_bad=args.super_abort_if_bad)
-            if not get_choice('{} exists. Overwrite?'.format(path)):
+        except IOError as e:
-                continue
+            # human friendly
-
+            print('{}: {}'.format(os.path.basename(f), str(e)))
-        if i.is_file:
+            exit(1)
-            extract(i, args, path)
+
-        else:
+        path = os.path.join(args.out_dir, i.default_out)
-            repack(i, args, path)
+        if os.path.exists(path) and not args.overwrite:
-
+            if not get_choice('{} exists. Overwrite?'.format(path)):
-
+                continue
-if __name__ == '__main__':
+
-    main()
+        if i.is_file:
+            extract(i, args, path)
+        else:
+            repack(i, args, path)
+
+
+if __name__ == '__main__':
+    main()