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

PyCriCodecsEx/usm.py

@@ -6,8 +6,6 @@ from functools import cached_property
 
 from PyCriCodecsEx.chunk import *
 from PyCriCodecsEx.utf import UTF, UTFBuilder
-from PyCriCodecsEx.adx import ADX
-from PyCriCodecsEx.hca import HCA
 try:
     import ffmpeg
 except ImportError:
@@ -22,6 +20,7 @@ import tempfile
 # code for a complete breakdown of the USM format.
 
 class USMCrypt:
+    """USM related crypto functions"""
     videomask1: bytearray
     videomask2: bytearray
     audiomask: bytearray
@@ -162,6 +161,13 @@ class FFmpegCodec:
     avbps: int
 
     def __init__(self, stream: str | bytes):
+        """Initialize FFmpegCodec with a media stream, gathering metadata and frame info.
+
+        Args:
+            stream (str | bytes): The media stream to process.
+        NOTE:
+            A temp file maybe created for probing only. Which will be deleted after use.
+        """
         if type(stream) == str:
             self.filename = stream
         else:
@@ -317,249 +323,11 @@ class MPEG1Codec(FFmpegCodec):
         super().__init__(stream)
         assert self.format == "mpegvideo", "must be m1v format (mpegvideo)."
 
-class HCACodec(HCA):
-    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, quality: CriHcaQuality = CriHcaQuality.High, key=0, subkey=0, **kwargs):
-        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"):
-        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 save(self, filepath: str):
-        """Saves the decoded WAV audio to filepath"""
-        with open(filepath, "wb") as f:
-            f.write(self.decode())
-
-class ADXCodec(ADX):
-    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, bitdepth: int = 4, **kwargs):
-        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"):
-        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 save(self, filepath: str):
-        """Saves the encoded ADX audio to filepath"""
-        with open(filepath, "wb") as f:
-            f.write(self.decode(self.adx))
-
+from PyCriCodecsEx.hca import HCACodec
+from PyCriCodecsEx.adx import ADXCodec
 
 class USM(USMCrypt):
-    """USM class for extracting infromation and data from a USM file."""
+    """Use this class to extract infromation and data from a USM file."""
 
     filename: str
     decrypt: bool
@@ -574,7 +342,7 @@ class USM(USMCrypt):
 
     metadata: list
 
-    def __init__(self, filename, key: str | int = None):
+    def __init__(self, filename : str | BinaryIO, key: str | int = None):
         """Loads a USM file into memory and prepares it for processing.
 
         Args:
@@ -590,7 +358,10 @@ class USM(USMCrypt):
         self._load_file()
     
     def _load_file(self):
-        self.stream = open(self.filename, "rb")
+        if type(self.filename) == str:
+            self.stream = open(self.filename, "rb")
+        else:
+            self.stream = self.filename
         self.stream.seek(0, 2)
         self.size = self.stream.tell()
         self.stream.seek(0)
@@ -705,7 +476,7 @@ class USM(USMCrypt):
 
     @property
     def streams(self):
-        """[Type (@SFV, @SFA), Filename, Raw stream data]"""
+        """Generator of Tuple[Stream Type ("@SFV", "@SFA"), File name, Raw stream data]"""
         for stream in self.CRIDObj.dictarray[1:]:
             filename, stmid, chno = stream["filename"][1], stream["stmid"][1], stream["chno"][1]
             stmid = int.to_bytes(stmid, 4, 'big', signed='False')
@@ -729,7 +500,7 @@ class USM(USMCrypt):
         stream.filename = sfname
         return stream
 
-    def get_audios(self) -> List[HCACodec]:
+    def get_audios(self) -> List[ADXCodec | HCACodec]:
         """Create a list of audio codecs from the available streams."""
         match self.audio_codec:
             case ADXCodec.AUDIO_CODEC:
@@ -740,115 +511,73 @@ class USM(USMCrypt):
                 return []
 
 class USMBuilder(USMCrypt):
-    """USM class for building USM files."""
+    """Use this class to build USM files."""
     video_stream: VP9Codec | H264Codec | MPEG1Codec
-    
-    enable_audio: bool
     audio_streams: List[HCACodec | ADXCodec]
 
-    key: int
-    encrypt: bool
-    encrypt_audio: bool
+    key: int = None
+    encrypt: bool = False
+    encrypt_audio: bool = False
 
-    audio_codec: int
-    # !!: TODO Quality settings
     def __init__(
         self,
-        video: str,
-        audio: List[str] | str = None,
         key = None,
-        audio_codec=HCACodec.AUDIO_CODEC,
-        encrypt_audio: bool = False,
+        encrypt_audio = False
     ) -> None:
         """Initialize the USMBuilder from set source files.
 
         Args:
