PyCriCodecsEx 0.0.1__cp310-cp310-macosx_11_0_arm64.whl → 0.0.3__cp310-cp310-macosx_11_0_arm64.whl

PyCriCodecsEx/cpk.py

@@ -8,7 +8,7 @@ from concurrent.futures import ProcessPoolExecutor, as_completed
 from tempfile import NamedTemporaryFile
 import CriCodecsEx
 
-def worker_do_compression(src : str, dst: str):
+def _worker_do_compression(src : str, dst: str):
     with open(src, "rb") as fsrc, open(dst, "wb") as fdst:
         data = fsrc.read()
         compressed = CriCodecsEx.CriLaylaCompress(data)
@@ -48,6 +48,7 @@ class _TOC():
         self.table = UTF(self.stream.read()).table
 
 class CPK:
+    """Use this class to load CPK file table-of-content, and read files from them on-demand."""
     magic: bytes
     encflag: int
     packet_size: int
@@ -55,7 +56,12 @@ class CPK:
     stream: BinaryIO
     tables: dict
     filename: str
-    def __init__(self, filename) -> None:
+    def __init__(self, filename : str | BinaryIO) -> None:
+        """Loads a CPK archive's table-of-content and ready for file reading.
+
+        Args:
+            filename (str | BinaryIO): The path to the CPK file or a BinaryIO stream containing the CPK data.
+        """
         if type(filename) == str:
             self.filename = filename
             self.stream = FileIO(filename)
@@ -184,7 +190,7 @@ class CPKBuilder:
     ContentSize: int
     EnabledDataSize: int
     EnabledPackedSize: int
-    outfile: str
+    outfile: BinaryIO
     init_toc_len: int # This is a bit of a redundancy, but some CPK's need it.
 
     in_files : list[tuple[str, str, bool]] # (source path, dest filename, compress or not)
@@ -249,14 +255,13 @@ class CPKBuilder:
         
         self.in_files.append((src, dst, compress))
 
-    def _writetofile(self, header) -> None:
-        with open(self.outfile, "wb") as out:
-            out.write(header)
-            for i, ((path, _), (filename, file_size, pack_size)) in enumerate(zip(self.os_files, self.files)):
-                src = open(path, 'rb').read()
-                out.write(src)
-                out.write(bytes(0x800 - pack_size % 0x800))
-                self.progress_cb("Write %s" % os.path.basename(filename), i + 1, len(self.files))
+    def _writetofile(self, header) -> None:        
+        self.outfile.write(header)
+        for i, ((path, _), (filename, file_size, pack_size)) in enumerate(zip(self.os_files, self.files)):
+            src = open(path, 'rb').read()
+            self.outfile.write(src)
+            self.outfile.write(bytes(0x800 - pack_size % 0x800))
+            self.progress_cb("Write %s" % os.path.basename(filename), i + 1, len(self.files))
 
     def _populate_files(self, parallel : bool):
         self.files = []
@@ -271,7 +276,7 @@ class CPKBuilder:
                 futures = []
                 for (src, _, _), (dst, compress) in zip(self.in_files,self.os_files):
                     if compress:
-                        futures.append(exec.submit(worker_do_compression, src, dst))
+                        futures.append(exec.submit(_worker_do_compression, src, dst))
                 for i, fut in as_completed(futures):
                     try:
                         fut.result()
@@ -281,7 +286,7 @@ class CPKBuilder:
         else:
             for i, ((src, _, _), (dst, compress)) in enumerate(zip(self.in_files,self.os_files)):
                     if compress:
-                        worker_do_compression(src, dst)
+                        _worker_do_compression(src, dst)
                         self.progress_cb("Compress %s" % os.path.basename(src), i + 1, len(self.in_files))
         for (src, filename, _) , (dst, _) in zip(self.in_files,self.os_files):
             file_size = os.stat(src).st_size         
@@ -299,12 +304,12 @@ class CPKBuilder:
                 pass
         self.os_files = []
 
-    def save(self, outfile : str, parallel : bool = False):
+    def save(self, outfile : str | BinaryIO, parallel : bool = False):
         """Build and save the bundle into a file
 
 
         Args:
-            outfile (str): The output file path.
+            outfile (str | BinaryIO): The output file path or a writable binary stream.
             parallel (bool, optional): Whether to use parallel processing for file compression (if at all used). Defaults to False.
 
         NOTE: 
@@ -313,6 +318,8 @@ class CPKBuilder:
         """
         assert self.in_files, "cannot save empty bundle"
         self.outfile = outfile
+        if type(outfile) == str:
+            self.outfile = open(outfile, "wb")
         self._populate_files(parallel)
         if self.encrypt:
             encflag = 0
