cmdline args, texture compression caching

handlers/GenericFile.py

@@ -53,7 +53,7 @@ class GenericFile(object):
         self.size = len(ret)
         return ret
 
-    def repack(self, manifest, data_blob, progress):
+    def repack(self, manifest, data_blob, progress, recache):
         if progress:
             print(self.name)
         elem = etree.SubElement(manifest, self.packed_name)

handlers/GenericFolder.py

@@ -53,7 +53,7 @@ class GenericFolder():
 
         return cls(ifs, name, time, files, folders)
 
-    def repack(self, manifest, data_blob, progress):
+    def repack(self, manifest, data_blob, progress, recache):
         if self.name:
             manifest = etree.SubElement(manifest, self.packed_name)
             manifest.attrib['__type'] = 's32'
@@ -62,7 +62,7 @@ class GenericFolder():
                 print(self.name)
 
         for name, entry in chain(self.folders.items(), self.files.items()):
-            entry.repack(manifest, data_blob, progress)
+            entry.repack(manifest, data_blob, progress, recache)
 
     def tostring(self, indent = 0):
         ret = ''

handlers/ImageFile.py

@@ -1,5 +1,7 @@
 from io import BytesIO
 from struct import unpack, pack
+from os.path import getmtime, isfile, join, dirname
+from os import utime, mkdir
 
 from PIL import Image
 import lxml.etree as etree
@@ -83,24 +85,14 @@ class ImageFile(GenericFile):
         im.save(b, format = 'PNG')
         return b.getvalue()
 
-    def repack(self, manifest, data_blob, progress):
+    def repack(self, manifest, data_blob, progress, recache):
         if progress:
             print(self.name)
 
-        data = self.load()
-
-        im = Image.open(BytesIO(data))
-        if self.format == 'argb8888rev':
-            data = im.tobytes('raw', 'BGRA')
-        else:
-            raise NotImplementedError('Unknown format {}'.format(self.format))
-
         if self.compress == 'avslz':
-            o = data
+            data = self.read_cache(recache)
-            uncompressed_size = len(data)
+        else:
-            data = lz77.compress(data)
+            data = self._load_im()
-            compressed_size = len(data)
-            data = pack('>I', uncompressed_size) + pack('>I', compressed_size) + data
 
         # offset, size, timestamp
         elem = etree.SubElement(manifest, self.packed_name)
@@ -108,3 +100,45 @@ class ImageFile(GenericFile):
         elem.text = '{} {} {}'.format(len(data_blob.getvalue()), len(data), self.time)
         data_blob.write(data)
 
+    def _load_im(self):
+        data = self.load()
+
+        im = Image.open(BytesIO(data))
+        if im.mode != 'RGBA':
+            im = im.convert('RGBA')
+        if self.format == 'argb8888rev':
+            data = im.tobytes('raw', 'BGRA')
+        else:
+            raise NotImplementedError('Unknown format {}'.format(self.format))
+
+        return data
+
+    def write_cache(self, data):
+        cache = join(dirname(self.path), '_cache', self.name)
+        self._mkdir(dirname(cache))
+        with open(cache, 'wb') as f:
+            f.write(data)
+        utime(cache, (self.time, self.time))
+
+    def read_cache(self, recache):
+        cache = join(dirname(self.path), '_cache', self.name)
+        if isfile(cache) and not recache:
+            mtime = int(getmtime(cache))
+            if self.time <= mtime:
+                with open(cache, 'rb') as f:
+                    return f.read()
+        print('Not cached/out of date, compressing')
+        data = self._load_im()
+        uncompressed_size = len(data)
+        data = lz77.compress(data)
+        compressed_size = len(data)
+        data = pack('>I', uncompressed_size) + pack('>I', compressed_size) + data
+        self.write_cache(data)
+        return data
+
+    def _mkdir(self, dir):
+        try:
+            mkdir(dir)
+        except FileExistsError:
+            pass
+

handlers/lz77.py

@@ -3,6 +3,8 @@ from builtins import bytes
 from struct import unpack, pack
 from io import BytesIO
 
+from tqdm import tqdm
+
 WINDOW_SIZE = 0x1000
 WINDOW_MASK = WINDOW_SIZE - 1
 THRESHOLD = 3
@@ -57,7 +59,9 @@ def match_window(in_data, offset):
 
     return None
 
-def compress(input):
+def compress(input, progress = True):
+    pbar = tqdm(total = len(input), leave = False, unit = 'b', unit_scale = True,
+                desc = 'Compressing', disable = not progress)
     compressed = bytearray()
     input = bytes([0]*WINDOW_SIZE) + bytes(input)
     input_size = len(input)
