From 9bd51ae6c7693c6bb970b58786537d2923139aec Mon Sep 17 00:00:00 2001
From: Mc-muffin <8714476+Mc-muffin@users.noreply.github.com>
Date: Tue, 16 Jan 2024 22:30:13 -0500
Subject: [PATCH] Rework isotool dumping so it's less cringe

---
 isotool.py | 287 ++++++++++++++++++++++++++++-------------------------
 1 file changed, 151 insertions(+), 136 deletions(-)

diff --git a/isotool.py b/isotool.py
index 6418b8d..eacaf47 100644
--- a/isotool.py
+++ b/isotool.py
@@ -5,15 +5,18 @@ from pathlib import Path
 from dataclasses import dataclass
 import typing
 
-SCRIPT_VERSION = "1.1"
+SCRIPT_VERSION = "1.2"
 SECTOR_SIZE = 0x800
 SYSTEM_AREA_SIZE = 0x10 * SECTOR_SIZE
+LAYER0_PVD_LOCATION = SYSTEM_AREA_SIZE
+
 
 @dataclass
 class FileListData:
     path: Path
     inode: int
     lba: int
+    size: int = 0
 
 
 @dataclass
@@ -25,6 +28,7 @@ class FileListInfo:
 def main():
     print(f"pyPS2 ISO Rebuilder v{SCRIPT_VERSION}")
     print("Original by Raynê Games")
+    print()
 
     args = get_arguments()
 
@@ -114,168 +118,179 @@ def check_pvd(fp: typing.BinaryIO, pvd_loc: int) -> bool:
         return False
 
 
+def dump_dir_records(iso: typing.BinaryIO, pvd_loc: int, pvd_off: int) -> FileListInfo:
+    path_parts = []
+    record_ends = []
+    record_pos = []
+    file_info = FileListInfo([], 0)
+
+    # get the root directory record off the PVD
+    iso.seek(pvd_loc + 0x9E)
+    dr_data_pos, dr_data_len = struct.unpack("<I4xI", iso.read(12))
+    dr_data_pos *= SECTOR_SIZE
+    dr_data_pos += pvd_off
+    record_ends.append(dr_data_pos + dr_data_len)
+    record_pos.append(0)
+
+    iso.seek(dr_data_pos)
+
+    # Traverse directory records recursively
+    # Did I mention that I won't do function recursion?
+    while True:
+        # Have we reached the end of current dir record?
+        if iso.tell() >= record_ends[-1]:
+            if len(record_ends) == 1:
+                # If it's the last one, we finished
+                break
+            else:
+                # Otherwise keep reading the previous one
+                record_ends.pop()
+                path_parts.pop()
+                iso.seek(record_pos.pop() + pvd_off)
+                continue
+
+        # Parse the record
+        inode = iso.tell()
+        dr_len = struct.unpack("<B", iso.read(1))[0]
+        dr_blob = iso.read(dr_len - 1)
+
+        (
+            dr_ea_len,
+            dr_data_pos,
+            dr_data_len,
+            dr_flags,
+            dr_inter,
+            dr_volume,
+            dr_name_len,
+        ) = struct.unpack_from("<BI4xI4x7xBHH2xB", dr_blob)
+
+        assert dr_ea_len == 0, "ISOs with extra attributes are not supported!"
+        assert dr_inter == 0, "Interleaved ISOs are not supported!"
+        assert dr_volume == 1, "multi-volume ISOs are not supported!"
+        assert (dr_flags & 0b1000000) == 0, "4GiB+ files are not supported!"
+
+        # the dir records always en on even addresses
+        if (iso.tell() % 2) != 0:
+            iso.read(1)
+
+        dr_data_pos *= 0x800
+        dr_data_pos += pvd_off
+
+        dr_name = dr_blob[32 : 32 + dr_name_len]
+
+        # record with these names are the '.' and '..'
+        # directories respectively, so skip them
+        if dr_name == b"\x00" or dr_name == b"\x01":
+            continue
+
+        dr_name = dr_name.decode()
+        if dr_name.endswith(";1"):
+            dr_name = dr_name[:-2]
+        path_parts.append(dr_name)
+
+        file_info.total_inodes += 1
+
+        # is it a directory?
+        if (dr_flags & 0b10) != 0:
+            # Go to its directory record
+            record_pos.append(iso.tell())
+            record_ends.append(dr_data_pos + dr_data_len)
+            iso.seek(dr_data_pos)
+            continue
+        else:
+            # Otherwise dump the file
+            fp = "/".join(path_parts)
+
+            file_info.files.append(
+                FileListData(Path(fp), inode, dr_data_pos, dr_data_len)
+            )
+            path_parts.pop()
+
+    return file_info
+
+
+def save_iso_files(
+    iso: typing.BinaryIO, file_info: FileListInfo, base_folder: Path
+) -> None:
+    for file in file_info.files:
+        print(f"SAVING {file.path.as_posix()}")
+
+        final_path = base_folder / file.path
+        final_path.parent.mkdir(exist_ok=True, parents=True)
+        iso.seek(file.lba)
+
+        with open(final_path, "wb+") as f:
+            f.write(iso.read(file.size))
+
+
 def dump_iso(iso_path: Path, filelist: Path, iso_files: Path) -> None:
     if iso_path.exists() is False:
         print(f"Could not to find '{iso_path.name}'!")
         return
 