@@ -358,7 +365,9 @@ class CPKBuilder:
             data = self.CPKdata.ljust(len(self.CPKdata) + (0x800 - len(self.CPKdata) % 0x800) - 6, b'\x00') + bytearray(b"(c)CRI") + self.ITOCdata
         self._writetofile(data)
         self._cleanup_files()
-    
+        if type(outfile) == str:
+            self.outfile.close()
+            
     def _generate_GTOC(self) -> bytearray:
         # NOTE: Practically useless
         # I have no idea why are those numbers here.

PyCriCodecsEx/hca.py

@@ -5,7 +5,7 @@ from array import array
 import CriCodecsEx
 
 from PyCriCodecsEx.chunk import *
-
+from PyCriCodecsEx.utf import UTFTypeValues, UTFBuilder
 HcaHeaderStruct = Struct(">4sHH")
 HcaFmtHeaderStruct = Struct(">4sIIHH")
 HcaCompHeaderStruct = Struct(">4sHBBBBBBBBBB")
@@ -17,6 +17,10 @@ HcaCiphHeaderStruct = Struct(">4sH")
 HcaRvaHeaderStruct = Struct(">4sf")
 
 class HCA:
+    """HCA class for decoding and encoding HCA files
+
+    **NOTE:** Direct usage of this class is not recommended, use the `HCACodec` wrapper instead.
+    """
     stream: BinaryIO
     hcastream: BinaryIO
     HcaSig: bytes
@@ -46,7 +50,14 @@ class HCA:
     table: array
     looping: bool
 
-    def __init__(self, stream: BinaryIO, key: int = 0, subkey: int = 0) -> None:
+    def __init__(self, stream: str | BinaryIO, key: int = 0, subkey: int = 0) -> None:
+        """Initializes the HCA encoder/decoder
+
+        Args:
+            stream (str | BinaryIO): Path to the HCA or WAV file, or a BinaryIO stream.
+            key (int, optional): HCA key. Defaults to 0.
+            subkey (int, optional): HCA subkey. Defaults to 0.
+        """
         if type(stream) == str:
             self.stream = FileIO(stream)
             self.hcastream = FileIO(stream)
@@ -66,10 +77,10 @@ class HCA:
         self.hcabytes: bytearray = b''
         self.enc_table: array = b''
         self.table: array = b''
-        self.Pyparse_header()
+        self._Pyparse_header()
     
 
-    def Pyparse_header(self) -> None:
+    def _Pyparse_header(self) -> None:
         self.HcaSig, self.version, self.header_size = HcaHeaderStruct.unpack(
             self.hcastream.read(HcaHeaderStruct.size)
         )
@@ -198,39 +209,35 @@ class HCA:
                     raise ValueError(f"WAV bitdepth of {self.fmtBitCount} is not supported, only 16 bit WAV files are supported.")
                 elif self.fmtSize != 16:
                     raise ValueError(f"WAV file has an FMT chunk of an unsupported size: {self.fmtSize}, the only supported size is 16.")
-                if self.stream.read(4) == b"smpl":
-                    self.stream.seek(-4, 1)
-                    self.looping = True
-                    # Will just be naming the important things here.
-                    smplsig, smplesize, _, _, _, _, _, _, _, self.LoopCount, _, _, _, self.LoopStartSample, self.LoopEndSample, _, _ = WavSmplHeaderStruct.unpack(
-                        self.stream.read(WavSmplHeaderStruct.size)
-                    )
-                    if self.LoopCount != 1:
-                        self.looping = False # Unsupported multiple looping points, so backtracks, and ignores looping data.
-                        self.stream.seek(-WavSmplHeaderStruct.size, 1)
-                        self.stream.seek(8 + smplesize, 1)
-                else:
-                    self.stream.seek(-4, 1)
-                    self.looping = False
-                if self.stream.read(4) == b"note": # There's no use for this on ADX.
-                    len = self.stream.read(4)
-                    self.stream.seek(len+4) # + 1? + padding maybe?
-                else:
-                    self.stream.seek(-4, 1)
-                if self.stream.read(4) == b"data":
-                    self.stream.seek(-4, 1)
-                    self.dataSig, self.dataSize = WavDataHeaderStruct.unpack(
-                        self.stream.read(WavDataHeaderStruct.size)
-                    )
-                else:
-                    raise ValueError("Invalid or an unsupported wav file.")
+                while (hdr := self.stream.read(4)):
+                    size = int.from_bytes(self.stream.read(4), 'little')
+                    size += (size & 1) # padding
+                    offset = self.stream.tell()
+                    match hdr:
+                        case b"smpl":
+                            self.stream.seek(-4, 1)
+                            self.looping = True
+                            # Will just be naming the important things here.
+                            smplsig, smplesize, _, _, _, _, _, _, _, self.LoopCount, _, _, _, self.LoopStartSample, self.LoopEndSample, _, _ = WavSmplHeaderStruct.unpack(
+                                self.stream.read(WavSmplHeaderStruct.size)
+                            )
+                            if self.LoopCount != 1:
+                                self.looping = False # Unsupported multiple looping points, so backtracks, and ignores looping data.
+                                self.stream.seek(-WavSmplHeaderStruct.size, 1)
+                                self.stream.seek(8 + smplesize, 1)
+                        case b"data":
+                            self.stream.seek(-4, 1)
+                            self.dataSig, self.dataSize = WavDataHeaderStruct.unpack(
+                                self.stream.read(WavDataHeaderStruct.size)
+                            )
+                    self.stream.seek(offset + size, 0)
         else:
             raise ValueError("Invalid HCA or WAV file.")
         self.stream.seek(0)
         self.hcastream.seek(0)
     
     def info(self) -> dict:
-        """ Returns info related to the input file. """
+        """Returns info related to the input file. """
         if self.filetype == "hca":
             return self.hca
         elif self.filetype == "wav":
@@ -238,6 +245,7 @@ class HCA:
             return wav
     
     def decode(self) -> bytes:
+        """Decodes the HCA or WAV file to WAV bytes. """
         if self.filetype == "wav":
             raise ValueError("Input type for decoding must be an HCA file.")
         self.hcastream.seek(0)
@@ -247,6 +255,7 @@ class HCA:
         return bytes(self.wavbytes)
     
     def encode(self, force_not_looping: bool = False, encrypt: bool = False, keyless: bool = False, quality_level: CriHcaQuality = CriHcaQuality.High) -> bytes:
+        """Encodes the WAV file to HCA bytes."""
         if self.filetype == "hca":
             raise ValueError("Input type for encoding must be a WAV file.")
         if force_not_looping == False:
@@ -260,21 +269,21 @@ class HCA:
         self.stream.seek(0)
         self.hcabytes = CriCodecsEx.HcaEncode(self.stream.read(), force_not_looping, quality_level.value)
         self.hcastream = BytesIO(self.hcabytes)
-        self.Pyparse_header()
+        self._Pyparse_header()
         if encrypt:
             if self.key == 0 and not keyless:
                 self.key = 0xCF222F1FE0748978 # Default key.
-            self.encrypt(self.key, keyless)
+            self._encrypt(self.key, keyless)
         return self.get_hca()
  
-    def encrypt(self, keycode: int, subkey: int = 0, keyless: bool = False) -> None:
+    def _encrypt(self, keycode: int, subkey: int = 0, keyless: bool = False) -> None:
         if(self.encrypted):
             raise ValueError("HCA is already encrypted.")
         self.encrypted = True
         enc = CriCodecsEx.HcaCrypt(self.get_hca(), 1, self.header_size, (1 if keyless else 56), keycode, subkey)
         self.hcastream = BytesIO(enc)
 
-    def decrypt(self, keycode: int, subkey: int = 0) -> None:
+    def _decrypt(self, keycode: int, subkey: int = 0) -> None:
         if(not self.encrypted):
             raise ValueError("HCA is already decrypted.")
         self.encrypted = False
@@ -282,21 +291,161 @@ class HCA:
         self.hcastream = BytesIO(dec)
 
     def get_hca(self) -> bytes:
-        """ Use this function to get the HCA file bytes after encrypting or decrypting. """
+        """Get the HCA file bytes after encrypting or decrypting. """
         self.hcastream.seek(0)
         fl: bytes = self.hcastream.read()
         self.hcastream.seek(0)
         return fl
     
     def get_frames(self):
-        """ Generator function to yield Frame number, and Frame data. """
+        """Generator function to yield Frame number, and Frame data. """
         self.hcastream.seek(self.header_size, 0)
         for i in range(self.hca['FrameCount']):
             yield (i, self.hcastream.read(self.hca['FrameSize']))
     
     def get_header(self) -> bytes:
-        """ Use this function to retrieve the HCA Header. """
+        """Get the HCA Header. """
         self.hcastream.seek(0)
         header = self.hcastream.read(self.header_size)
         self.hcastream.seek(0)
         return header
