server/net: prevent youtube-dl errors when downloading image links
This commit is contained in:
		
							parent
							
								
									c732e62844
								
							
						
					
					
						commit
						2b9a4ab786
					
				@ -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
 | 
			
		||||
 | 
			
		||||
@ -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):
 | 
			
		||||
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user