tools: Preserve superblock between unpacking and re-packing

extract_blocks.py

@@ -1,6 +1,7 @@
 import os
 import sys
 from fatx import FATX
+from fatx.blocks import SuperBlock
 
 if __name__ == "__main__":
     cwd = os.getcwd()
@@ -9,9 +10,9 @@
     os.mkdir(os.path.split(sys.argv[1])[1] + ".extract")
     os.chdir(os.path.split(sys.argv[1])[1] + ".extract")
     with open("superblock.bin", "w+b") as f:
-        f.write(fs.read(FATX.SUPERBLOCK_SIZE))
+        f.write(fs.read(SuperBlock.SUPERBLOCK_SIZE))
     with open("fat.bin", "w+b") as f:
-        fat_size = FATX.Filesystem._calc_fat_size(size)
+        fat_size = FATX.Filesystem._calc_fat_size(size, 512*32)
         f.write(fs.read(fat_size))
     for i in range(0, 10):
         with open("chunk({0}).bin".format(i), "w+b") as f:

fatx/FATX.py

@@ -45,11 +45,15 @@ def __init__(self, file: str, sector_size: int = 512):
         self.root = RootObject(DirectoryEntryList(cluster, 1))
 
     @classmethod
-    def new(cls, size: int, file: str, sector_size: int = 512):
+    def new(cls, size: int, file: str, sector_size: int = 512, sb: bytes = None):
         self = cls.__new__(cls)
         self.f = open(file, "w+b")
 
-        self.sb = SuperBlock.new(sector_size)
+        if sb:
+            self.sb = SuperBlock(sb, sector_size)
+            print("Yupp found old superblock")
+        else:
+            self.sb = SuperBlock.new(sector_size)
         self.fat_size = self._calc_fat_size(size, self.sb.cluster_size)
         self.fat = FAT.new(self.fat_size)
         root_dl = DirectoryEntryList(b"\xFF" * 64, 1)

fatx/blocks.py

@@ -24,9 +24,14 @@ def __init__(self, sb, sector_size):
             raise BaseException(
                 "SuperBlock is not " + str(self.SUPERBLOCK_SIZE) + " bytes long"
             )
-        self.name, self.volume, self.cluster_num, self.fatcopies = struct.unpack(
-            "<4sIIh4082x", sb
-        )
+
+        (
+            self.name,
+            self.volume,
+            self.cluster_num,
+            self.fatcopies,
+            self.padding,
+        ) = struct.unpack("<4sIIh4082s", sb)
         try:
             self.name = self.name.decode("ascii")
         except UnicodeDecodeError:
@@ -45,6 +50,7 @@ def new(cls, sector_size):
         self.cluster_num = 32
         self.cluster_size = self.cluster_num * self.SECTOR_SIZE
         self.fatcopies = 1
+        self.padding = 4082*b'\xFF'
         return self
 
     def pack(self):
@@ -54,7 +60,7 @@ def pack(self):
             self.volume,
             self.cluster_num,
             self.fatcopies,
-            4082 * b"\xFF",
+            self.padding,
         )
 
     def __str__(self):
@@ -91,7 +97,7 @@ def __init__(self, raw_clustermap):
         # slice up the fat table
         while len(self.f) > 0:
             entry = int.from_bytes(self.f[: self.size], "little")
-            self.f = self.f[self.size :]
+            self.f = self.f[self.size:]
             self.clustermap.append(entry)
 
         if self.size == 2:
@@ -209,8 +215,10 @@ def linkClusterChain(self, cc):
         self.setEntryType(index, EntryType.FATX_CLUSTER_END)
 
     @classmethod
-    def new(cls, size):
-        self = cls(size * b"\x00")
+    def new(cls, fat_size: int):
+        self = cls.__new__(cls)
+        self.size = 2 if fat_size < (0xFFF5 * 2) else 4
+        self.clustermap = [0x0000] * (fat_size // self.size)
         if self.size == 2:
             self.clustermap[0] = 0xFFF8
             self.clustermap[1] = 0xFFFF

pack.py

@@ -4,17 +4,29 @@
 from fatx import FATX
 
 
+def check_for_superblock(path):
+    # check if a superblock is given, if so, prepare to re-import it
+    path = os.path.join(path, ".FATX-on-a-snake/superblock.bin")
+    if os.path.exists(path):
+        with open(path, 'rb') as f:
+            sb = f.read()
+            return sb
+    else:
+        return None
+
+
 def walkfs(fs):
     with os.scandir(".") as it:
         for entry in it:
             if entry.is_file():
                 with open(entry.name, "rb") as f:
                     fs.import_file(entry.name, f.read())
             else:
-                fs.create_dir(entry.name)
-                os.chdir(entry.name)
-                walkfs(fs.get(entry.name))
-                os.chdir("..")
+                if entry.name != ".FATX-on-a-snake":
+                    fs.create_dir(entry.name)
+                    os.chdir(entry.name)
+                    walkfs(fs.get(entry.name))
+                    os.chdir("..")
 
 
 if __name__ == "__main__":
@@ -48,12 +60,15 @@ def walkfs(fs):
     size = args.size
     src = args.src[0]
     file = args.image[0]
+    sb = None
+    FATX.READ_ONLY = False
+
     if not os.path.isdir(src):
         sys.exit("Fatal: src-dir is not a valid directory")
     if os.path.exists(file):
         sys.exit("Fatal: target file already exists")
 
-    FATX.READ_ONLY = False
-    fs = FATX.Filesystem.new(size, file, args.sector_size)
+    sb = check_for_superblock(src)
+    fs = FATX.Filesystem.new(size, file, args.sector_size, sb)
     os.chdir(src)
     walkfs(fs.root)

unpack.py

@@ -56,3 +56,10 @@ def walkfs(obj: FatxObject):
     root = fs.root
     os.chdir(dest)
     print("Unpacked {0} files.".format(walkfs(root)))
+    os.mkdir(".FATX-on-a-snake")
+    os.chdir(".FATX-on-a-snake")
+    f = open("superblock.bin", "wb")
+    f.write(fs.sb.pack())
+    f.close()
+    print("Saved superblock as {0}/{1}/{2}".format(
+        dest, ".FATX-on-a-snake", "superblock.bin"))