-            video (str): The path to the video file. The video source format will be used to map accordingly to the ones Sofdec use.
-                - MPEG1 (with M1V container): MPEG1 Codec (Sofdec Prime)
-                - H264 (with H264 container): H264 Codec
-                - VP9 (with IVF container): VP9 Codec
-            audio (List[str] | str, optional): The path(s) to the audio file(s). Defaults to None.
             key (str | int, optional): The encryption key. Either int64 or a hex string. Defaults to None.
-            audio_codec (int, optional): The audio codec to use. Defaults to HCACodec.AUDIO_CODEC.
-            encrypt_audio (bool, optional): Whether to encrypt the audio. Defaults to False.
+            encrypt_audio (bool, optional): Whether to also encrypt the audio. Defaults to False.
         """
-        self.audio_codec = audio_codec
-        self.encrypt = False
-        self.enable_audio = False
-        self.encrypt_audio = encrypt_audio
-        self.key = 0
-        if encrypt_audio and not key:
-            raise ValueError("Cannot encrypt Audio without key.")
         if key:
             self.init_key(key)
             self.encrypt = True
-        self.load_video(video)
+        self.encrypt_audio = encrypt_audio
         self.audio_streams = []
-        if audio:
-            self.load_audio(audio)
-            self.enable_audio = True
-
-    def load_video(self, video):
-        temp_stream = FFmpegCodec(video)
-        self.video_stream = None
-        match temp_stream.stream["codec_name"]:
-            case "h264":
-                self.video_stream = H264Codec(video)
-            case "vp9":
-                self.video_stream = VP9Codec(video)
-            case "mpeg1video":
-                self.video_stream = MPEG1Codec(video)
-        assert self.video_stream, (
-            "fail to match suitable video codec. Codec=%s"
-            % temp_stream.stream["codec_name"]
-        )
 
-    def load_audio(self, audio):
-        self.audio_filenames = []
-        if type(audio) == list:
-            count = 0
-            for track in audio:
-                if type(track) == str:
-                    self.audio_filenames.append(os.path.basename(track))
-                else:
-                    self.audio_filenames.append("{:02d}.sfa".format(count))
-                    count += 1
-        else:
-            if type(audio) == str:
-                self.audio_filenames.append(os.path.basename(audio))
-            else:
-                self.audio_filenames.append("00.sfa")
+    def add_video(self, video : str | H264Codec | VP9Codec | MPEG1Codec):
+        """Sets the video stream from the specified video file.
 
