Update afstool.py

empty blocks handling & fd_last_attribute_type now handled by pack
2025-11-04 01:40:30 +01:00 · 2022-02-07 14:20:10 +01:00
parent a322ed9d54
commit 76f7486172
1 changed files with 70 additions and 56 deletions
--- a/afstool.py
+++ b/afstool.py
@@ -9,7 +9,7 @@ import re
 import time


-__version__ = "0.1.0"
+__version__ = "0.1.1"
 __author__ = "rigodron, algoflash, GGLinnk"
 __license__ = "MIT"
 __status__ = "developpement"
@@ -20,8 +20,6 @@ class AfsInvalidFileLenError(Exception): pass
 class AfsEmptyAfsError(Exception): pass
 class AfsInvalidFilenameDirectoryLengthError(Exception): pass
 class AfsInvalidAfsFolderError(Exception): pass
-
-
 # Tested:
 class AfsInvalidMagicNumberError(Exception): pass
 class AfsInvalidFilesRebuildStrategy(Exception): pass
@@ -38,6 +36,8 @@ class AfsFdOffsetOffsetValueError(Exception): pass
 class AfsFdOffsetValueError(Exception): pass
 class AfsFdLastAttributeTypeValueError(Exception): pass
 class AfsFdOffsetCollisionError(Exception): pass
+class AfsEmptyBlockValueError(Exception): pass
+class AfsEmptyBlockAlignError(Exception): pass


 class FilenameResolver:
@@ -113,7 +113,7 @@ class Afs:
        return self.__filenamedirectory[fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN:fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+32].split(b"\x00")[0].decode("utf-8")
    def __get_file_fdlast(self, fileindex:int):
        return int.from_bytes(self.__filenamedirectory[fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+44:fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+48], "little")
-    def __get_mtime(self, fileindex:int):
+    def __get_file_mtime(self, fileindex:int):
        mtime_data = self.__filenamedirectory[fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+32:fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+44]
        year   = int.from_bytes(mtime_data[0:2], "little")
        month  = int.from_bytes(mtime_data[2:4], "little")
@@ -122,14 +122,9 @@ class Afs:
        minute = int.from_bytes(mtime_data[8:10], "little")
        second = int.from_bytes(mtime_data[10:12], "little")
        return time.mktime(datetime(year=year, month=month, day=day, hour=hour, minute=minute, second=second).timetuple())
-    def __patch_file_len(self, fileindex:int, file_len:int):
-        # Patch file_len in the FD
-        if self.__filenamedirectory:
-            if self.__get_file_len(fileindex) == self.__filenamedirectory[fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+44:fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+48]:
-                self.__filenamedirectory[fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+44:fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+48] = file_len.to_bytes(4, "little")
-        # Patch file_len in the TOC
+    def __patch_file_len(self, fileindex:int, file_len:int): # Patch file_len in the TOC
        self.__tableofcontent[Afs.HEADER_LEN+fileindex*8+4:Afs.HEADER_LEN+fileindex*8+8] = file_len.to_bytes(4, "little")
-    def __patch_mtime(self, fileindex:int, mtime):
+    def __patch_file_mtime(self, fileindex:int, mtime):
        mtime = datetime.fromtimestamp(mtime)
        self.__filenamedirectory[Afs.FILENAMEDIRECTORY_ENTRY_LEN*fileindex+32:Afs.FILENAMEDIRECTORY_ENTRY_LEN*fileindex+44] = \
            mtime.year.to_bytes(2,"little")+ \
@@ -138,9 +133,21 @@ class Afs:
            mtime.hour.to_bytes(2,"little")+ \
            mtime.minute.to_bytes(2,"little")+\
            mtime.second.to_bytes(2,"little")
