psptool: Improve Zen generation parsing

psptool/directory.py

@@ -31,6 +31,22 @@ class Directory(NestedBuffer):
     ENTRY_SIZE = DirectoryEntry.ENTRY_SIZE
     FILE_CLASS = File
 
+    ZEN_GENERATION_IDS = {'Zen 1'  : [b'\x00\x09\xBC', b'\x00\x0A\xBC'],
+                          'Zen 2'  : [b'\x05\x0B\xBC', b'\x01\x0A\xBC'],
+                          'Zen 3'  : [b'\x01\x0C\xBC', b'\x00\x0C\xBC'],
+                          'Zen 4'  : [b'\x04\x0D\xBC', b'\x0B\x0D\xBC'],
+                          'Zen 4/5': [b'\x03\x0D\xBC']
+                         }
+
+    @classmethod
+    def get_possible_zen_generation(cls, zen_generation_id):
+        zen_generation = 'unknown'
+        for possible_zen_generation in cls.ZEN_GENERATION_IDS:
+            if zen_generation_id in cls.ZEN_GENERATION_IDS[possible_zen_generation]:
+                zen_generation = possible_zen_generation
+
+        return zen_generation
+
     class ParseError(Exception):
         pass
 
@@ -43,6 +59,7 @@ def create_directories_if_not_exist(cls, offset, fet, zen_generation=None) -> Li
         # Recursively return or create and return found directories
 
         if offset in fet.psptool.directories_by_offset:
+            fet.psptool.directories_by_offset[offset].update_zen_generation(fet, zen_generation)
             return [fet.psptool.directories_by_offset[offset]]
         else:
             # 1. Create the immediate directory in front of us
@@ -69,6 +86,14 @@ def create_directories_if_not_exist(cls, offset, fet, zen_generation=None) -> Li
             for tertiary_directory_offset in directory.tertiary_directory_offsets:
                 directory_body = fet.rom.get_bytes(tertiary_directory_offset, 32)
                 actual_tertiary_offset = int.from_bytes(directory_body[16:20], 'little')
+                zen_generation_id = directory_body[21:24]
+                zen_generation = cls.get_possible_zen_generation(zen_generation_id)
+                if zen_generation == 'unknown':
+                    fet.psptool.ph.print_warning(f"Unknown {zen_generation_id=}")
+
+                zen_generation_id = hex(int.from_bytes(directory_body[20:24], 'little'))
+                zen_generation += f' (PSP ID {zen_generation_id})'
+
                 # Resolve one more indirection
                 tertiary_directories = cls.create_directories_if_not_exist(actual_tertiary_offset, fet, zen_generation)
                 created_directories += tertiary_directories
@@ -198,6 +223,14 @@ def update_entry_fields(self, file: File, type_, size, offset):
         # 3. Update checksum
         self.update_checksum()
 
+    def update_zen_generation(self, fet, zen_generation):
+        if zen_generation is not None:
+            if zen_generation not in self.zen_generation:
+                self.zen_generation += '\n' + zen_generation
+                for offset in self.secondary_directory_offsets:
+                    dir = fet.psptool.directories_by_offset[offset]
+                    dir.update_zen_generation(fet, zen_generation)
+
 
 class BiosDirectory(Directory):
     DIRECTORY_MAGICS = [b'$BHD', b'$BL2']

psptool/fet.py

@@ -19,12 +19,6 @@
 
 from typing import List
 
-
-ZEN_GENERATION_IDS = {'Zen 1': [b'\x00\x09\xBC', b'\x00\x0A\xBC'],
-                      'Zen 2': [b'\x05\x0B\xBC', b'\x01\x0A\xBC'],
-                      'Zen 3': [b'\x01\x0C\xBC', b'\x00\x0C\xBC']}
-
-
 class EmptyFet(Exception):
     pass
 
@@ -121,13 +115,13 @@ def _parse_combo_dir(self, dir_addr):
 
             # entry_addr += self.blob_offset
             zen_generation_id = combo_dir[i*16+5:i*16+8]
-            zen_generation = 'unknown'
-            for possible_zen_generation in ZEN_GENERATION_IDS:
-                if zen_generation_id in ZEN_GENERATION_IDS[possible_zen_generation]:
-                    zen_generation = possible_zen_generation
+            zen_generation = Directory.get_possible_zen_generation(zen_generation_id)
             if zen_generation == 'unknown':
                 self.psptool.ph.print_warning(f"Unknown {zen_generation_id=}")
 
+            zen_generation_id = hex(combo_dir[i*16+4:i*16+8])
+            zen_generation += f' (PSP ID {zen_generation_id})'
+
             results.append((entry_addr, zen_generation))
 
         return results