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

This commit is contained in:
Shyam Sunder 2021-01-07 08:28:22 -05:00
parent c732e62844
commit 2b9a4ab786
2 changed files with 20 additions and 4 deletions

View file

@ -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

View file

@ -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):