added stuff

read_gzp.py

@@ -1,10 +1,13 @@
 import logging
 import argparse
 import struct
+import os.path
+
+logger = None
 
 
 def main():
-    objects = []
+    global logger
     _ch = logging.StreamHandler()
     _ch.setLevel("DEBUG")
     _formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -13,10 +16,28 @@ def main():
     logger.addHandler(_ch)
     logger.setLevel("DEBUG")
     parser = argparse.ArgumentParser()
-    parser.add_argument("gzp_file", help="gzp file to read")
+    parser.add_argument("gzp_file_directory", help="gzp file or directory to read")
+    parser.add_argument("output_dir", help="output directory to send all extracted files")
+    parser.add_argument("--extensions", help="extensions to extract, separated by a comma (ex: 'tga,wav' will extract only .tga and .wav files)")
     args = parser.parse_args()
-    gzp_file = args.gzp_file
+    gzp_file = args.gzp_file_directory
+    output_dir = args.output_dir
+    extensions = args.extensions
 
+    if extensions:
+        extensions = extensions.split(",")
+
+    if os.path.isdir(gzp_file):
+        logger.info("Reading .gzp files in "+gzp_file)
+        files = os.listdir(gzp_file)
+        for file in files:
+            if file.endswith(".gzp"):
+                extract_gzp(file, extensions, output_dir)
+    else:
+        extract_gzp(gzp_file, extensions, output_dir)
+
+
+def extract_gzp(gzp_file, extensions, output_dir):
     logger.info("Opening " + gzp_file)
     with open(gzp_file, "rb") as gzp_fp:
         checksum = read_int(gzp_fp)
@@ -42,13 +63,23 @@ def main():
             logger.info("Reading index " + str(index))
             compressed_size = read_int(gzp_fp)
             original_size = read_int(gzp_fp)
-            # original_size = read_int(gzp_fp)
             file_time = read_int(gzp_fp)
             content_offset = read_int(gzp_fp) + 16
             compression = read_byte(gzp_fp)  # compression: 1 if compressed else 0
             name_length = read_byte(gzp_fp)
             name = read_bytes(gzp_fp, name_length).decode("utf8").strip('\x00')
-            logger.info(name + " compression: "+str(compression) + ", filesize: "+sizeof_fmt(original_size) + ", start: "+str(content_offset))
+            logger.info(name + " compression: " + str(compression) + ", filesize: " + sizeof_fmt(
+                original_size) + ", start: " + str(content_offset))
+
+            if extensions:
+                extract_file = False
+                for extension in extensions:
+                    if name.endswith(extension):
+                        extract_file = True
+                        break
+                if not extract_file:
+                    logger.info("File " + name + " does not match any of wanted extension. Skipping...")
+                    continue
 
             curr_pos = gzp_fp.tell()
             gzp_fp.seek(content_offset)
@@ -60,11 +91,11 @@ def main():
                 buffer = decompress(buffer, original_size)
 
             logger.info("Writing file " + name)
-            with open(name, "wb") as entry_fp:
+            if not os.path.exists(output_dir):
+                os.mkdir(output_dir)
+            with open(os.path.join(output_dir, name), "wb") as entry_fp:
                 entry_fp.write(buffer)
 
-            # separator = read_byte(gzp_fp)
-
 
 def sizeof_fmt(num, suffix='B'):
     for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
@@ -89,42 +120,40 @@ def read_int(fp):
 def decompress(buffer, finalsize):
     i = 0
     j = 0
-    decByte = 0
-    decBits = 8
-    buffStart = 0xFEE
+    dec_byte = 0
+    dec_bits = 8
+    buff_start = 0xFEE
     res = bytearray(finalsize)
 
     if finalsize == 0:
         return res
 
     while j < finalsize:
-        if decBits == 8:
-            decByte = buffer[i]
+        if dec_bits == 8:
+            dec_byte = buffer[i]
             i += 1
-            decBits = 0
+            dec_bits = 0
 
-        if (decByte >> decBits & 1) == 0:
-            decPos = ((buffer[i] + ((buffer[i + 1] & 0xF0) << 4) - buffStart - j) & 0xFFF) - 0x1000 + j
-            decLen = (buffer[i + 1] & 0xF) + 3
+        if (dec_byte >> dec_bits & 1) == 0:
+            dec_pos = ((buffer[i] + ((buffer[i + 1] & 0xF0) << 4) - buff_start - j) & 0xFFF) - 0x1000 + j
+            dec_len = (buffer[i + 1] & 0xF) + 3
             i += 2
 
-            while decLen > 0:
-                if decPos >= 0:
-                    # print("j:" +str(j))
-                    # print("decPos: "+str(decPos))
-                    res[j] = res[decPos]
+            while dec_len > 0:
+                if dec_pos >= 0:
+                    res[j] = res[dec_pos]
                 else:
                     res[j] = 32
 
                 j += 1
-                decPos += 1
-                decLen -= 1
+                dec_pos += 1
+                dec_len -= 1
         else:
             res[j] = buffer[i]
             i += 1
             j += 1
 
-        decBits += 1
+        dec_bits += 1
 
     return res