+    def __patch_fdlasts(self, fileindex:int, fd_last_attribute_type): # Patch FD last attributes according to the type
+        if type(fd_last_attribute_type) == int: # every entry has the same const value
+            self.__filenamedirectory[fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+44:fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+48] = fd_last_attribute_type.to_bytes(4, "little")
+        elif fd_last_attribute_type == "length": # 
+            self.__filenamedirectory[fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+44:fileindex*Afs.FILENAMEDIRECTORY_ENTRY_LEN+48] = self.__get_file_len(fileindex).to_bytes(4, "little")
+        elif fd_last_attribute_type == "offset-length":
+            # every odd index is changed according to the TOC lengths values with the serie: 0->updated_index=1 1->updated_index=3 2->updated_index=5
+            # updated_index = index*2+1 with index*2+1 < self.__file_count
+            updated_fdlast_index = fileindex*2+1
+            if updated_fdlast_index < self.__file_count:
+                self.__filenamedirectory[updated_fdlast_index*Afs.FILENAMEDIRECTORY_ENTRY_LEN+44:updated_fdlast_index*Afs.FILENAMEDIRECTORY_ENTRY_LEN+48] = self.__get_file_len(fileindex).to_bytes(4, "little")
+        # fd_last_attribute_type == unknown
    def __pad(self, data:bytes):
-        if len(data) % self.ALIGN != 0:
-            data += b"\x00" * (self.ALIGN - (len(data) % self.ALIGN))
+        if len(data) % Afs.ALIGN != 0:
+            data += b"\x00" * (Afs.ALIGN - (len(data) % Afs.ALIGN))
        return data
    def __clean_filenamedirectory(self):
        self.__filenamedirectory = None
@@ -194,7 +201,7 @@ class Afs:

            # Test if filename is correct by very basic pattern matching
            pattern = re.compile(b"^(?=.{32}$)[^\x00]+\x00+$")
-            for i in range(0, self.__file_count):
+            for i in range(self.__file_count):
                if not pattern.fullmatch(self.__filenamedirectory[i*Afs.FILENAMEDIRECTORY_ENTRY_LEN:i*Afs.FILENAMEDIRECTORY_ENTRY_LEN+32]):
                    self.__clean_filenamedirectory()
                    return False
@@ -214,7 +221,7 @@ class Afs:
            self.__filenamedirectory_len = self.__get_filenamedirectory_len()
            if self.__filenamedirectory_len != len(self.__filenamedirectory):
                raise AfsInvalidFilenameDirectoryLengthError("Error - Tableofcontent filenamedirectory length does not match real filenamedirectory length.")
-    def __print(self, title:str, lines_tuples, columns:list = list(range(0,7)), infos:str = ""):
+    def __print(self, title:str, lines_tuples, columns:list = list(range(7)), infos:str = ""):
        stats_buffer = "#"*100+f"\n# {title}\n"+"#"*100+f"\n{infos}|"+"-"*99+"\n"
        if 0 in columns: stats_buffer += "| Index    ";
        if 1 in columns: stats_buffer += "| b offset ";
@@ -228,11 +235,12 @@ class Afs:
            stats_buffer += line if type(line) == str else "| "+" | ".join(line)+"\n"
        print(stats_buffer, end='')
    # end offset not included (0,1) -> len=1
+    # return a list of offsets where files and sys files begin
    def __get_offsets_map(self):
        # offsets_map is used to check next used offset when updating files
        # we also check if there is intersect between files
        offsets_map = [(0, len(self.__tableofcontent))]
-        for i in range(0, self.__file_count):
+        for i in range(self.__file_count):
            file_offset = self.__get_file_offset(i)
            offsets_map.append( (file_offset, file_offset + self.__get_file_len(i)) )
        if self.__filenamedirectory:
@@ -252,10 +260,10 @@ class Afs:
    def __get_formated_map(self):
        files_map = [("SYS TOC ", "00000000", f"{len(self.__tableofcontent):08x}", f"{len(self.__tableofcontent):08x}", "SYS TOC"+' '*12, "SYS TOC ", "SYS TOC")]

-        for i in range(0, self.__file_count):
+        for i in range(self.__file_count):
            file_offset = self.__get_file_offset(i)
            file_len    = self.__get_file_len(i)