-        self.audio_streams = []
-        codec = None
-        match self.audio_codec:
-            case HCACodec.AUDIO_CODEC:
-                codec = HCACodec
-            case ADXCodec.AUDIO_CODEC:
-                codec = ADXCodec
-        assert codec, (
-            "fail to match suitable audio codec given option: %s" % self.audio_codec
-        )
-        if type(audio) == list:
-            for track in audio:
-                if type(track) == str:
-                    fn = os.path.basename(track)
-                else:
-                    fn = "{:02d}.sfa".format(count)
-                hcaObj = codec(track, fn, key=self.key)
-                self.audio_streams.append(hcaObj)
+        USMs only support one video stream. Consecutive calls to this method will replace the existing video stream.
+
+        When `video` is str - it will be treated as a file path. The video source format will be used to map accordingly to the ones Sofdec use.
+            - MPEG1 (with M1V container): MPEG1 Codec (Sofdec Prime)
+            - H264 (with H264 container): H264 Codec
+            - VP9 (with IVF container): VP9 Codec
+            
+        Args:
+            video (str | FFmpegCodec): The path to the video file or an FFmpegCodec instance.
+        """
+        if isinstance(video, str):
+            temp_stream = FFmpegCodec(video)
+            self.video_stream = None
+            match temp_stream.stream["codec_name"]:
+                case "h264":
+                    self.video_stream = H264Codec(video)
+                case "vp9":
+                    self.video_stream = VP9Codec(video)
+                case "mpeg1video":
+                    self.video_stream = MPEG1Codec(video)
+            assert self.video_stream, (
+                "fail to match suitable video codec. Codec=%s"
+                % temp_stream.stream["codec_name"]
+            )
         else:
-            if type(audio) == str:
-                fn = os.path.basename(audio)
-            else:
-                fn = "00.sfa"
-            hcaObj = codec(audio, fn, key=self.key)
-            self.audio_streams.append(hcaObj)
+            self.video_stream = video
 
+    def add_audio(self, audio : ADXCodec | HCACodec):
+        """Append the audio stream(s) from the specified audio file(s).
+
+        Args:
+            audio (ADXCodec | HCACodec): The path(s) to the audio file(s).
+        """
+        self.audio_streams.append(audio)
 
     def build(self) -> bytes:
+        """Build the USM payload"""
         SFV_list = self.video_stream.generate_SFV(self)
-        if self.enable_audio:
+        if self.audio_streams:
             SFA_chunks = [s.generate_SFA(i, self) for i, s in enumerate(self.audio_streams) ]
         else:
             SFA_chunks = []
@@ -928,7 +657,7 @@ class USMBuilder(USMCrypt):
         )
         CRIUSF_DIR_STREAM.append(video_dict)
 
-        if self.enable_audio:
+        if self.audio_streams:
             chno = 0
             for stream in self.audio_streams:
                 avbps = stream.avbps
@@ -936,7 +665,7 @@ class USMBuilder(USMCrypt):
                 minbuf += 27860
                 audio_dict = dict(
                     fmtver=(UTFTypeValues.uint, 0),
-                    filename=(UTFTypeValues.string, self.audio_filenames[chno]),
+                    filename=(UTFTypeValues.string, stream.filename),
                     filesize=(UTFTypeValues.uint, stream.filesize),
                     datasize=(UTFTypeValues.uint, 0),
                     stmid=(
@@ -1018,7 +747,7 @@ class USMBuilder(USMCrypt):
 
         audio_metadata = []
         audio_headers = []
-        if self.enable_audio:
+        if self.audio_streams:
             chno = 0
             for stream in self.audio_streams:
                 metadata = stream.get_metadata()
@@ -1141,7 +870,7 @@ class USMBuilder(USMCrypt):
         VIDEO_HDRINFO = gen_video_hdr_info(len(VIDEO_SEEKINFO))
 
         total_len = sum([len(x) for x in SFV_list]) + first_chk_ofs
-        if self.enable_audio:
+        if self.audio_streams:
             sum_len = 0
             for stream in SFA_chunks:
                 for x in stream:
@@ -1181,7 +910,7 @@ class USMBuilder(USMCrypt):
 
         # Header chunks
         header += VIDEO_HDRINFO
-        if self.enable_audio:
+        if self.audio_streams:
             header += b''.join(audio_headers)            
         SFV_END = USMChunkHeader.pack(
             USMChunckHeaderType.SFV.value,
@@ -1202,7 +931,7 @@ class USMBuilder(USMCrypt):
         header += SFV_END
 
         SFA_chk_END  = b'' # Maybe reused
-        if self.enable_audio:
+        if self.audio_streams:
             SFA_chk_END  = b''.join([
                 USMChunkHeader.pack(
                     USMChunckHeaderType.SFA.value,
@@ -1223,7 +952,7 @@ class USMBuilder(USMCrypt):
         header += SFA_chk_END # Ends audio_headers
         header += VIDEO_SEEKINFO
 
-        if self.enable_audio:
+        if self.audio_streams:
             header += b''.join(audio_metadata)
         SFV_END = USMChunkHeader.pack(
             USMChunckHeaderType.SFV.value,