From 400b7d3d6edd7728157724f3338d3a9f36095d4f Mon Sep 17 00:00:00 2001
From: Hellow <74311245+HeIIow2@users.noreply.github.com>
Date: Thu, 22 Jun 2023 14:30:26 +0200
Subject: [PATCH] added the config option how long to wait after yt returns 404

---
 src/music_kraken/connection/connection.py   | 15 ++++++++++-----
 src/music_kraken/pages/youtube.py           | 10 ++++------
 src/music_kraken/utils/config/connection.py |  7 +++++++
 src/music_kraken/utils/shared.py            |  1 +
 4 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/music_kraken/connection/connection.py b/src/music_kraken/connection/connection.py
index cd739c4..e1166c2 100644
--- a/src/music_kraken/connection/connection.py
+++ b/src/music_kraken/connection/connection.py
@@ -22,7 +22,8 @@ class Connection:
             logger: logging.Logger = logging.getLogger("connection"),
             header_values: Dict[str, str] = None,
             accepted_response_codes: Set[int] = None,
-            semantic_not_found: bool = True
+            semantic_not_found: bool = True,
+            sleep_after_404: float = 0.0
     ):
         if proxies is None:
             proxies = PROXIES_LIST
@@ -39,6 +40,7 @@ class Connection:
 
         self.ACCEPTED_RESPONSE_CODES = accepted_response_codes or {200}
         self.SEMANTIC_NOT_FOUND = semantic_not_found
+        self.sleep_after_404 = sleep_after_404
 
         self.session = requests.Session()
         self.session.headers = self.get_header(**self.HEADER_VALUES)
@@ -86,9 +88,11 @@ class Connection:
             headers: dict,
             refer_from_origin: bool = True,
             raw_url: bool = False,
-            wait_on_403: bool = True,
+            sleep_after_404: float = None,
             **kwargs
     ) -> Optional[requests.Response]:
+        if sleep_after_404 is None:
+            sleep_after_404 = self.sleep_after_404
         if try_count >= self.TRIES:
             return
 
@@ -127,9 +131,9 @@ class Connection:
             self.LOGGER.warning(f"{self.HOST.netloc} responded wit {r.status_code} "
                                 f"at {url}. ({try_count}-{self.TRIES})")
             self.LOGGER.debug(r.content)
-            if wait_on_403:
-                self.LOGGER.warning(f"Waiting for 5 seconds.")
-                time.sleep(5)
+            if sleep_after_404 != 0:
+                self.LOGGER.warning(f"Waiting for {sleep_after_404} seconds.")
+                time.sleep(sleep_after_404)
 
         self.rotate()
 
@@ -140,6 +144,7 @@ class Connection:
             url=url,
             timeout=timeout,
             headers=headers,
+            sleep_after_404=sleep_after_404,
             **kwargs
         )
 
diff --git a/src/music_kraken/pages/youtube.py b/src/music_kraken/pages/youtube.py
index 1ef1323..3d63d7c 100644
--- a/src/music_kraken/pages/youtube.py
+++ b/src/music_kraken/pages/youtube.py
@@ -20,7 +20,7 @@ from ..objects import (
 )
 from ..connection import Connection
 from ..utils.support_classes import DownloadResult
-from ..utils.shared import YOUTUBE_LOGGER, INVIDIOUS_INSTANCE, BITRATE, ENABLE_SPONSOR_BLOCK, PIPED_INSTANCE
+from ..utils.shared import YOUTUBE_LOGGER, INVIDIOUS_INSTANCE, BITRATE, ENABLE_SPONSOR_BLOCK, PIPED_INSTANCE, SLEEP_AFTER_YOUTUBE_403
 
 
 """
@@ -148,7 +148,8 @@ class YouTube(Page):
 
         self.download_connection: Connection = Connection(
             host="https://www.youtube.com/",
-            logger=self.LOGGER
+            logger=self.LOGGER,
+            sleep_after_404=SLEEP_AFTER_YOUTUBE_403
         )
         
         # the stuff with the connection is, to ensure sponsorblock uses the proxies, my programm does
@@ -422,11 +423,8 @@ class YouTube(Page):
 
         endpoint = audio_format["url"]
 
-        self.download_connection.stream_into(endpoint, target, description=desc, raw_url=True)
+        return self.download_connection.stream_into(endpoint, target, description=desc, raw_url=True)
 
-        if self.download_connection.get(endpoint, stream=True, raw_url=True):
-            return DownloadResult(total=1)
-        return DownloadResult(error_message=f"Streaming to the file went wrong: {endpoint}, {str(target.file_path)}")
 
     def get_skip_intervals(self, song: Song, source: Source) -> List[Tuple[float, float]]:
         if not ENABLE_SPONSOR_BLOCK:
diff --git a/src/music_kraken/utils/config/connection.py b/src/music_kraken/utils/config/connection.py
index 43fc782..aca2c6d 100644
--- a/src/music_kraken/utils/config/connection.py
+++ b/src/music_kraken/utils/config/connection.py
@@ -100,6 +100,12 @@ class ConnectionSection(Section):
                         "Hidden instances (.onion) will only work, when setting 'tor=true'.",
             value="https://pipedapi.kavin.rocks"
         )
+
+        self.SLEEP_AFTER_YOUTUBE_403 = FloatAttribute(
+            name="sleep_after_youtube_403",
+            description="The time to wait, after youtube returned 403 (in seconds)",
+            value="20"
+        )
         
         self.ALL_YOUTUBE_URLS = UrlListAttribute(
             name="youtube_url",
@@ -126,6 +132,7 @@ class ConnectionSection(Section):
             self.SHOW_DOWNLOAD_ERRORS_THRESHOLD,
             self.INVIDIOUS_INSTANCE,
             self.PIPED_INSTANCE,
+            self.SLEEP_AFTER_YOUTUBE_403,
             self.ALL_YOUTUBE_URLS,
             self.SPONSOR_BLOCK
         ]
diff --git a/src/music_kraken/utils/shared.py b/src/music_kraken/utils/shared.py
index 32ba24a..b8dbe55 100644
--- a/src/music_kraken/utils/shared.py
+++ b/src/music_kraken/utils/shared.py
@@ -124,3 +124,4 @@ have fun :3
 FFMPEG_BINARY: Path = PATHS_SECTION.FFMPEG_BINARY.object_from_value
 
 HASNT_YET_STARTED: bool = MISC_SECTION.HASNT_YET_STARTED.object_from_value
+SLEEP_AFTER_YOUTUBE_403: float = CONNECTION_SECTION.SLEEP_AFTER_YOUTUBE_403.object_from_value