fixup! fix: make functions of mp4.py and save.py asynchronous

src/api.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 from ssl import SSLError
+from typing import Optional
 
 import httpx
 import regex
@@ -167,7 +168,7 @@ async def get_song_info(song_id: str, token: str, storefront: str, lang: str):
 @retry(retry=retry_if_exception_type((httpx.HTTPError, SSLError, FileNotFoundError)),
        wait=wait_random_exponential(multiplier=1, max=60),
        stop=stop_after_attempt(retry_times), before_sleep=before_sleep_log(logger, logging.WARN))
-async def get_song_lyrics(song_id: str, storefront: str, token: str, dsid: str, account_token: str, lang: str) -> str:
+async def get_song_lyrics(song_id: str, storefront: str, token: str, dsid: str, account_token: str, lang: str) -> Optional[str]:
     async with request_lock:
         req = await client.get(f"https://amp-api.music.apple.com/v1/catalog/{storefront}/songs/{song_id}/lyrics",
                                params={"l": lang},
@@ -175,7 +176,10 @@ async def get_song_lyrics(song_id: str, storefront: str, token: str, dsid: str,
                                         "X-Dsid": dsid},
                                cookies={f"mz_at_ssl-{dsid}": account_token})
         result = SongLyrics.model_validate(req.json())
-        return result.data[0].attributes.ttml
+        if result.data:
+            return result.data[0].attributes.ttml
+        else:
+            return None
 
 
 @alru_cache

src/decrypt.py

@@ -18,7 +18,12 @@ retry_count = {}
 async def decrypt(info: SongInfo, keys: list[str], manifest: Datum, device: Device) -> bytes:
     async with device.decryptLock:
         logger.info(f"Decrypting song: {manifest.attributes.artistName} - {manifest.attributes.name}")
-        reader, writer = await asyncio.open_connection(device.host, device.fridaPort)
+        try:
+            reader, writer = await asyncio.open_connection(device.host, device.fridaPort)
+        except ConnectionRefusedError:
+            logger.warning(f"Failed to connect to device {device.device.serial}, re-injecting")
+            device.restart_inject_frida()
+            raise RetryableDecryptException
         decrypted = bytes()
         last_index = 255
         for sample in info.samples:

src/models/song_lyrics.py

@@ -22,4 +22,4 @@ class Datum(BaseModel):
 
 
 class SongLyrics(BaseModel):
-    data: List[Datum]
+    data: Optional[List[Datum]] = None

src/mp4.py

