usm: Default HCA codec decryption key should be the USM key itself (#4)

Author: mos9527

.vscode/launch.json

@@ -78,6 +78,13 @@
             "request": "launch",
             "module": "Tests.test_CPK",
             "justMyCode": false
-        },   
+        },  
+        {
+            "name": "Temp Issue 4",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "Tests.issue4",
+            "justMyCode": false
+        } 
     ]
 }

PyCriCodecsEx/usm.py

@@ -20,7 +20,7 @@ class USMCrypt:
     videomask1: bytearray
     videomask2: bytearray
     audiomask: bytearray
-
+    usm_key: int = 0
     def init_key(self, key: str):
         if type(key) == str:
             if len(key) <= 16:
@@ -36,6 +36,7 @@ def init_key(self, key: str):
             raise ValueError(
                 "Invalid key format, must be either a string or an integer."
             )
+        self.usm_key = int.from_bytes(key1, "big") | (int.from_bytes(key2, "big") << 32)
         t = bytearray(0x20)
         t[0x00:0x09] = [
             key1[3],
@@ -506,13 +507,19 @@ def get_video(self) -> VP9Codec | H264Codec | MPEG1Codec:
         stream.filename = sfname
         return stream
 
-    def get_audios(self) -> List[ADXCodec | HCACodec]:
-        """Create a list of audio codecs from the available streams."""
+    def get_audios(self, hca_key = 0, hca_subkey = 0) -> List[ADXCodec | HCACodec]:
+        """Create a list of audio codecs from the available streams.
+        
+        Args:
+            hca_key (int, optional): The HCA decryption key. Either int64 or a hex string. Defaults to 0 - in which
+                                     case the key for USM (if used) would also be used for HCA decryption.
+            hca_subkey (int, optional): The HCA decryption subkey. Either int64 or a hex string. Defaults to 0.
+        """
         match self.audio_codec:
             case ADXCodec.AUDIO_CODEC:
                 return [ADXCodec(s[2], s[1]) for s in self.streams if s[0] == USMChunckHeaderType.SFA.value]
             case HCACodec.AUDIO_CODEC:
-                return [HCACodec(s[2], s[1]) for s in self.streams if s[0] == USMChunckHeaderType.SFA.value] # HCAs are never encrypted in USM
+                return [HCACodec(s[2], s[1], key=hca_key or self.usm_key, subkey=hca_subkey) for s in self.streams if s[0] == USMChunckHeaderType.SFA.value] # HCAs are never encrypted in USM
             case _:
                 return []
 

Tests/issue4.py

@@ -0,0 +1,41 @@
+# USM Sample from Digimon Story: Cyber Sleuth (PC)
+# Uses same key for both HCA and USM decryption
+# HCA Key {2897314143465725881}, // 283553DCE3FD5FB9
+# USM Key 2897314143465725881
+from . import sample_file_path, temp_file_path
+from PyCriCodecsEx.usm import USM, USMBuilder, ADXCodec, HCACodec
+
+def test_usm_decode_and_mux():
+    usm = USM(temp_file_path('S01_B.usm'), 2897314143465725881)
+    audio = usm.get_audios()
+    video = usm.get_video()
+    audio = audio[0] if audio else None
+    # Mux into MP4
+    import ffmpeg, os
+    def mux_av(video_src: str, audio_src: str, output: str, delete: bool = False):
+        (        
+            ffmpeg.output(
+                ffmpeg.input(video_src), 
+                ffmpeg.input(audio_src),
+                output, 
+                vcodec='copy',
+                acodec='copy',
+            ).overwrite_output()
+        ).run()
+        if delete:
+            print('* Cleaning up temporary files')        
+            os.unlink(video_src)
+            os.unlink(audio_src)
+        print(f'* Result available at: {output}')
+    saved_video = temp_file_path('tmp_video.mp4')
+    saved_audio = temp_file_path('tmp_audio.wav')
+    saved_hca = temp_file_path('tmp_audio.hca')
+    result = temp_file_path('muxed_result1.mp4')
+    open(saved_hca,'wb').write(audio.get_hca())
+    video.save(saved_video)
+    audio.save(saved_audio)
+    mux_av(saved_video, saved_audio, result)
+    print('Remux Done.')
+
+if __name__ == "__main__":
+    test_usm_decode_and_mux()