From 34a94fdcee16f77156c4891a909d9dbe9b0a7bc2 Mon Sep 17 00:00:00 2001
From: Mc-muffin <8714476+Mc-muffin@users.noreply.github.com>
Date: Tue, 13 Feb 2024 09:03:04 -0500
Subject: [PATCH] Make scpk more accurate to the game

---
 pythonlib/formats/scpk.py | 168 +++++++++++++++++++++++++-------------
 1 file changed, 113 insertions(+), 55 deletions(-)

diff --git a/pythonlib/formats/scpk.py b/pythonlib/formats/scpk.py
index 64cc8fa..c7e56ac 100644
--- a/pythonlib/formats/scpk.py
+++ b/pythonlib/formats/scpk.py
@@ -3,103 +3,161 @@ from dataclasses import dataclass
 from pathlib import Path
 
 from ..formats.FileIO import FileIO
+from ..formats.pak import Pak
 from ..utils import comptolib
 
 MAGIC = b"SCPK"
 
+MAP_FLAG = 0x1
+CHR_FLAG = 0x2
+SCE_FLAG = 0x4
+UNK_FLAG = 0x8
+
 
 @dataclass
-class scpk_file():
+class scpk_file:
     is_compressed: bool
     type: int
     data: bytes
 
 
-class Scpk():
+class Scpk:
     def __init__(self) -> None:
-        self.unk1 = 0
-        self.flags = 0
-        self.files = []
-        self._rsce = b""
-        self._rsce_pos = 0
+        self.map: bytes = b""
+        self._map_comp_type = 0
+        self.chars: dict[int, Pak] = dict()
+        self.rsce: bytes = b""
+        self._rsce_comp_type = 0
+        self.unk_file: bytes = b""
+
+        self.version = 1
 
-    
     @staticmethod
-    def from_path(path: Path) -> 'Scpk':
+    def from_path(path: Path) -> "Scpk":
         with FileIO(path) as f:
             if f.read(4) != MAGIC:
                 raise ValueError("Not an SCPK file!")
-            
+
             self = Scpk()
-            self.unk1 = f.read_uint16()
-            self.flags = f.read_uint16()
+            self.version = f.read_uint16()
+            flags = f.read_uint16()
             file_amount = f.read_uint32()
-            assert self.unk1 == 1, "scpk unk1 is not 1!"  # version?
+
+            # It's not checked by the game
+            assert self.version == 1, "scpk version is not 1!"
             assert f.read_uint32() == 0, "scpk padding is not zero!"  # padding?
-            self.files = []
 
             sizes = []
             for _ in range(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)
+            cursor = f.tell()
 
-                if len(data) > 8 and data[0:8] == b"THEIRSCE":
-                    self._rsce = data
-                    self._rsce_pos = i
+            if flags & MAP_FLAG:
+                size = sizes.pop(0)
+                self.map = f.read_at(cursor, size)
+                self._map_comp_type = self.map[0]
+                self.map = comptolib.decompress_data(self.map)
+                cursor += size
 
-                self.files.append(scpk_file(is_compressed, c_type, data))
+            if flags & CHR_FLAG:
+                f.seek(cursor)
+                total_chars = f.read_uint16()
+                char_ids = []
+                for _ in range(total_chars):
+                    char_ids.append(f.read_uint16())
+                cursor += sizes.pop(0)
+
+                for id in char_ids:
+                    size = sizes.pop(0)
+                    self.chars[id] = Pak.from_path(f.read_at(cursor, size), 1)
+                    cursor += size
+
+            if flags & SCE_FLAG:
+                size = sizes.pop(0)
+                self.rsce = f.read_at(cursor, size)
+                self._rsce_comp_type = self.rsce[0]
+                self.rsce = comptolib.decompress_data(self.rsce)
+                cursor += size
+
+            if flags & UNK_FLAG:
+                size = sizes.pop(0)
+                assert size == 4
+                self.unk_file = f.read_at(cursor, 1)
+                cursor += size
 
         return self
-    
 
-    def to_bytes(self):
+    def to_bytes(self) -> bytes:
         out = MAGIC
-        out += struct.pack("<H", self.unk1)
-        out += struct.pack("<H", self.flags)
-        out += struct.pack("<I", len(self.files))
+        out += struct.pack("<H", self.version)
+        out += struct.pack("<H", self.get_flags())
+        out += struct.pack("<I", self.get_total_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)
 
-            temp_len = len(blobs[-1])
-            if (temp_len % 4) != 0:
-                blobs[-1] = blobs[-1] + (b"#" * (0x4 - ((temp_len) % 0x4)))
-        
+        if self.map:
+            blob = comptolib.compress_data(self.map, version=self._map_comp_type)
+            blobs.append(_pad_blob(blob, 4, b"#"))
+
+        if self.chars:
+            blob = struct.pack("<H", len(self.chars))
+
+            for id in self.chars.keys():
+                blob += struct.pack("<H", id)
+
+            blobs.append(_pad_blob(blob, 4, b"#"))
+
+            for chr in self.chars.values():
+                blobs.append(_pad_blob(chr.to_bytes(), 4, b"\x00"))
+
+        if self.rsce:
+            blob = comptolib.compress_data(self.rsce, version=self._rsce_comp_type)
+            blobs.append(_pad_blob(blob, 4, b"#"))
+
+        if self.unk_file:
+            blobs.append(_pad_blob(self.unk_file, 4, b"#"))
+
         # add sizes
-        for l in [len(x) for x in blobs]:
-            out += struct.pack("<I", l)
+        for size in [len(x) for x in blobs]:
+            out += struct.pack("<I", size)
 
         # 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
+    def get_total_files(self) -> int:
+        total_files = 0
+        if self.map:
+            total_files += 1
+        if self.chars:
+            total_files += 1
+            total_files += len(self.chars)
+        if self.rsce:
+            total_files += 1
+        if self.unk_file:
+            total_files += 1
+
+        return total_files
+
+    def get_flags(self) -> int:
+        total_files = 0
+        if self.map:
+            total_files |= MAP_FLAG
+        if self.chars:
+            total_files |= CHR_FLAG
+        if self.rsce:
+            total_files |= SCE_FLAG
+        if self.unk_file:
+            total_files |= UNK_FLAG
+
+        return total_files
 
 
-    def __getitem__(self, item):
-        return self.files[item]
-
-    def __len__(self):
-        return len(self.files)
+def _pad_blob(blob: bytes, pad_to: int, pad_char: bytes) -> bytes:
+    if (len(blob) % pad_to) != 0:
+        blob = blob + (pad_char * (pad_to - ((len(blob)) % pad_to)))
+    return blob