server/net: prevent youtube-dl errors when downloading image links

2 files changed, 20 insertions, 4 deletions

diff --git a/server/szurubooru/func/net.py b/server/szurubooru/func/net.py
index e00a6c3..9dff3c4 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 19e31b6..c5b4c73 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):