+
+
+class HCACodec(HCA):
+    """Use this class for encoding and decoding HCA files, from and to WAV."""
+    CHUNK_INTERVAL = 64
+    BASE_FRAMERATE = 2997 # dt = CHUNK_INTERVAL / BASE_FRAMERATE
+    AUDIO_CODEC = 4
+    METADATA_COUNT = 1
+
+    filename: str
+
+    chnls: int
+    sampling_rate: int
+    total_samples: int
+    avbps: int
+
+    filesize: int
+
+    def __init__(self, stream: str | bytes, filename: str = "default.hca", quality: CriHcaQuality = CriHcaQuality.High, key=0, subkey=0, **kwargs):
+        """Initializes the HCA encoder/decoder
+
+        Args:
+            stream (str | bytes): Path to the HCA or WAV file, or a BinaryIO stream. WAV files will be automatically encoded with the given settings first.
+            filename (str, optional): USM filename. Defaults to "default.hca".
+            quality (CriHcaQuality, optional): Encoding quality. Defaults to CriHcaQuality.High.
+            key (int, optional): HCA key. Defaults to 0.
+            subkey (int, optional): HCA subkey. Defaults to 0.
+        """
+        self.filename = filename
+        super().__init__(stream, key, subkey)
+        if self.filetype == "wav":
+            self.encode(
+                force_not_looping=True,
+                encrypt=key != 0,
+                keyless=False,
+                quality_level=quality
+            )
+        self.hcastream.seek(0, 2)
+        self.filesize = self.hcastream.tell()
+        self.hcastream.seek(0)
+
+        if self.filetype == "wav":
+            self.chnls = self.fmtChannelCount
+            self.sampling_rate = self.fmtSamplingRate
+            self.total_samples = int(self.dataSize // self.fmtSamplingSize)
+        else:
+            self.chnls = self.hca["ChannelCount"]
+            self.sampling_rate = self.hca["SampleRate"]
+            self.total_samples = self.hca["FrameCount"]
+        # I don't know how this is derived so I am putting my best guess here. TODO
+        self.avbps = int(self.filesize / self.chnls)
+
+    def generate_SFA(self, index: int, builder):
+        # USMBuilder usage
+        current_interval = 0
+        padding = (
+            0x20 - (self.hca["HeaderSize"] % 0x20)
+            if self.hca["HeaderSize"] % 0x20 != 0
+            else 0
+        )
+        SFA_chunk = USMChunkHeader.pack(
+            USMChunckHeaderType.SFA.value,
+            self.hca["HeaderSize"] + 0x18 + padding,
+            0,
+            0x18,
+            padding,
+            index,
+            0,
+            0,
+            0,
+            current_interval,
+            self.BASE_FRAMERATE,
+            0,
+            0,
+        )
+        SFA_chunk += self.get_header().ljust(self.hca["HeaderSize"] + padding, b"\x00")
+        res = []
+        res.append(SFA_chunk)
+        for i, frame in enumerate(self.get_frames(), start=1):
+            padding = (
+                0x20 - (self.hca["FrameSize"] % 0x20)
+                if self.hca["FrameSize"] % 0x20 != 0
+                else 0
+            )
+            SFA_chunk = USMChunkHeader.pack(
+                USMChunckHeaderType.SFA.value,
+                self.hca["FrameSize"] + 0x18 + padding,
+                0,
+                0x18,
+                padding,
+                index,
+                0,
+                0,
+                0,
+                current_interval,
+                self.BASE_FRAMERATE,
+                0,
+                0,
+            )
+            SFA_chunk += frame[1].ljust(self.hca["FrameSize"] + padding, b"\x00")
+            current_interval = round(i * self.CHUNK_INTERVAL)            
+            res.append(SFA_chunk)
+        else:
+            SFA_chunk = USMChunkHeader.pack(
+                USMChunckHeaderType.SFA.value,
+                0x38,
+                0,
+                0x18,
+                0,
+                index,
+                0,
+                0,
+                2,
+                0,
+                30,
+                0,
+                0,
+            )
+            SFA_chunk += b"#CONTENTS END   ===============\x00"
+            res[-1] += SFA_chunk
+
+        return res
+
+    def get_metadata(self):
+        payload = [dict(hca_header=(UTFTypeValues.bytes, self.get_header()))]
+        p = UTFBuilder(payload, table_name="AUDIO_HEADER")
+        p.strings = b"<NULL>\x00" + p.strings
+        return p.bytes()
+
+    def get_encoded(self):
+        """Gets the encoded HCA audio data."""
+        self.hcastream.seek(0)
+        res = self.hcastream.read()
+        self.hcastream.seek(0)
+        return res
+    
+    def save(self, filepath: str):
+        """Saves the decoded WAV audio to filepath"""
+        with open(filepath, "wb") as f:
+            f.write(self.decode())