-            file_date   = datetime.fromtimestamp(self.__get_mtime(i)).strftime("%Y-%m-%d %H:%M:%S") if self.__filenamedirectory else " "*19
+            file_date   = datetime.fromtimestamp(self.__get_file_mtime(i)).strftime("%Y-%m-%d %H:%M:%S") if self.__filenamedirectory else " "*19
            filename    = self.__get_file_name(i) if self.__filenamedirectory else f"{i:08}"
            fdlast      = f"{self.__get_file_fdlast(i):08x}" if self.__filenamedirectory else " "*8
            files_map.append((f"{i:08x}", f"{file_offset:08x}", f"{file_offset + file_len:08x}", f"{file_len:08x}", file_date, fdlast, filename))
@@ -265,13 +273,13 @@ class Afs:
                f"{self.__filenamedirectory_offset + len(self.__filenamedirectory):08x}", \
                f"{len(self.__filenamedirectory):08x}", "SYS FD"+' '*13, "SYS FD  ", "SYS FD"))
        return files_map
-    def __get_fd_last_attribute_type(self):
+    def __get_fdlast_type(self):
        # Try to get the type of FD last attribute
        length_type = True
        offset_length_type = True
        constant_type = self.__get_file_fdlast(0)

-        for i in range(0, self.__file_count):
+        for i in range(self.__file_count):
            fd_last_attribute = self.__get_file_fdlast(i)
            if fd_last_attribute != self.__get_file_len(i):
                length_type = None
@@ -282,7 +290,7 @@ class Afs:
        if length_type: return "length"
        if offset_length_type: return "offset-length"
        if constant_type: return f"0x{constant_type:x}"
-        logging.info("unknown FD last attribute type.")
+        logging.info("Unknown FD last attribute type.")
        return "unknown"
    def __write_rebuild_config(self, sys_path:Path, resolver:FilenameResolver):
        config = ConfigParser(allow_no_value=True) # allow_no_value to allow adding comments
@@ -296,12 +304,12 @@ class Afs:
            config.add_section("FilenameDirectory")
            config.set("FilenameDirectory", "toc_offset_of_fd_offset", f"0x{self.__filenamedirectory_offset_offset:x}")
            config.set("FilenameDirectory", "fd_offset", f"0x{self.__filenamedirectory_offset:x}")
-            config.set("FilenameDirectory", "fd_last_attribute_type", self.__get_fd_last_attribute_type())
+            config.set("FilenameDirectory", "fd_last_attribute_type", self.__get_fdlast_type())
        config.write((sys_path / "afs_rebuild.conf").open("w"))

        rebuild_csv = ""
        # generate and save afs_rebuild.csv
-        for i in range(0, self.__file_count):
+        for i in range(self.__file_count):
            filename = self.__get_file_name(i) if self.__filenamedirectory else f"{i:08}"
            unpacked_filename = resolver.resolve_from_index(i, filename) if self.__filenamedirectory else f"{i:08}"
            rebuild_csv += f"{unpacked_filename}/0x{i:x}/0x{self.__get_file_offset(i):x}/{filename}\n"
@@ -327,7 +335,7 @@ class Afs:
            (sys_path / "tableofcontent.bin").write_bytes(self.__tableofcontent)

            logging.info(f"Extracting {self.__file_count} files.")
-            for i in range(0, self.__file_count):
+            for i in range(self.__file_count):
                file_offset = self.__get_file_offset(i)
                file_len    = self.__get_file_len(i)
                filename    = resolver.resolve_new(i, self.__get_file_name(i)) if self.__filenamedirectory else f"{i:08}"
@@ -337,7 +345,7 @@ class Afs:
                (root_path / filename).write_bytes(afs_file.read(file_len))

                if self.__filenamedirectory:
-                    mtime = self.__get_mtime(i)
+                    mtime = self.__get_file_mtime(i)
                    os.utime(root_path / filename, (mtime, mtime))

            if self.__filenamedirectory:
@@ -353,34 +361,37 @@ class Afs:
        root_path = folder_path / "root"

        self.__loadsys_from_folder(sys_path)
-
        resolver = FilenameResolver(sys_path)
-
        offsets_map = self.__get_offsets_map()
