Update pzztool.py

pzztool.py

@@ -1,12 +1,16 @@
 #!/usr/bin/env python3
-__version__ = "1.0"
+__version__ = "1.1"
 __author__  = "rigodron, algoflash, GGLinnk"
 __OriginalAutor__ = "infval"
 
-from os import listdir, path, stat
+from os import listdir
 from pathlib import Path
 from struct import unpack, pack
+from math import ceil
 
+BIT_COMPRESSION_FLAG = 0x40000000
+FILE_LENGTH_MASK = 0x3FFFFFFF
+CHUNK_SIZE = 0x800
 
 def pzz_decompress(compressed_bytes: bytes):
     uncompressed_bytes = bytearray()
@@ -49,7 +53,6 @@ def pzz_decompress(compressed_bytes: bytes):
 
     return uncompressed_bytes
 
-
 def bytes_align(bout: bytes):
     success = False
     while not success:
@@ -58,7 +61,6 @@ def bytes_align(bout: bytes):
         if hex(address).endswith("00"):
             break
 
-
 def pzz_compress(b):
     bout = bytearray()
     size_b = len(b) // 2 * 2
@@ -134,32 +136,31 @@ def pzz_compress(b):
 
     return bout
 
-
+def pzz_unpack(pzz_path):
-def pzz_unpack(path, dir_path):
     # Script BMS pour les pzz de ps2 (GioGio's adventure) -> https://zenhax.com/viewtopic.php?f=9&t=8724&p=39437#p39437
+    unpacked_pzz_path = Path( Path(pzz_path).stem )
+    unpacked_pzz_path.mkdir(exist_ok=True)
 
-    with open(path, "rb") as f:
+    with open(pzz_path, "rb") as pzz_file:
         # file_count reçoit le nombre de fichiers présent dans le PZZ :
-        file_count, = unpack(">I", f.read(4)) # On lit les 4 premiers octets (uint32 big-endian)
+        file_count, = unpack(">I", pzz_file.read(4)) # On lit les 4 premiers octets (uint32 big-endian)
 
         # files_descriptors reçoit un tuple avec l'ensemble des descripteurs de fichiers (groupes d'uint32 big-endian)
-        files_descriptors = unpack(">{}I".format(file_count), f.read(file_count * 4))
+        files_descriptors = unpack(">{}I".format(file_count), pzz_file.read(file_count * 4))
 
         print("File count:", file_count)
 
-        offset = 0x800
+        offset = CHUNK_SIZE
         for i, file_descriptor in enumerate(files_descriptors): # on parcours le tuple de descripteurs de fichiers
 
-            is_compressed = (file_descriptor & 0x40000000) != 0 # Le bit 30 correspond au flag de compression (bits numérotés de 0 à 31)
+            is_compressed = (file_descriptor & BIT_COMPRESSION_FLAG) != 0 # Le bit 30 correspond au flag de compression (bits numérotés de 0 à 31)
-            print(file_descriptor)
 
-            # file_descriptor reçoit maintenant les 30 premiers bits : (la taille / 0x800) 
+            # file_descriptor reçoit maintenant les 30 premiers bits : (la taille / CHUNK_SIZE) 
-            file_descriptor &= 0x3FFFFFFF 
+            file_descriptor &= FILE_LENGTH_MASK
-            print(file_descriptor)
 
             # file_len reçoit la taille du fichier
-            # la taille du fichier est un multiple de 0x800, on paddera avec des 0 jusqu'au fichier suivant
+            # la taille du fichier est un multiple de CHUNK_SIZE, on paddera avec des 0 jusqu'au fichier suivant
-            file_len = file_descriptor * 0x800 # file_len contient alors la taille du fichier en octets
+            file_len = file_descriptor * CHUNK_SIZE # file_len contient alors la taille du fichier en octets
             
             # Si la taille est nulle, on passe au descripteur de fichier suivant
             if file_len == 0:
@@ -171,98 +172,68 @@ def pzz_unpack(path, dir_path):
                 comp_str = "_compressed"
 
             # On forme le nom du nouveau fichier que l'on va extraire
