diff --git a/handlers/GenericFile.py b/handlers/GenericFile.py index 112eead..4fb807d 100644 --- a/handlers/GenericFile.py +++ b/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) diff --git a/handlers/GenericFolder.py b/handlers/GenericFolder.py index 4fe7fcf..186fb4b 100644 --- a/handlers/GenericFolder.py +++ b/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 = '' diff --git a/handlers/ImageFile.py b/handlers/ImageFile.py index d8a1bff..0c78bcd 100644 --- a/handlers/ImageFile.py +++ b/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 - uncompressed_size = len(data) - data = lz77.compress(data) - compressed_size = len(data) - data = pack('>I', uncompressed_size) + pack('>I', compressed_size) + data + data = self.read_cache(recache) + else: + data = self._load_im() # 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 + diff --git a/handlers/lz77.py b/handlers/lz77.py index e5bde72..4d08359 100644 --- a/handlers/lz77.py +++ b/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): diff --git a/ifstools.py b/ifstools.py index 567add2..b92c5dc 100644 --- a/ifstools.py +++ b/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.default_out = join(dirname(path), out) - self.ifs_out = path + self.ifs_out = basename(path) + self.default_out = splitext(self.ifs_out)[0] + '_ifs' 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 - self.ifs_out = path.replace('_ifs', '.ifs') + path = path.rstrip('/\\') + '/' + 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 - if len(sys.argv) < 2: - print('ifstools filename.ifs OR folder_ifs') - exit() - i = IFS(sys.argv[1]) - if i.is_file: - i.extract_all() - else: - i.repack() + parser = argparse.ArgumentParser(description='Unpack/pack IFS files and textures') + parser.add_argument('files', metavar='file.ifs|folder_ifs', type=str, nargs='+', + help='files/folders to process. Files will be unpacked, folders will be repacked') + parser.add_argument('--recache', action='store_true', help='ignore texture cache, recompress all') + 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(args.progress, args.recurse) + else: + i.repack(args.progress, args.recache)