diff --git a/server/szurubooru/func/net.py b/server/szurubooru/func/net.py index e00a6c35..9dff3c45 100644 --- a/server/szurubooru/func/net.py +++ b/server/szurubooru/func/net.py @@ -7,6 +7,7 @@ from threading import Thread from typing import Any, Dict, List from szurubooru import config, errors +from szurubooru.func import mime logger = logging.getLogger(__name__) _dl_chunk_size = 2 ** 15 @@ -22,8 +23,12 @@ class DownloadTooLargeError(DownloadError): def download(url: str, use_video_downloader: bool = False) -> bytes: assert url + youtube_dl_error = None if use_video_downloader: - url = _get_youtube_dl_content_url(url) + try: + url = _get_youtube_dl_content_url(url) or url + except errors.ThirdPartyError as ex: + youtube_dl_error = ex request = urllib.request.Request(url) if config.config["user_agent"]: @@ -41,11 +46,18 @@ def download(url: str, use_video_downloader: bool = False) -> bytes: content_buffer += chunk except urllib.error.HTTPError as ex: raise DownloadError(url) from ex + + if ( + youtube_dl_error + and mime.get_mime_type(content_buffer) == "application/octet-stream" + ): + raise youtube_dl_error + return content_buffer def _get_youtube_dl_content_url(url: str) -> str: - cmd = ["youtube-dl", "--format", "best"] + cmd = ["youtube-dl", "--format", "best", "--no-playlist"] if config.config["user_agent"]: cmd.extend(["--user-agent", config.config["user_agent"]]) cmd.extend(["--get-url", url]) @@ -85,6 +97,6 @@ def _post_to_webhook(webhook: str, payload: Dict[str, Any]) -> int: f"Webhook {webhook} returned {res.status} {res.reason}" ) return res.status - except urllib.error.URLError as e: - logger.warning(f"Unable to call webhook {webhook}: {str(e)}") + except urllib.error.URLError as ex: + logger.warning(f"Unable to call webhook {webhook}: {ex}") return 400 diff --git a/server/szurubooru/tests/func/test_net.py b/server/szurubooru/tests/func/test_net.py index 19e31b6b..c5b4c73e 100644 --- a/server/szurubooru/tests/func/test_net.py +++ b/server/szurubooru/tests/func/test_net.py @@ -85,6 +85,10 @@ def test_too_large_download(url): "https://upload.wikimedia.org/wikipedia/commons/a/ad/Utah_teapot.png", # noqa: E501 "cfadcbdeda1204dc1363ee5c1969191f26be2e41", ), + ( + "https://i.imgur.com/GPgh0AN.jpg", + "26861a4663fedae48e5beed3eec5156ded20640f", + ), ], ) def test_content_download(url, expected_sha1):