-    iso_files.mkdir(parents=True, exist_ok=True)
-
     with open(iso_path, "rb") as iso:
         # Sanity check
-        assert check_pvd(iso, 0x8000), "No valid PVD found in Layer0!"
+        assert check_pvd(iso, LAYER0_PVD_LOCATION), "No valid PVD found in Layer0!"
 
         # Test dual-layer-dness
-        iso.seek(0x8050)
-        layer0_sector_count = struct.unpack("<I", iso.read(4))[0]
-        iso.seek(0, io.SEEK_END)
-        iso_sector_count = iso.tell() // SECTOR_SIZE
-        layer1_pvd_pos = layer0_sector_count * SECTOR_SIZE
         has_second_layer = False
 
-        if iso_sector_count != layer0_sector_count:
+        iso.seek(LAYER0_PVD_LOCATION + 0x50)
+        pvd0_sector_count = struct.unpack("<I", iso.read(4))[0]
+        iso.seek(0, io.SEEK_END)
+        iso_sector_count = iso.tell() // SECTOR_SIZE
+        pvd1_pos = pvd0_sector_count * SECTOR_SIZE
+
+        if iso_sector_count != pvd0_sector_count:
             # sector count of the PVD disagree with the file
             # check for another PVD at volume end
-            has_second_layer =  check_pvd(iso, layer1_pvd_pos)
+            has_second_layer = check_pvd(iso, pvd1_pos)
 
             if has_second_layer:
                 print("< Dual layer ISO Detected >")
+                print()
             else:
                 print("WARNING: Iso data suggest this is a double layer image")
-                print("         but no valid PVD was found for Layer1, iso might be corrupt")
+                print(
+                    "         but no valid PVD was found for Layer1, iso might be corrupt"
+                )
+                print()
 
+        layer0_data = dump_dir_records(iso, LAYER0_PVD_LOCATION, 0)
+        layer0_data.files.sort(
+            key=lambda x: x.lba
+        )  # The files are ordered based on their disc position
 
-        path_parts = []
-        record_ends = []
-        record_pos = []
-        file_info = FileListInfo([], 0)
-
-        # get the root directory record off the PVD in Layer1
         if has_second_layer:
-            iso.seek(layer1_pvd_pos + 0x9E)
-            dr_data_pos, dr_data_len = struct.unpack("<I4xI", iso.read(12))
-            dr_data_pos *= SECTOR_SIZE
-            record_ends.append(dr_data_pos + dr_data_len + layer1_pvd_pos - (SECTOR_SIZE * 0x10))
-            record_pos.append(dr_data_pos)
-        
-        # get the root directory record off the PVD in Layer0
-        iso.seek(0x809E)
-        dr_data_pos, dr_data_len = struct.unpack("<I4xI", iso.read(12))
-        dr_data_pos *= SECTOR_SIZE
-        record_ends.append(dr_data_pos + dr_data_len)
-        if has_second_layer:
-            record_pos.append(record_pos[-1])
+            layer1_data = dump_dir_records(iso, pvd1_pos, pvd1_pos - SYSTEM_AREA_SIZE)
+            layer1_data.files.sort(
+                key=lambda x: x.lba
+            )  # The files are ordered based on their disc position
         else:
-            record_pos.append(0)
-        iso.seek(dr_data_pos)
-        print(layer1_pvd_pos)
+            layer1_data = FileListInfo([], 0)
 
-        in_layer1 = False
-        read_offset = 0
+        # save files
+        save_iso_files(iso, layer0_data, iso_files)
 