@@ -77,9 +81,11 @@ def compress(input):
                     buf.extend(pack('>H', info))
                     bit = 0
                     current_pos += length
+                    pbar.update(length)
                 else:
                     buf.append(input[current_pos])
                     current_pos += 1
+                    pbar.update(1)
                     bit = 1
             flag_byte = (flag_byte >> 1) | ((bit & 1) << 7)
         compressed.append(flag_byte)
@@ -88,6 +94,7 @@ def compress(input):
     compressed.append(0)
     compressed.append(0)
 
+    pbar.close()
     return bytes(compressed)
 
 def compress_dummy(input):

ifstools.py

@@ -4,6 +4,7 @@ from io import BytesIO
 import hashlib
 import lxml.etree as etree
 from time import time as unixtime
+import argparse
 
 from kbinxml.kbinxml import KBinXML
 from kbinxml.bytebuffer import ByteBuffer
@@ -27,9 +28,8 @@ class IFS:
             raise IOError('Input path does not exist')
 
     def _load_ifs(self, path):
-        out = splitext(basename(path))[0] + '_ifs'
+        self.ifs_out = basename(path)
-        self.default_out = join(dirname(path), out)
+        self.default_out = splitext(self.ifs_out)[0] + '_ifs'
-        self.ifs_out = path
 
         with open(path, 'rb') as f:
             self.file = f.read()
@@ -52,8 +52,9 @@ class IFS:
         assert self.tree_size == self._tree_size()
 
     def _load_dir(self, path):
-        self.default_out = path
+        path = path.rstrip('/\\') + '/'
-        self.ifs_out = path.replace('_ifs', '.ifs')
+        self.default_out = dirname(path)
+        self.ifs_out = self.default_out.replace('_ifs', '.ifs')
 
         self.file_version = FILE_VERSION
         self.time = int(getmtime(path))
@@ -79,7 +80,10 @@ class IFS:
         tree['files'] = files
         tree['folders'] = []
         for dir in dirs:
-            tree['folders'].append(self._create_dir_tree_recurse(walker))
+            subdir = self._create_dir_tree_recurse(walker)
+            # this should probably be moved to TexFolder.py
+            if basename(subdir['path']) != '_cache':
+                tree['folders'].append(subdir)
 
         return tree
 
@@ -97,7 +101,7 @@ class IFS:
                 f.write(self.manifest.to_text().encode('utf8'))
         self._extract_tree(self.tree, progress, recurse)
 
-    def repack(self, progress = True, path = None):
+    def repack(self, progress = True, recache = False, path = None):
         if path is None:
             path = self.ifs_out
         data_blob = BytesIO()
@@ -106,7 +110,7 @@ class IFS:
         manifest_info = etree.SubElement(self.manifest.xml_doc, '_info_')
 
         # the important bit
-        self.tree.repack(self.manifest.xml_doc, data_blob, progress)
+        self.tree.repack(self.manifest.xml_doc, data_blob, progress, recache)
         data = data_blob.getvalue()
 
         data_md5 = etree.SubElement(manifest_info, 'md5')
@@ -167,7 +171,7 @@ class IFS:
             self._save_with_time(out, data, f.time)
             if recurse and f.name.endswith('.ifs'):
                 i = IFS(out)
-                i.extract_all()
+                i.extract_all(progress, recurse)
 
         for name, f in tree.folders.items():
             self._extract_tree(f, progress, recurse, join(dir, f.name))
@@ -176,7 +180,6 @@ class IFS:
         timestamp = tree.time if tree.time else self.time
         utime(outdir, (timestamp, timestamp))
 
-
     def _mkdir(self, dir):
         try:
             mkdir(dir)
@@ -196,12 +199,21 @@ class IFS:
         utime(filename, (time,time))
 
 if __name__ == '__main__':
-    import sys
+    parser = argparse.ArgumentParser(description='Unpack/pack IFS files and textures')
-    if len(sys.argv) < 2:
+    parser.add_argument('files', metavar='file.ifs|folder_ifs', type=str, nargs='+',
-        print('ifstools filename.ifs OR folder_ifs')
+                       help='files/folders to process. Files will be unpacked, folders will be repacked')
-        exit()
+    parser.add_argument('--recache', action='store_true', help='ignore texture cache, recompress all')
-    i = IFS(sys.argv[1])
+    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()
+
+    for f in args.files:
+        i = IFS(f)
+        path = None
         if i.is_file:
-        i.extract_all()
+            i.extract_all(args.progress, args.recurse)
         else:
-        i.repack()
+            i.repack(args.progress, args.recache)