update scpk.py

pythonlib/formats/scpk.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+import struct
 from ..formats.FileIO import FileIO
 from ..utils import comptolib
 
@@ -12,35 +13,81 @@ class scpk_file():
     data: bytes
 
 
-class Scpk(FileIO):
-    def __init__(self, path="") -> None:
-        super().__init__(path, "r+b", "<")
-        super().__enter__()
-        if self.read(4) != MAGIC:
-            raise ValueError("Not an SCPK file!")
-
-        self.unk1 = self.read_uint16()
-        self.unk2 = self.read_uint16()
-        file_amount = self.read_uint32()
-        self.read_uint32()  # padding?
+class Scpk():
+    def __init__(self) -> None:
+        self.unk1 = 0
+        self.unk2 = 0
         self.files = []
+        self._rsce = b""
+        self._rsce_pos = 0
 
-        sizes = []
-        for _ in range(file_amount):
-            sizes.append(self.read_uint32())
+    
+    def from_path(self, path="") -> typing.Self:
+        with FileIO(path) as f:
+            if f.read(4) != MAGIC:
+                raise ValueError("Not an SCPK file!")
 
-        for size in sizes:
-            data = self.read(size)
-            is_compressed = comptolib.is_compressed(data)
-            c_type = 0
-            if is_compressed:
-                c_type = data[0]
-                data = comptolib.decompress_data(data)
+            self.unk1 = f.read_uint16()
+            self.unk2 = f.read_uint16()
+            self.file_amount = f.read_uint32()
+            assert f.read_uint32() == 0, "scpk padding is not zero!"  # padding?
+            self.files = []
 
-            if len(data) > 8 and data[0:8] == b"THEIRSCE":
-                self.rsce = data
+            sizes = []
+            for _ in range(self.file_amount):
+                sizes.append(f.read_uint32())
+
+            for i, size in enumerate(sizes):
+                data = f.read(size)
+                is_compressed = comptolib.is_compressed(data)
+                c_type = 0
+                if is_compressed:
+                    c_type = data[0]
+                    data = comptolib.decompress_data(data)
+
+                if len(data) > 8 and data[0:8] == b"THEIRSCE":
+                    self._rsce = data
+                    self._rsce_pos = i
+
+                self.files.append(scpk_file(is_compressed, c_type, data))
+
+        return self
+    
+
+    def to_bytes(self):
+        out = MAGIC
+        out += struct.pack("<H", self.unk1)
+        out += struct.pack("<H", self.unk2)
+        out += struct.pack("<I", len(self.files))
+        out += struct.pack("<I", 0)
+
+        blobs = []
+        for blob in self.files:
+            if blob.is_compressed:
+                blobs.append(comptolib.compress_data(blob.data, version=blob.type))
+            else:
+                blobs.append(blob.data)
+        
+        # add sizes
+        for l in [len(x) for x in blobs]:
+            out += struct.pack("<I", l)
+
+        # add files
+        for blob in blobs:
+            out += blob
+
+        return out
+    
+
+    @property
+    def rsce(self):
+        return self._rsce
+    
+    @rsce.setter
+    def rsce(self, value):
+        self._rsce = value
+        self.files[self._rsce_pos].data = value
 
-            self.files.append(scpk_file(is_compressed, c_type, data))
 
     def __getitem__(self, item):
         return self.files[item]