Use use_prefetch_key flag; update song staging

gamdl/downloader/amdecrypt.py

@@ -1584,7 +1584,7 @@ async def _decrypt_track_hex(
     input_path: str,
     decryption_key: str,
     handler_type: bytes,
-    legacy: bool = False,
+    use_prefetch_key: bool = False,
     use_track_key_for_all_descriptions: bool = False,
     file_backed: bool = False,
 ) -> DecryptedTrack:
@@ -1594,7 +1594,7 @@ async def _decrypt_track_hex(
 
     if use_track_key_for_all_descriptions:
         keys = {sample.desc_index: track_key for sample in track_info.samples}
-    elif handler_type == b"soun" and legacy:
+    elif handler_type == b"soun" and use_prefetch_key:
         keys = {0: track_key}
     elif handler_type == b"soun":
         keys = {0: DEFAULT_SONG_DECRYPTION_KEY, 1: track_key}
@@ -1602,10 +1602,10 @@ async def _decrypt_track_hex(
         keys = {sample.desc_index: track_key for sample in track_info.samples}
 
     enc_info = track_info.encryption_info or EncryptionInfo(
-        scheme_type="cenc" if legacy else "cbcs"
+        scheme_type="cenc" if use_prefetch_key else "cbcs"
     )
     enc_info_per_desc = None
-    if track_info.moov_data and not legacy:
+    if track_info.moov_data and not use_prefetch_key:
         enc_info_per_desc = await asyncio.to_thread(
             _extract_encryption_info_per_stsd,
             track_info.moov_data,
@@ -1659,14 +1659,14 @@ async def decrypt_file_hex(
     input_audio_path: str,
     decryption_key_video: str | None = None,
     input_video_path: str | None = None,
-    legacy: bool = False,
+    use_prefetch_key: bool = False,
 ) -> DecryptedMedia:
     """Decrypt audio and optional video with raw AES hex keys."""
     audio = await _decrypt_track_hex(
         input_audio_path,
         decryption_key_audio,
         b"soun",
-        legacy,
+        use_prefetch_key,
         use_track_key_for_all_descriptions=input_video_path is not None,
         file_backed=input_video_path is not None,
     )

gamdl/downloader/song.py

@@ -4,7 +4,7 @@ import structlog
 
 from ..interface.enums import CoverFormat
 from ..interface.types import AppleMusicMedia, DecryptionKeyAv
-from .amdecrypt import decrypt_file_hex, decrypt_wrapper, write_decrypted_media
+from .amdecrypt import decrypt_file_hex, decrypt_wrapper, write_decrypted_media
 from .base import AppleMusicBaseDownloader
 from .types import DownloadItem
 
@@ -55,39 +55,39 @@ class AppleMusicSongDownloader:
         self,
         input_path: str,
         output_path: str,
-        media_id: str,
-        fairplay_key: str,
-    ) -> None:
-        decrypted_media = await decrypt_wrapper(
-            self.base.interface.base.wrapper_url + "/decrypt",
-            media_id,
-            input_path,
-            fairplay_key_audio=fairplay_key,
-        )
-        await write_decrypted_media(decrypted_media, output_path)
+        media_id: str,
+        fairplay_key: str,
+    ) -> None:
+        decrypted_media = await decrypt_wrapper(
+            self.base.interface.base.wrapper_url + "/decrypt",
+            media_id,
+            input_path,
+            fairplay_key_audio=fairplay_key,
+        )
+        await write_decrypted_media(decrypted_media, output_path)
 
     async def _decrypt_amdecrypt_hex(
         self,
         input_path: str,
         output_path: str,
-        decryption_key: str,
-        legacy: bool = False,
-    ) -> None:
-        decrypted_media = await decrypt_file_hex(
-            decryption_key,
-            input_path,
-            legacy=legacy,
-        )
-        await write_decrypted_media(decrypted_media, output_path)
+        decryption_key: str,
+        use_prefetch_key: bool = False,
+    ) -> None:
+        decrypted_media = await decrypt_file_hex(
+            decryption_key,
+            input_path,
+            use_prefetch_key=use_prefetch_key,
+        )
+        await write_decrypted_media(decrypted_media, output_path)
 
     async def stage(
         self,
         encrypted_path: str,
         staged_path: str,
-        decryption_key: DecryptionKeyAv,
-        legacy: bool,
         media_id: str,
-        fairplay_key: str,
+        web_song_codec: bool,
+        decryption_key: DecryptionKeyAv | None = None,
+        fairplay_key: str = None,
     ):
         log = logger.bind(
             action="stage_song",
@@ -96,20 +96,20 @@ class AppleMusicSongDownloader:
             staged_path=staged_path,
         )
 
-        if self.base.interface.base.use_wrapper and not legacy:
+        if decryption_key:
+            await self._decrypt_amdecrypt_hex(
+                encrypted_path,
+                staged_path,
+                decryption_key.audio_track.key,
+                web_song_codec,
+            )
+        else:
             await self._decrypt_amdecrypt(
                 encrypted_path,
                 staged_path,
                 media_id,
                 fairplay_key,
             )
-        else:
-            await self._decrypt_amdecrypt_hex(
-                encrypted_path,
-                staged_path,
-                decryption_key.audio_track.key,
-                legacy,
-            )
 
         log.debug("success")
 
@@ -161,9 +161,9 @@ class AppleMusicSongDownloader:
         await self.stage(
             encrypted_path,
             download_item.staged_path,
+            download_item.media.media_id,
+            download_item.media.stream_info.audio_track.web_song_codec,
             download_item.media.decryption_key,
-            download_item.media.stream_info.audio_track.legacy,
-            download_item.media.media_metadata["id"],
             download_item.media.stream_info.audio_track.fairplay_key,
         )