-        # Traverse directory records recursively
-        # Did I mention that I won't do function recursion?
-        while True:
-            # Have we reached the end of current dir record?
-            if iso.tell() >= record_ends[-1]:
-                if len(record_ends) == 2 and has_second_layer:
-                    print("")
-                    print("< Dumping Second Layer >")
-                    print("")
-                    path_parts.append("")
-                    read_offset = layer1_pvd_pos - (SECTOR_SIZE * 0x10)
-                    in_layer1 = True
-                
-                if len(record_ends) == 1:
-                    # If it's the last one, we finished
-                    break
-                else:
-                    # Otherwise keep reading the previous one
-                    record_ends.pop()
-                    path_parts.pop()
-                    iso.seek(record_pos.pop() + read_offset)
-                    continue
+        if has_second_layer:
+            print("\n< SECOND LAYER >\n")
 
-            # Parse the record
-            inode = iso.tell()
-            print(inode)
-            dr_len = struct.unpack("<B", iso.read(1))[0]
-            dr_blob = iso.read(dr_len - 1)
-
-            (
-                dr_ea_len,
-                dr_data_pos,
-                dr_data_len,
-                dr_flags,
-                dr_inter,
-                dr_volume,
-                dr_name_len,
-            ) = struct.unpack_from("<BI4xI4x7xBHH2xB", dr_blob)
-
-            assert dr_ea_len == 0, "ISOs with extra attributes are not supported!"
-            assert dr_inter == 0, "Interleaved ISOs are not supported!"
-            assert dr_volume == 1, "multi-volume ISOs are not supported!"
-            assert (dr_flags & 0b1000000) == 0, "4GiB+ files are not supported!"
-
-            # the dir records always en on even addresses
-            if (iso.tell() % 2) != 0:
-                iso.read(1)
-
-            dr_data_pos *= 0x800
-
-            dr_name = dr_blob[32 : 32 + dr_name_len]
-
-            # record with these names are the '.' and '..'
-            # directories respectively, so skip them
-            if dr_name == b"\x00" or dr_name == b"\x01":
-                continue
-
-            dr_name = dr_name.decode()
-            if dr_name.endswith(";1"):
-                dr_name = dr_name[:-2]
-            path_parts.append(dr_name)
-
-            file_info.total_inodes += 1
-
-            # is it a directory?
-            if (dr_flags & 0b10) != 0:
-                # Go to its directory record
-                record_pos.append(iso.tell())
-                record_ends.append(dr_data_pos + dr_data_len)
-
-                # if has_second_layer:
-                #     if in_layer1:
-                #         fp = iso_files / "Layer1"
-                #     else:
-                #         fp = iso_files / "Layer0"
-                # else:
-                #     fp = iso_files
-                
-                fp = iso_files / "/".join(path_parts)
-                fp.mkdir(exist_ok=True, parents=True)
-                iso.seek(dr_data_pos + read_offset)
-                continue
-            else:
-                # Otherwise dump the file
-                fp = "/".join(path_parts)
-                print(f"saving {fp}")
-
-                save_pos = iso.tell()
-                with open(iso_files / fp, "wb+") as f:
-                    iso.seek(dr_data_pos + read_offset)
-                    f.write(iso.read(dr_data_len))
-                iso.seek(save_pos)
-
-                file_info.files.append(FileListData(Path(fp), inode, dr_data_pos + read_offset))
-                path_parts.pop()
-
-        # The filelist file has the files ordered based on their disc position
-        file_info.files = sorted(file_info.files, key=lambda x: x.lba)
+        save_iso_files(iso, layer1_data, iso_files)
 
+        # Save filelist
         with open(filelist, "w", encoding="utf8") as f:
-            for d in file_info.files:
-                f.write(f"|{d.inode}||{iso_files.name}/{d.path}|\n")
-            f.write(f"//{file_info.total_inodes}")
+            if has_second_layer:
+                f.write(f"//{len(layer0_data.files)}\n")
+
+            for d in layer0_data.files:
+                f.write(f"|{d.inode}||{iso_files.name}/{d.path.as_posix()}|\n")
+            f.write(f"//{layer0_data.total_inodes}")
+
+            if has_second_layer:
+                f.write("\n")
+                for d in layer1_data.files:
+                    f.write(f"|{d.inode}||{iso_files.name}/{d.path.as_posix()}|\n")
+                f.write(f"//{layer1_data.total_inodes}")
 
 
 def rebuild_iso(