-            filename = "{}_{:03}{}".format(Path(path).stem, i, comp_str)
+            filename = "{}_{:03}{}".format(Path(pzz_path).stem, i, comp_str)
-            file_path = (Path(dir_path) / filename).with_suffix(".dat")
+            file_path = (Path(unpacked_pzz_path) / filename).with_suffix(".dat")
 
             print("Offset: {:010} - {}".format(offset, file_path))
 
             # On se positionne au début du fichier dans l'archive
-            f.seek(offset)
+            pzz_file.seek(offset)
             # On extrait notre fichier
-            file_path.write_bytes(f.read(file_len))
+            file_path.write_bytes(pzz_file.read(file_len))
 
             # Enfin, on ajoute la taille du fichier afin de pointer sur le fichier suivant
-            # La taille du fichier étant un multiple de 0x800, on aura complété les 2048 octets finaux avec des 0x00
+            # La taille du fichier étant un multiple de CHUNK_SIZE, on aura complété les 2048 octets finaux avec des 0x00
             offset += file_len
 
-def pzz_pack(src, dir_path):
+def pzz_pack(src_path):
-    bout = bytearray()
+    # On récupère les fichiers du dossier à compresser
-    filebout = bytearray()
+    src_files = listdir(src_path)
-    file_count = 0;
-    files = []
 
-    linkPath = path.normpath(dir_path)
+    # On récupère le nombre total de fichiers
-    linkFiles = [f for f in listdir(linkPath) if path.isfile(path.join(linkPath, f))]
+    file_count = int(src_files[-1].split("_")[1][0:3]) + 1
 
-    for file in linkFiles:
+    print(str(file_count) + " files to pack in " + str(src_path.with_suffix(".pzz")))
-        if (str(src)[12:-18] in file):
-            file_count += 1
-            files.append(file)
 
-    is_odd_number = (file_count % 2) != 0
+    with src_path.with_suffix(".pzz").open("wb") as pzz_file :
+        # On écrit file_count au début de header
+        pzz_file.write(file_count.to_bytes(4, byteorder='big'))
 
-    if (file_count == 6 or file_count == 12):
+        # On écrit les file_descriptor dans le header du PZZ pour chaque fichier
-        file_count += 4
+        last_index = 0 # permet d'ajouter les file_descriptor=NULL
-        for i, file in enumerate(files):
+        for src_file_name in src_files :
-            count = int(0x40 << 24) + int(path.getsize(linkPath + "/" + file) / 0x800)
+            index = int(src_file_name.split("_")[1][0:3])
+            is_compressed = ( len(src_file_name.split("_compressed")) > 1 )
 
-            if (i == 1 or i == 3 or i == 5 or i == 7):
+            # On ajoute les file_descriptor=NULL
-                filebout.extend(b"\x00\x00\x00\x00")
+            while(last_index < index):
-                filebout.extend(pack(">I", count))
+                pzz_file.write(b"\x00\x00\x00\x00")
-            else:
+                last_index += 1
-                filebout.extend(pack(">I", count))
 
-        file_count = pack(">I", file_count)
+            # file_descriptor = arrondi supérieur de la taille / CHUNK_SIZE
-        bout.extend(file_count)
+            file_descriptor = ceil( (src_path / src_file_name).stat().st_size / CHUNK_SIZE)
-        bout.extend(filebout)
 
-    elif (file_count == 6 or file_count == 14):
+            # On ajoute le flag de compression au file_descriptor
-        file_count += 2
+            if is_compressed :
-        for i, file in enumerate(files):
+                file_descriptor |= BIT_COMPRESSION_FLAG
-            count = int(0x40 << 24) + int(path.getsize(linkPath + "/" + file) / 0x800)
 
-            if (i == 1 or i == 3):
+            # On ecrit le file_descriptor
-                filebout.extend(b"\x00\x00\x00\x00")
+            pzz_file.write(file_descriptor.to_bytes(4, byteorder='big'))
-                filebout.extend(pack(">I", count))
+            last_index += 1
-            else:
-                filebout.extend(pack(">I", count))
 
-        file_count = pack(">I", file_count)
+        # On se place à la fin du header PZZ
-        bout.extend(file_count)
+        pzz_file.seek(CHUNK_SIZE)
-        bout.extend(filebout)
 
