(archives.pi_studios)(#191) bringing Bpk in spec

bsp_tool/archives/pi_studios.py

@@ -4,73 +4,67 @@
 import struct
 from typing import List
 
-from ..utils.binary import read_struct
+from ..branches.base import Struct
+from ..utils import binary
 from . import base
 
 
 class Bpk(base.Archive):
-    filename: str
+    ext = "*.bpk"
     headers: List[CentralHeader]
     files: List[(LocalHeader, bytes)]
 
-    def __init__(self, filename: str):
-        self.filename = filename
-        with open(filename, "rb") as bpk_file:
-            one, num_headers = struct.unpack(">2I", bpk_file.read(8))
-            self.headers = [
-                CentralHeader.from_stream(bpk_file)
-                for i in range(num_headers)]
-            self.files = list()
-            for header in self.headers:
-                assert header.size >= 0x48
-                bpk_file.seek(header.offset)
-                local_header = LocalHeader.from_stream(bpk_file)
-                data = bpk_file.read(header.size - 0x48)[1:-5]
-                # NOTE: trimmed bytes are always 0
-                self.files.append((local_header, data))
+    def __init__(self):
+        self.headers = list()
+        self.files = list()
 
     def __repr__(self) -> str:
-        return f"<BPK '{self.filename}' {len(self.headers)} files @ {id(self):016X}>"
-
-    # TODO: requires filenames (reversing CentralHeader hashes?)
-
-    def extract(self, filename: str, path: str = None) -> str:
-        raise NotImplementedError()
+        descriptor = f"{len(self.headers)} files"
+        return f"<{self.__class__.__name__} {descriptor} @ {id(self):016X}>"
 
     def namelist(self) -> List[str]:
+        # TODO: get filenames somehow
+        # -- reverse hashes in CentralHeader?
         raise NotImplementedError()
 
     def read(self, filename: str) -> bytes:
         raise NotImplementedError()
 
-
-class CentralHeader:
+    @classmethod
+    def from_stream(cls, stream: io.BytesIO) -> Bpk:
+        out = cls()
+        one, num_headers = binary.read_struct(stream, ">2I")
+        out.headers = [
+            CentralHeader.from_stream(stream)
+            for i in range(num_headers)]
+        assert all(ch.one == 1 for ch in out.headers)
+        for header in out.headers:
+            assert header.size >= 0x48
+            stream.seek(header.offset)
+            local_header = LocalHeader.from_stream(stream)
+            data = stream.read(header.size - 0x48)[1:-5]
+            # NOTE: trimmed bytes are always 0
+            out.files.append((local_header, data))
+        return out
+
+
+class CentralHeader(Struct):
     key: int  # filename hash?
     offset: int
-    data_size: int  # matches text length if uncompressed ascii
+    data_size: int  # matches text length if uncompressed plaintext
+    one: int  # always 1
     size: int
-
-    def __init__(self, key, offset, data_size, size):
-        self.key = key
-        self.offset = offset
-        self.data_size = data_size
-        self.size = size
+    __slots__ = ["key", "offset", "data_size", "one", "size"]
+    _format = ">Q4I"
 
     def __repr__(self) -> str:
-        return f"EntryHeader(0x{self.key:016X}, 0x{self.offset:08X}, 0x{self.data_size:06X}, 0x{self.size:06X})"
-
-    def as_bytes(self) -> bytes:
-        return struct.pack(">Q4I", self.key, self.offset, self.data_size, 1, self.size)
-
-    @classmethod
-    def from_bytes(cls, raw: bytes) -> CentralHeader:
-        key, offset, data_size, one, size = struct.unpack(">Q4I", raw)
-        assert one == 1
-        return cls(key, offset, data_size, size)
-
-    @classmethod
-    def from_stream(cls, stream) -> CentralHeader:
-        return cls.from_bytes(stream.read(24))
+        attrs = ", ".join([
+            f"key=0x{self.key:016X}",
+            f"offset=0x{self.offset:08X}",
+            f"data_size=0x{self.data_size:06X}",
+            f"one={self.one}",
+            f"size=0x{self.size:06X}"])
+        return f"EntryHeader({attrs})"
 
 
 class LocalHeader:
@@ -106,10 +100,10 @@ def from_bytes(cls, raw: bytes) -> LocalHeader:
     @classmethod
     def from_stream(cls, stream) -> LocalHeader:
         assert stream.read(8) == b"\x0F\xF5\x12\xEE\x01\x03\x00\x00"
-        assert read_struct(stream, ">5I") == (0, 0, 0x8000, 0x8000, 0)
-        uncompressed_size = read_struct(stream, ">I")
-        assert read_struct(stream, ">I") == 0
-        unknown_1 = read_struct(stream, ">I")
-        assert read_struct(stream, ">I") == 0x8000
-        unknown_2 = read_struct(stream, ">7I")
+        assert binary.read_struct(stream, ">5I") == (0, 0, 0x8000, 0x8000, 0)
+        uncompressed_size = binary.read_struct(stream, ">I")
+        assert binary.read_struct(stream, ">I") == 0
+        unknown_1 = binary.read_struct(stream, ">I")
+        assert binary.read_struct(stream, ">I") == 0x8000
+        unknown_2 = binary.read_struct(stream, ">7I")
         return cls(uncompressed_size, unknown_1, *unknown_2)