From 10067fe160987f9878018293993b735637cb1866 Mon Sep 17 00:00:00 2001
From: WorldObservationLog <wolc@duck.com>
Date: Tue, 21 May 2024 22:33:41 +0800
Subject: [PATCH] fix: enhance stability

---
 src/decrypt.py    | 27 +++++++++++++++++++--------
 src/exceptions.py |  4 ++++
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/src/decrypt.py b/src/decrypt.py
index bf5536f..04b9aee 100644
--- a/src/decrypt.py
+++ b/src/decrypt.py
@@ -2,16 +2,18 @@ import asyncio
 import logging
 
 from loguru import logger
-from tenacity import retry, retry_if_exception_type, stop_after_attempt, before_sleep_log
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, before_sleep_log, RetryCallState
 
 from src.adb import Device
-from src.exceptions import DecryptException
+from src.exceptions import DecryptException, RetryableDecryptException
 from src.models.song_data import Datum
 from src.mp4 import SongInfo, SampleInfo
 from src.types import defaultId, prefetchKey
 
+retry_count = {}
 
-@retry(retry=retry_if_exception_type(DecryptException), stop=stop_after_attempt(3),
+
+@retry(retry=retry_if_exception_type(RetryableDecryptException), stop=stop_after_attempt(3),
        before_sleep=before_sleep_log(logger, logging.WARN))
 async def decrypt(info: SongInfo, keys: list[str], manifest: Datum, device: Device) -> bytes:
     async with device.decryptLock:
@@ -34,10 +36,19 @@ async def decrypt(info: SongInfo, keys: list[str], manifest: Datum, device: Devi
             last_index = sample.descIndex
             try:
                 result = await decrypt_sample(writer, reader, sample)
-            except DecryptException as e:
-                writer.write(bytes([0, 0, 0, 0]))
-                writer.close()
-                raise e
+            except RetryableDecryptException as e:
+                if 0 <= retry_count.get(device.device.serial, 0) < 3 or 4 <= retry_count.get(device.device.serial, 0) < 6:
+                    logger.warning(f"Failed to decrypt song: {manifest.attributes.artistName} - {manifest.attributes.name}, retrying")
+                    writer.write(bytes([0, 0, 0, 0]))
+                    writer.close()
+                    raise e
+                elif retry_count == 3:
+                    logger.warning(f"Failed to decrypt song: {manifest.attributes.artistName} - {manifest.attributes.name}, re-injecting")
+                    device.restart_inject_frida()
+                    raise e
+                else:
+                    logger.error(f"Failed to decrypt song: {manifest.attributes.artistName} - {manifest.attributes.name}")
+                    raise DecryptException
             decrypted += result
         writer.write(bytes([0, 0, 0, 0]))
         writer.close()
@@ -49,5 +60,5 @@ async def decrypt_sample(writer: asyncio.StreamWriter, reader: asyncio.StreamRea
     writer.write(sample.data)
     result = await reader.read(len(sample.data))
     if not result:
-        raise DecryptException
+        raise RetryableDecryptException
     return result
diff --git a/src/exceptions.py b/src/exceptions.py
index 0f45cab..d3a9046 100644
--- a/src/exceptions.py
+++ b/src/exceptions.py
@@ -24,3 +24,7 @@ class NotTimeSyncedLyricsException(Exception):
 
 class CodecNotFoundException(Exception):
     ...
+
+
+class RetryableDecryptException(Exception):
+    ...