-    elif is_odd_number:
+        # On écrit tous les fichiers à la suite du header
-        file_count += 1
+        for src_file_name in src_files :
-        for i, file in enumerate(files):
+            is_compressed = ( len(src_file_name.split("_compressed")) > 1 )
-            count = int(0x40 << 24) + int(path.getsize(linkPath + "/" + file) / 0x800)
             
-            if (i == 1):
+            with (src_path / src_file_name).open("rb") as src_file :
-                filebout.extend(b"\x00\x00\x00\x00")
+                pzz_file.write( src_file.read() )
-                filebout.extend(pack(">I", count))
-            else:
-                filebout.extend(pack(">I", count))
 
-        file_count = pack(">I", file_count)
+                # Si le fichier n'est pas compressé, on ajoute le padding pour correspondre à un multiple de CHUNK_SIZE
-        bout.extend(file_count)
+                if not is_compressed and (src_file.tell() % CHUNK_SIZE) > 0:
-        bout.extend(filebout)
+                    pzz_file.write( b"\x00" * (CHUNK_SIZE - (src_file.tell() % CHUNK_SIZE)) )
-
-    success = False
-
-    while not success:
-        bout.extend(b"\x00\x00")
-        address = len(bout)
-        if hex(address).endswith("800"):
-            break
-    for file in files:
-        filebout = open(linkPath + "/" + file, "rb")
-        data = filebout.read()
-        bout.extend(data)
-    filename = "{}".format(str(src)[12:-19])
-    p = (Path(dir_path) / filename).with_suffix(".pzz")
-    p.write_bytes(bout)
-
-def pzz_test():
-    print(pack(">I", int(0x40 << 24) + int(stat(linkPath + "/" + file).st_size) / 0x800))
 
 def get_argparser():
     import argparse
@@ -271,17 +242,15 @@ def get_argparser():
     parser.add_argument('input_path',  metavar='INPUT',  help='only relative if -bu, -bc, -bd, p')
     parser.add_argument('output_path', metavar='OUTPUT', help='directory if -u, -bu, -bc, -bd')
     group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument('-u', '--unpack', action='store_true', help='PZZ files from AFS')
+    group.add_argument('-u', '--unpack',            action='store_true', help='-u source_pzz.pzz : Unpack the pzz in source_pzz folder')
-    group.add_argument('-c', '--compress', action='store_true')
+    group.add_argument('-p', '--pack',              action='store_true', help="-p source_folder : Pack source_folder in source_folder.pzz")
+    group.add_argument('-c', '--compress',          action='store_true', help='')
     group.add_argument('-d', '--decompress',        action='store_true', help='Unpacked files from PZZ')
     group.add_argument('-bu', '--batch-unpack',     action='store_true', help='INPUT relative pattern; e.g. AFS_DATA\\*.pzz')
     group.add_argument('-bc', '--batch-compress',   action='store_true', help='INPUT relative pattern; e.g. AFS_DATA\\*.bin')
     group.add_argument('-bd', '--batch-decompress', action='store_true', help='INPUT relative pattern; e.g. AFS_DATA\\*_compressed.dat')
-    group.add_argument('-p', '--pack', action='store_true')
-    group.add_argument('-t', '--test', action='store_true')
     return parser
 
-
 if __name__ == '__main__':
     import sys
     parser = get_argparser()
@@ -318,14 +287,10 @@ if __name__ == '__main__':
                 print("! Wrong PZZ file")
     elif args.pack:
         print("### Pack")
-        p_output.mkdir(exist_ok=True)
+        pzz_pack(p_input)
-        pzz_pack(p_input, p_output)
-    elif args.test:
-        pzz_test()
     elif args.unpack:
         print("### Unpack")
-        p_output.mkdir(exist_ok=True)
+        pzz_unpack(p_input)
-        pzz_unpack(p_input, p_output)
     #elif args.batch_pack:
     #    pass
     elif args.batch_unpack:
@@ -336,5 +301,3 @@ if __name__ == '__main__':
         for filename in p.glob(args.input_path):
             print(filename)
             pzz_unpack(filename, p_output)
-
-