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

PyCriCodecsEx/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.1"
+__version__ = "0.0.3"

PyCriCodecsEx/acb.py

@@ -1,18 +1,17 @@
-from struct import iter_unpack
-from typing import BinaryIO, List
-from io import BytesIO
-from PyCriCodecsEx.chunk import *
-from PyCriCodecsEx.utf import UTF, UTFBuilder, UTFViewer
-from PyCriCodecsEx.awb import AWB, AWBBuilder
-from PyCriCodecsEx.hca import HCA
-from copy import deepcopy
-import os
-
 # Credit:
 # - github.com/vgmstream/vgmstream which is why this is possible at all
 # - Original work by https://github.com/Youjose/PyCriCodecs
 # See Research/ACBSchema.py for more details.
 
+from typing import Generator, List, Tuple, BinaryIO
+from PyCriCodecsEx.chunk import *
+from PyCriCodecsEx.utf import UTF, UTFBuilder, UTFViewer
+from PyCriCodecsEx.hca import HCACodec
+from PyCriCodecsEx.adx import ADXCodec
+from PyCriCodecsEx.awb import AWB, AWBBuilder
+from dataclasses import dataclass
+from copy import deepcopy
+
 class CueNameTable(UTFViewer):
     CueIndex: int
     CueName: str
@@ -20,11 +19,13 @@ class CueNameTable(UTFViewer):
 
 class CueTable(UTFViewer):
     CueId: int
+    Length: int
     ReferenceIndex: int
     ReferenceType: int
 
 
 class SequenceTable(UTFViewer):
+    NumTracks : int
     TrackIndex: bytes
     Type: int
 
@@ -65,29 +66,204 @@ class ACBTable(UTFViewer):
     TrackTable: List[TrackTable]
     WaveformTable: List[WaveformTable]
 
+    @staticmethod
+    def decode_tlv(data : bytes):
+        pos = 0
+        while pos < len(data):
+            tag = data[pos : pos + 2]
+            length = data[pos + 3]            
+            value = data[pos + 4 : pos + 4 + length]
+            pos += 3 + length
+            yield (tag, value)
+
+    def waveform_of_track(self, index: int):
+        tlv = self.decode_tlv(self.TrackEventTable[index])
+        def noteOn(data: bytes):
+            # Handle note on event
+            tlv_type, tlv_index = AcbTrackCommandNoteOnStruct.unpack(data[:AcbTrackCommandNoteOnStruct.size])
+            match tlv_type:
+                case 0x02: # Synth
+                    yield from self.waveform_of_synth(tlv_index)
+                case 0x03: # Sequence
+                    yield from self.waveform_of_sequence(tlv_index)
+                # Ignore others silently                
+        for code, data in tlv:
+            match code:
+                case 2000:
+                    yield from noteOn(data)
+                case 2003:
+                    yield from noteOn(data)            
+
+    def waveform_of_sequence(self, index : int):
+        seq = self.SequenceTable[index]
+        for i in range(seq.NumTracks):
+            track_index = int.from_bytes(seq.TrackIndex[i*2:i*2+2], 'big')
+            yield self.WaveformTable[track_index]
+
+    def waveform_of_synth(self, index: int):        
+        item_type, item_index = AcbSynthReferenceStruct.unpack(self.SynthTable[index].ReferenceItems)
+        match item_type:
+            case 0x00: # No audio
+                return
+            case 0x01: # Waveform
+                yield self.WaveformTable[item_index]
+            case 0x02: # Yet another synth...
+                yield from self.waveform_of_synth(item_index)
+            case 0x03: # Sequence
+                yield from self.waveform_of_sequence(item_index)
+            case _:
+                raise NotImplementedError(f"Unknown synth reference type: {item_type} at index {index}")
+
+    def waveform_of(self, index : int) -> List["WaveformTable"]:
+        cue = next(filter(lambda c: c.CueId == index, self.CueTable), None)
+        assert cue, "cue of index %d not found" % index
+        match cue.ReferenceType:
+            case 0x01:
+                return [self.WaveformTable[index]]
+            case 0x02:
+                return list(self.waveform_of_synth(index))
+            case 0x03:
+                return list(self.waveform_of_sequence(index))
+            case 0x08:
+                raise NotImplementedError("BlockSequence type not implemented yet")
+            case _:
+                raise NotImplementedError(f"Unknown cue reference type: {cue.ReferenceType}")
+
+@dataclass(frozen=True)
+class CueItem:
+    CueId: int
+    CueName: str
+    Length: float
+    Waveforms: list[int] # List of waveform IDs
 
 class ACB(UTF):