@@ -61,7 +61,7 @@ async def extract_media(m3u8_url: str, codec: str, song_metadata: SongMetadata,
     return stream.segment_map[0].absolute_uri, keys, selected_codec
 
 
-def extract_song(raw_song: bytes, codec: str) -> SongInfo:
+async def extract_song(raw_song: bytes, codec: str) -> SongInfo:
     tmp_dir = TemporaryDirectory()
     mp4_name = uuid.uuid4().hex
     raw_mp4 = Path(tmp_dir.name) / Path(f"{mp4_name}.mp4")
@@ -118,7 +118,7 @@ def extract_song(raw_song: bytes, codec: str) -> SongInfo:
     return SongInfo(codec=codec, raw=raw_song, samples=samples, nhml=raw_nhml, decoderParams=decoder_params)
 
 
-def encapsulate(song_info: SongInfo, decrypted_media: bytes, atmos_convent: bool) -> bytes:
+async def encapsulate(song_info: SongInfo, decrypted_media: bytes, atmos_convent: bool) -> bytes:
     tmp_dir = TemporaryDirectory()
     name = uuid.uuid4().hex
     media = Path(tmp_dir.name) / Path(name).with_suffix(".media")
@@ -170,7 +170,7 @@ def encapsulate(song_info: SongInfo, decrypted_media: bytes, atmos_convent: bool
     return final_song
 
 
-def write_metadata(song: bytes, metadata: SongMetadata, embed_metadata: list[str], cover_format: str) -> bytes:
+async def write_metadata(song: bytes, metadata: SongMetadata, embed_metadata: list[str], cover_format: str) -> bytes:
     tmp_dir = TemporaryDirectory()
     name = uuid.uuid4().hex
     song_name = Path(tmp_dir.name) / Path(f"{name}.m4a")
@@ -199,7 +199,7 @@ def write_metadata(song: bytes, metadata: SongMetadata, embed_metadata: list[str
 # There are suspected errors in M4A files encapsulated by MP4Box and GPAC,
 # causing some applications to be unable to correctly process Metadata (such as Android.media, Salt Music)
 # Using FFMPEG re-encapsulating solves this problem
-def fix_encapsulate(song: bytes) -> bytes:
+async def fix_encapsulate(song: bytes) -> bytes:
     tmp_dir = TemporaryDirectory()
     name = uuid.uuid4().hex
     song_name = Path(tmp_dir.name) / Path(f"{name}.m4a")
@@ -217,7 +217,7 @@ def fix_encapsulate(song: bytes) -> bytes:
 # FFMPEG will overwrite maxBitrate in DecoderConfigDescriptor
 # Using raw song's esds box to fix it
 # see also https://trac.ffmpeg.org/ticket/4894
-def fix_esds_box(raw_song: bytes, song: bytes) -> bytes:
+async def fix_esds_box(raw_song: bytes, song: bytes) -> bytes:
     tmp_dir = TemporaryDirectory()
     name = uuid.uuid4().hex
     esds_name = Path(tmp_dir.name) / Path(f"{name}.atom")

src/rip.py

@@ -48,7 +48,10 @@ async def rip_song(song: Song, auth_params: GlobalAuthParams, codec: str, config
                                f"Use storefront {auth_params.storefront.upper()} to get lyrics")
             lyrics = await get_song_lyrics(song.id, auth_params.storefront, auth_params.accountAccessToken,
                                            auth_params.dsid, auth_params.accountToken, config.region.language)
-            song_metadata.lyrics = lyrics
+            if lyrics:
+                song_metadata.lyrics = lyrics
+            else:
+                logger.warning(f"Unable to get lyrics of song: {song_metadata.artist} - {song_metadata.title}")
         if config.m3u8Api.enable and codec == Codec.ALAC and not specified_m3u8:
             m3u8_url = await get_m3u8_from_api(config.m3u8Api.endpoint, song.id, config.m3u8Api.enable)
             if m3u8_url:
@@ -76,15 +79,15 @@ async def rip_song(song: Song, auth_params: GlobalAuthParams, codec: str, config
         logger.info(f"Downloading song: {song_metadata.artist} - {song_metadata.title}")
         codec = get_codec_from_codec_id(codec_id)
         raw_song = await download_song(song_uri)
-        song_info = extract_song(raw_song, codec)
+        song_info = await extract_song(raw_song, codec)
         decrypted_song = await decrypt(song_info, keys, song_data, device)
-        song = encapsulate(song_info, decrypted_song, config.download.atmosConventToM4a)
+        song = await encapsulate(song_info, decrypted_song, config.download.atmosConventToM4a)
         if not if_raw_atmos(codec, config.download.atmosConventToM4a):
-            metadata_song = write_metadata(song, song_metadata, config.metadata.embedMetadata, config.download.coverFormat)
-            song = fix_encapsulate(metadata_song)
+            metadata_song = await write_metadata(song, song_metadata, config.metadata.embedMetadata, config.download.coverFormat)
+            song = await fix_encapsulate(metadata_song)
             if codec == Codec.AAC or codec == Codec.AAC_DOWNMIX or codec == Codec.AAC_BINAURAL:
-                song = fix_esds_box(song_info.raw, song)
-        filename = save(song, codec, song_metadata, config.download, playlist)
+                song = await fix_esds_box(song_info.raw, song)
+        filename = await save(song, codec, song_metadata, config.download, playlist)
         logger.info(f"Song {song_metadata.artist} - {song_metadata.title} saved!")
         if config.download.afterDownloaded:
             command = config.download.afterDownloaded.format(filename=filename)

src/save.py

@@ -7,7 +7,7 @@ from src.models import PlaylistInfo
 from src.utils import ttml_convent_to_lrc, get_song_name_and_dir_path, get_suffix
 
 
-def save(song: bytes, codec: str, metadata: SongMetadata, config: Download, playlist: PlaylistInfo = None):
+async def save(song: bytes, codec: str, metadata: SongMetadata, config: Download, playlist: PlaylistInfo = None):
     song_name, dir_path = get_song_name_and_dir_path(codec, config, metadata, playlist)
     if not dir_path.exists() or not dir_path.is_dir():
         os.makedirs(dir_path.absolute())