+
+        if self.__filenamedirectory:
+            fd_last_attribute_type = self.__get_fdlast_type()
+            if fd_last_attribute_type[:2] == "0x":
+                fd_last_attribute_type = int(fd_last_attribute_type, 16)
+
        with afs_path.open("wb") as afs_file:
            # We update files
-            for i in range(0, self.__file_count):
+            for i in range(self.__file_count):
                file_offset = self.__get_file_offset(i)
                file_len    = self.__get_file_len(i)
-                filename    = self.__get_file_name(i) if self.__filenamedirectory else f"{i:08}"
-                filename    = resolver.resolve_from_index(i, filename)
+                filename    = resolver.resolve_from_index(i, self.__get_file_name(i) if self.__filenamedirectory else f"{i:08}")

                file_path = root_path / filename
                new_file_len = file_path.stat().st_size
                
                if new_file_len != file_len:
-                    next_offset = None
                    # If no FD, we can raise AFS length without constraint
                    if offsets_map.index(file_offset) + 1 < len(offsets_map):
                        next_offset = offsets_map[offsets_map.index(file_offset)+1]
-                    if next_offset:
                        if file_offset + new_file_len > next_offset:
                            raise AfsInvalidFileLenError(f"File {file_path} as a new file_len (0x{new_file_len:x}) > next file offset (0x{next_offset:x}). "\
                                "This means that we have to rebuild the AFS using -r and changing offset of all next files and this could lead to bugs if the main dol use AFS relative file offsets.")
                    self.__patch_file_len(i, new_file_len)
+                    if self.__filenamedirectory:
+                        self.__patch_fdlasts(i, fd_last_attribute_type)
                # If there is a filenamedirectory we update mtime:
                if self.__filenamedirectory:
-                    self.__patch_mtime(i, round(file_path.stat().st_mtime))
+                    self.__patch_file_mtime(i, round(file_path.stat().st_mtime))
                logging.debug(f"Packing {file_path} 0x{file_offset:x}:0x{file_offset+new_file_len:x} in AFS.")
                afs_file.seek(file_offset)
                afs_file.write(self.__pad(file_path.read_bytes()))
@@ -437,6 +448,7 @@ class Afs:

        csv_files_lists = []
        reserved_indexes = []
+        empty_blocks_list = []

        # We parse the file csv and verify entries retrieving length for files
        if (sys_path / "afs_rebuild.csv").is_file():
@@ -444,7 +456,6 @@ class Afs:
                line_splited = line.split('/')
                if len(line_splited) == 4:
                    unpacked_filename = line_splited[0]
-
                    index = None
                    if files_rebuild_strategy in ["index", "mixed"]:
                        if line_splited[1] != "auto":
@@ -477,7 +488,13 @@ class Afs:

                    files_paths.remove( root_path / unpacked_filename )
                elif len(line_splited) == 2: # empty block
-                    raise Exception(f"Error - Empty blocks not implemented yet in afs_rebuild.csv: \"{line}\"")
+                    if line_splited[0][:2] != "0x" or line_splited[1][:2] != "0x" or len(line_splited[0]) < 3 or len(line_splited[1]) < 3:
+                        raise AfsEmptyBlockValueError(f"Error - Invalid empty block values: \"{line}\"")
+                    offset = int(line_splited[0][2:], 16)
+                    length = int(line_splited[1][2:], 16)
+                    if offset % Afs.ALIGN > 0 or length % Afs.ALIGN > 0:
+                        raise AfsEmptyBlockAlignError(f"Error - Invalid empty block offset or length in afs_rebuild.csv: \"{line}\" - offset and length must be aligned to 0x800.")
+                    empty_blocks_list.append([None, None, offset, None, length])
                else:
                    raise AfsInvalidFieldsCountError(f"Error - Invalid entry fields count in afs_rebuild.csv: \"{line}\"")

@@ -485,28 +502,25 @@ class Afs:
        # available_space_ranges is then used to put files that have an adapted length
        # max_offset is used here to find memory collisions between files and next available space
        available_space_ranges = []