-    """An ACB is basically a giant @UTF table. Use this class to extract any ACB, and potentially modifiy it in place."""
-    def __init__(self, filename) -> None:
-        super().__init__(filename,recursive=True)
+    """Use this class to read, and modify ACB files in memory."""
+    def __init__(self, stream : str | BinaryIO) -> None:
+        """Loads an ACB file from the given stream.
+
+        Args:
+            stream (str | BinaryIO): The path to the ACB file or a BinaryIO stream containing the ACB data.
+        """
+        super().__init__(stream, recursive=True)
 
     @property
     def payload(self) -> dict:
-        """Retrives the only top-level UTF table dict within the ACB file."""
+        """Retrives the only UTF table dict within the ACB file."""
         return self.dictarray[0]
 
     @property
     def view(self) -> ACBTable:
-        """Returns a view of the ACB file, with all known tables mapped to their respective classes."""
+        """Returns a view of the ACB file, with all known tables mapped to their respective classes.
+        
+        * Use this to interact with the ACB payload instead of `payload` for helper functions, etc"""
         return ACBTable(self.payload)
 
-    # TODO: Extraction routines
-    # See Research/ACBSchema.py. vgmstream presented 4 possible permutations of subsong retrieval.
+    @property
+    def name(self) -> str:
+        """Returns the name of the ACB file."""
+        return self.view.Name
+
+    @property
+    def awb(self) -> AWB:
+        """Returns the AWB object associated with the ACB."""
+        return AWB(self.view.AwbFile)
+
+    def get_waveforms(self) -> List[HCACodec | ADXCodec | Tuple[AcbEncodeTypes, int, int, int,  bytes]]:
+        """Returns a list of decoded waveforms.
+
+        Item may be a codec (if known), or a tuple of (Codec ID, Channel Count, Sample Count, Sample Rate, Raw data).
+        """
+        CODEC_TABLE = {
+            AcbEncodeTypes.ADX: ADXCodec,
+            AcbEncodeTypes.HCA: HCACodec,
+            AcbEncodeTypes.HCAMX: HCACodec,
+        }
+        awb = self.awb
+        wavs = []        
+        for wav in self.view.WaveformTable:
+            encode = AcbEncodeTypes(wav.EncodeType)
+            codec = (CODEC_TABLE.get(encode, None))
+            if codec:
+                wavs.append(codec(awb.get_file_at(wav.MemoryAwbId)))
+            else:
+                wavs.append((encode, wav.NumChannels, wav.NumSamples, wav.SamplingRate, awb.get_file_at(wav.MemoryAwbId)))
+        return wavs
+
+    def set_waveforms(self, value: List[HCACodec | ADXCodec | Tuple[AcbEncodeTypes, int, int, int, bytes]]):
+        """Sets the waveform data.
+
+        Input item may be a codec (if known), or a tuple of (Codec ID, Channel Count, Sample Count, Sample Rate, Raw data).
+
+        NOTE: Cue duration is not set. You need to change that manually.
+        """
+        WAVEFORM = self.view.WaveformTable[0]._payload.copy()
+        encoded = []
+        tables = self.view.WaveformTable
+        tables.clear()
+        for i, codec in enumerate(value):
+            if type(codec) == HCACodec:
+                encoded.append(codec.get_encoded())
+                tables.append(WaveformTable(WAVEFORM.copy()))
+                entry = tables[-1]
+                entry.EncodeType = AcbEncodeTypes.HCA.value
+                entry.NumChannels = codec.chnls
+                entry.NumSamples = codec.total_samples
+                entry.SamplingRate = codec.sampling_rate
+            elif type(codec) == ADXCodec:
+                encoded.append(codec.get_encoded())
+                tables.append(WaveformTable(WAVEFORM.copy()))
+                entry = tables[-1]
+                entry.EncodeType = AcbEncodeTypes.ADX.value
+                entry.NumChannels = codec.chnls
+                entry.NumSamples = codec.total_samples
+                entry.SamplingRate = codec.sampling_rate                
+            elif isinstance(codec, tuple):
+                e_type, e_channels, e_samples, e_rate, e_data = codec
+                encoded.append(e_data)
+                tables.append(WaveformTable(WAVEFORM.copy()))
+                entry = tables[-1]
+                entry.EncodeType = e_type.value
+                entry.NumChannels = e_channels
+                entry.NumSamples = e_samples
+                entry.SamplingRate = e_rate
+            else:
+                raise TypeError(f"Unsupported codec type: {type(codec)}")
+            tables[-1].MemoryAwbId = i
+        awb = self.awb
+        self.view.AwbFile = AWBBuilder(encoded, awb.subkey, awb.version, align=awb.align).build()
+        pass
+
+    @property
+    def cues(self) -> Generator[CueItem, None, None]:
+        """Returns a generator of **read-only** Cues.
+
+        Cues reference waveform bytes by their AWB IDs, which can be accessed via `waveforms`.
+        To modify cues, use the `view` property instead.
+        """
+        for name, cue in zip(self.view.CueNameTable, self.view.CueTable):
+            waveforms = self.view.waveform_of(cue.CueId)
+            yield CueItem(cue.CueId, name.CueName, cue.Length / 1000.0, [waveform.MemoryAwbId for waveform in waveforms])
 
 class ACBBuilder:
+    """Use this class to build ACB files from an existing ACB object."""
     acb: ACB
 
     def __init__(self, acb: ACB) -> None:
+        """Initializes the ACBBuilder with an existing ACB object.
+        
+        Args:
+            acb (ACB): The ACB object to build from.
+
+        Building ACB from scratch isn't planned for now since:
+                
+            * We don't know how SeqCommandTable TLVs work. This is the biggest issue.
+            * Many fields are unknown or not well understood
+            - Games may expect AcfReferenceTable, Asiac stuff etc to be present for their own assets in conjunction
+                with their own ACF table. Missing these is not a fun debugging experience.
+            * ACB tables differ a LOT from game to game (e.g. Lipsync info), contary to USM formats.
+
+        Maybe one day I'll get around to this. But otherwise starting from nothing is a WONTFIX for now.
+        """
         self.acb = acb
 
     def build(self) -> bytes:
@@ -95,6 +271,9 @@ class ACBBuilder:
 
         The object may be modified in place before building, which will be reflected in the output binary.
         """
-        payload = deepcopy(self.acb.dictarray)
-        binary = UTFBuilder(payload, encoding=self.acb.encoding, table_name=self.acb.table_name)
+        # Check whether all AWB indices are valid
+        assert all(
+            waveform.MemoryAwbId < self.acb.awb.numfiles for waveform in self.acb.view.WaveformTable
+        ), "one or more AWB indices are out of range"
+        binary = UTFBuilder(self.acb.dictarray, encoding=self.acb.encoding, table_name=self.acb.table_name)
         return binary.bytes()

PyCriCodecsEx/adx.py

@@ -1,8 +1,13 @@
+from typing import BinaryIO
+from io import BytesIO
+from PyCriCodecsEx.chunk import *
 import CriCodecsEx
-
 class ADX:
-    """ADX Module for decoding and encoding ADX files, pass the either `adx file` or `wav file` in bytes to either `decode` or `encode` respectively."""  
-    
+    """ADX class for decoding and encoding ADX files, pass the either `adx file` or `wav file` in bytes to either `decode` or `encode` respectively.
+
+    **NOTE:** Direct usage of this class is not recommended, use the `ADXCodec` wrapper instead.
+    """
+
     # Decodes ADX to WAV.
     @staticmethod
     def decode(data: bytes) -> bytes:
@@ -13,4 +18,138 @@ class ADX:
     @staticmethod
     def encode(data: bytes, BitDepth = 0x4, Blocksize = 0x12, Encoding = 3, AdxVersion = 0x4, Highpass_Frequency = 0x1F4, Filter = 0, force_not_looping = False) -> bytes:
         """ Encodes WAV to ADX. """
-        return CriCodecsEx.AdxEncode(bytes(data), BitDepth, Blocksize, Encoding, Highpass_Frequency, Filter, AdxVersion, force_not_looping)
+        return CriCodecsEx.AdxEncode(bytes(data), BitDepth, Blocksize, Encoding, Highpass_Frequency, Filter, AdxVersion, force_not_looping)
+    
+class ADXCodec(ADX):
+    """Use this class for encoding and decoding ADX files, from and to WAV."""
+
+    CHUNK_INTERVAL = 99.9
+    BASE_FRAMERATE = 2997
+    # TODO: Move these to an enum
+    AUDIO_CODEC = 2
+    METADATA_COUNT = 0
+
+    filename : str
+    filesize : int
+
+    adx : bytes
+    header : bytes
+    sfaStream: BinaryIO    
+
+    AdxDataOffset: int
+    AdxEncoding: int
+    AdxBlocksize: int
+    AdxSampleBitdepth: int
+    AdxChannelCount: int
+    AdxSamplingRate: int
+    AdxSampleCount: int
+    AdxHighpassFrequency: int
+    AdxVersion: int
+    AdxFlags: int
+
+    chnls: int
+    sampling_rate: int
+    total_samples: int
+    avbps: int
+
+    def __init__(self, stream: str | bytes, filename: str = "default.adx", bitdepth: int = 4, **kwargs):
+        """Initializes the ADX encoder/decoder
+
+        Args:
+            stream (str | bytes): Path to the ADX or WAV file, or a BinaryIO stream. WAV files will be automatically encoded with the given settings first.
+            filename (str, optional): Output filename. Defaults to "default.adx".
+            bitdepth (int, optional): Audio bit depth within [2,15]. Defaults to 4.
+        """
+        if type(stream) == str:
+            self.adx = open(stream, "rb").read()
+        else:
+            self.adx = stream
+        self.filename = filename
+        self.filesize = len(self.adx)
+        magic = self.adx[:4]
+        if magic == b"RIFF":
+            self.adx = self.encode(self.adx, bitdepth, force_not_looping=True)        
+        self.sfaStream = BytesIO(self.adx)
+        header = AdxHeaderStruct.unpack(self.sfaStream.read(AdxHeaderStruct.size))
+        FourCC, self.AdxDataOffset, self.AdxEncoding, self.AdxBlocksize, self.AdxSampleBitdepth, self.AdxChannelCount, self.AdxSamplingRate, self.AdxSampleCount, self.AdxHighpassFrequency, self.AdxVersion, self.AdxFlags = header
+        assert FourCC == 0x8000, "either ADX or WAV is supported"
+        assert self.AdxVersion in {3,4}, "unsupported ADX version"
+        if self.AdxVersion == 4:
+            self.sfaStream.seek(4 + 4  * self.AdxChannelCount, 1)  # Padding + Hist values, they always seem to be 0.
+        self.sfaStream.seek(0)
+        self.chnls = self.AdxChannelCount
+        self.sampling_rate = self.AdxSamplingRate
+        self.total_samples = self.AdxSampleCount
+        self.avbps = int(self.filesize * 8 * self.chnls) - self.filesize
+    
+    def generate_SFA(self, index: int, builder):
+        # USMBuilder usage
+        current_interval = 0
+        stream_size = len(self.adx) - self.AdxBlocksize
+        chunk_size = int(self.AdxSamplingRate // (self.BASE_FRAMERATE / 100) // 32) * (self.AdxBlocksize * self.AdxChannelCount)
+        self.sfaStream.seek(0)
+        res = []
+        while self.sfaStream.tell() < stream_size:
+            if self.sfaStream.tell() > 0:
+                if self.sfaStream.tell() + chunk_size < stream_size:
+                    datalen = chunk_size
+                else:
+                    datalen = (stream_size - (self.AdxDataOffset + 4) - chunk_size) % chunk_size
+            else:
+                datalen = self.AdxDataOffset + 4
+            if not datalen:
+                break
+            padding = (0x20 - (datalen % 0x20) if datalen % 0x20 != 0 else 0)
+            SFA_chunk = USMChunkHeader.pack(
+                    USMChunckHeaderType.SFA.value,
+                    datalen + 0x18 + padding,
+                    0,
+                    0x18,
+                    padding,
+                    index,
+                    0,
+                    0,
+                    0,
+                    round(current_interval),
+                    self.BASE_FRAMERATE,
+                    0,
+                    0
+                    )
+            chunk_data = self.sfaStream.read(datalen)
+            if builder.encrypt_audio:
+                SFA_chunk = builder.AudioMask(chunk_data)
+            SFA_chunk += chunk_data.ljust(datalen + padding, b"\x00")            
+            current_interval += self.CHUNK_INTERVAL
+            res.append(SFA_chunk)
+        # ---
+        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):
+        return None
+
+    def get_encoded(self):
+        """Gets the encoded ADX audio data."""
+        return self.adx
+
+    def save(self, filepath: str):
+        """Saves the encoded ADX audio to filepath"""
+        with open(filepath, "wb") as f:
+            f.write(self.decode(self.adx))
+

PyCriCodecsEx/awb.py

@@ -19,16 +19,21 @@ class AWB:
     headersize: int
     id_alignment: int
 
-    def __init__(self, stream) -> None:
+    def __init__(self, stream : str | BinaryIO) -> None:
+        """Initializes the AWB object
+
+        Args:
+            stream (str | BinaryIO): Source file path or binary stream
+        """
         if type(stream) == str:
             self.stream = FileIO(stream)
             self.filename = stream
         else:
             self.stream = BytesIO(stream)
             self.filename = ""
-        self.readheader()
+        self._readheader()
     
-    def readheader(self):
+    def _readheader(self):
         # Reads header.
         magic, self.version, offset_intsize, self.id_intsize, self.numfiles, self.align, self.subkey = AWBChunkHeader.unpack(
             self.stream.read(AWBChunkHeader.size)
@@ -39,9 +44,9 @@ class AWB:
         # Reads data in the header.
         self.ids = list()
         self.ofs = list()
-        for i in iter_unpack(f"<{self.stringtypes(self.id_intsize)}", self.stream.read(self.id_intsize*self.numfiles)):
+        for i in iter_unpack(f"<{self._stringtypes(self.id_intsize)}", self.stream.read(self.id_intsize*self.numfiles)):
             self.ids.append(i[0])
-        for i in iter_unpack(f"<{self.stringtypes(offset_intsize)}", self.stream.read(offset_intsize*(self.numfiles+1))):
+        for i in iter_unpack(f"<{self._stringtypes(offset_intsize)}", self.stream.read(offset_intsize*(self.numfiles+1))):
             self.ofs.append(i[0] if i[0] % self.align == 0 else (i[0] + (self.align - (i[0] % self.align))))
         
         # Seeks to files offset.
@@ -51,21 +56,20 @@ class AWB:
         self.stream.seek(self.headersize, 0)
 
     def get_files(self):
-        """ Generator function to yield all data blobs from an AWB. """
+        """Generator function to yield all data blobs from an AWB. """
+        self.stream.seek(self.headersize, 0)
         for i in range(1, len(self.ofs)):
             data = self.stream.read((self.ofs[i]-self.ofs[i-1]))
             self.stream.seek(self.ofs[i], 0)
             yield data
     
-    def get_file_at(self, index):
-        """ Gets you a file at specific index. """
-        index += 1
+    def get_file_at(self, index) -> bytes:
+        """Gets you a file at specific index. """
         self.stream.seek(self.ofs[index], 0)
-        data = self.stream.read(self.ofs[index]-self.ofs[index-1])
-        self.stream.seek(self.headersize, 0) # Seeks back to headersize for getfiles.
+        data = self.stream.read(self.ofs[index + 1]-self.ofs[index])
         return data
 
-    def stringtypes(self, intsize: int) -> str:
+    def _stringtypes(self, intsize: int) -> str:
         if intsize == 1:
             return "B" # Probably impossible.
         elif intsize == 2:
@@ -78,7 +82,17 @@ class AWB:
             raise ValueError("Unknown int size.")
 
 class AWBBuilder:
+    """Use this class to build AWB files from a list of bytes."""
     def __init__(self, infiles: list[bytes], subkey: int = 0, version: int = 2, id_intsize = 0x2, align: int = 0x20) -> None:
+        """Initializes the AWB builder.
+
+        Args:
+            infiles (list[bytes]): List of bytes to be included in the AWB file.
+            subkey (int, optional): AWB subkey. Defaults to 0.
+            version (int, optional): AWB version. Defaults to 2.
+            id_intsize (hexadecimal, optional): Integer size (in bytes) for string lengths. Defaults to 0x2.
+            align (int, optional): Alignment. Defaults to 0x20.
+        """
         if version == 1 and subkey != 0:
             raise ValueError("Cannot have a subkey with AWB version of 1.")
         elif id_intsize not in [0x2, 0x4, 0x8]:
@@ -89,7 +103,7 @@ class AWBBuilder:
         self.subkey = subkey
         self.id_intsize = id_intsize
         
-    def stringtypes(self, intsize: int) -> str:
+    def _stringtypes(self, intsize: int) -> str:
         if intsize == 1:
             return "B" # Probably impossible.
         elif intsize == 2:
@@ -102,6 +116,7 @@ class AWBBuilder:
             raise ValueError("Unknown int size.")
 
     def build(self) -> bytes:
+        """Builds the AWB file from the provided infiles bytes."""
         size = 0
         ofs = []
         numfiles = 0
@@ -122,7 +137,7 @@ class AWBBuilder:
             b'AFS2', self.version, intsize, self.id_intsize, numfiles, self.align, self.subkey
         )
 
-        id_strsize = f"<{self.stringtypes(self.id_intsize)}"
+        id_strsize = f"<{self._stringtypes(self.id_intsize)}"
         for i in range(numfiles):
             header += pack(id_strsize, i)
         

PyCriCodecsEx/chunk.py

@@ -12,7 +12,8 @@ WavNoteHeaderStruct = Struct("<4sII")
 WavDataHeaderStruct = Struct("<4sI")
 AdxHeaderStruct = Struct(">HHBBBBIIHBB")
 AdxLoopHeaderStruct = Struct(">HHHHIIII")
-
+AcbSynthReferenceStruct = Struct(">HH")
+AcbTrackCommandNoteOnStruct = Struct(">HH")
 class USMChunckHeaderType(Enum):
     CRID  = b"CRID" # Header.
     SFSH  = b"SFSH" # SofDec1 Header?
@@ -45,11 +46,6 @@ class HCAType(Enum):
     HCA   = b"HCA\x00" # Header.
     EHCA  = b"\xC8\xC3\xC1\x00" # Encrypted HCA header.
 
-class VideoType(Enum):
-    IVF   = b"DKIF" # Header.
-    # H264  = b"" # Header.
-    # MPEG  = b"" # Header.
-
 # I saw some devs swap the unsigned/signed indexes. So I am not sure what's correct or not.
 # In my own experience, swapping those results in an incorrect signed values (should be unsigned) in ACB's/CPK's.
 # If someone were to change this, they must change 'stringtypes' function in UTF/UTFBuilder classes.
@@ -72,4 +68,25 @@ class CriHcaQuality(Enum):
     High    = 1
     Middle  = 2
     Low     = 3
-    Lowest  = 5
+    Lowest  = 5
+
+class AcbEncodeTypes(Enum):
+    ADX = 0
+    PCM = 1
+    HCA = 2
+    ADX_ALT = 3
+    WiiDSP = 4
+    NDSDSP = 5
+    HCAMX = 6
+    VAG = 7
+    ATRAC3 = 8
+    CWAV = 9
+    HEVAG = 10
+    ATRAC9 = 11
+    X360XMA = 12
+    DSP = 13
+    CWACDSP = 14
+    PS4HEVAG = 15
+    PS4ATRAC9 = 16
+    AACM4A = 17
+    SwitchOpus = 18