+        tmp_ranges = empty_blocks_list
        if files_rebuild_strategy in ["offset", "mixed"]:
-            # We have to sort offsets before merging to avoid complex algorithm
-            # TOC is already present to begin
-            for file_tuple in sorted(csv_files_lists, key=lambda x: (x[2] is not None, x[2])):
-                offset = file_tuple[2]
-                if offset is None:
-                    continue
-                if offset < max_offset:
-                    raise AfsOffsetCollisionError(f"Error - Offsets collision with offset \"0x{offset:x}\".")
-                elif offset > max_offset:
-                    available_space_ranges.append( [max_offset, offset] )
-                max_offset = int(ceil((offset + file_tuple[4]) / Afs.ALIGN)) * Afs.ALIGN
+            tmp_ranges += csv_files_lists
+
+        # We have to sort offsets before merging to avoid complex algorithm
+        # TOC is already present with max_offset
+        for file_tuple in sorted(tmp_ranges, key=lambda x: (x[2] is not None, x[2])):
+            offset = file_tuple[2]
+            if offset is None:
+                continue
+            if offset < max_offset:
+                raise AfsOffsetCollisionError(f"Error - Offsets collision with offset \"0x{offset:x}\".")
+            elif offset > max_offset:
+                available_space_ranges.append( [max_offset, offset] )
+            max_offset = int(ceil((offset + file_tuple[4]) / Afs.ALIGN)) * Afs.ALIGN

        for file_path in files_paths:
            csv_files_lists.append( [file_path.name, None, None, file_path.name, file_path.stat().st_size] )

-        # Now csv_files_lists contains all files for sys rebuild
-        # "auto"   -> csv_files_lists -> index & offsets is None
-        # "index"  -> csv_files_lists -> offsets is None
-        # "offset" -> csv_files_lists -> index is None  -> available_space_ranges to allocate file address space
-        # "mixed"  -> have to consider offsets & indexes -> available_space_ranges to allocate file address space
-        
        # sort by filename
        csv_files_lists.sort(key=lambda x: x[3])
        current_offset = max_offset
@@ -514,7 +528,7 @@ class Afs:
        # if index==None -> Assign an index not in reserved_indexes
        reserved_indexes.sort()
        next_index = 0
-        for i in range(0, len(csv_files_lists)):
+        for i in range(len(csv_files_lists)):
            if csv_files_lists[i][1] is None and files_rebuild_strategy in ["index", "mixed"] or files_rebuild_strategy in ["auto", "offset"]:
                for j in range(next_index, len(csv_files_lists)):
                    if j not in reserved_indexes:
@@ -525,10 +539,10 @@ class Afs:
        csv_files_lists.sort(key=lambda x: x[1])

        # if offset==None -> Assign an offset in available_space_ranges or at the end of file allocated space
-        for i in range(0, len(csv_files_lists)):
+        for i in range(len(csv_files_lists)):
            if files_rebuild_strategy in ["offset", "mixed"] and csv_files_lists[i][2] is None or files_rebuild_strategy in ["auto", "index"]:
                block_len = int(ceil(csv_files_lists[i][4] / Afs.ALIGN)) * Afs.ALIGN
-                for j in range(0, len(available_space_ranges)):
+                for j in range(len(available_space_ranges)):
                    available_block_len = int(ceil((available_space_ranges[j][1] - available_space_ranges[j][0]) / Afs.ALIGN)) * Afs.ALIGN
                    if block_len <= available_block_len:
                        csv_files_lists[i][2] = available_space_ranges[j][0]
@@ -551,7 +565,7 @@ class Afs:
        # Have to be sorted by index
        # current_offset contains now fd offset if not already set
        resolver = FilenameResolver(sys_path)
-        for i in range(0, len(csv_files_lists)):
+        for i in range(len(csv_files_lists)):
            self.__tableofcontent += csv_files_lists[i][2].to_bytes(4, "little") + csv_files_lists[i][4].to_bytes(4, "little")
            # unpacked_filename, index, offset, filename, file_length
            if self.__filenamedirectory_offset_offset: