diff --git a/test/test_utils.py b/test/test_utils.py index 971f68bbec..c44219cd2c 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -9,6 +9,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import contextlib import datetime as dt +import email.message import io import itertools import json @@ -48,6 +49,7 @@ from yt_dlp.utils import ( determine_ext, determine_file_encoding, dfxp2srt, + disposition_filename, encode_base_n, encode_compat_str, expand_path, @@ -126,6 +128,7 @@ from yt_dlp.utils import ( url_basename, url_or_none, urlencode_postdata, + urlhandle_detect_ext, urljoin, urshift, variadic, @@ -2123,6 +2126,32 @@ Line 1 self.assertEqual(get_compatible_ext( vcodecs=['av1'], acodecs=['mp4a'], vexts=['webm'], aexts=['m4a'], preferences=('webm', 'mkv')), 'mkv') + def test_disposition_filename(self): + test_cases = ( + (None, None), + ('attachment', None), + ('attachment; filename="video.mp4"', 'video.mp4'), + ('attachment; FILENAME="video.mp4"', 'video.mp4'), + ('inline; filename=video.mp4', 'video.mp4'), + ('attachment; filename=" video.mp4 "', 'video.mp4'), + ("attachment; filename*=UTF-8''r%C3%A9sum%C3%A9.mp4", 'résumé.mp4'), + ("attachment; filename=fallback.mp4; filename*=UTF-8''actual.webm", 'actual.webm'), + ("attachment; filename*=UTF-8''actual.webm; filename=fallback.mp4", 'actual.webm'), + ) + + for content_disposition, expected in test_cases: + headers = email.message.Message() + if content_disposition is not None: + headers['Content-Disposition'] = content_disposition + with self.subTest(content_disposition=content_disposition): + self.assertEqual(disposition_filename(headers), expected) + + def test_urlhandle_detect_ext(self): + headers = email.message.Message() + headers['Content-Disposition'] = "attachment; filename=fallback.mp4; filename*=UTF-8''actual.webm" + headers['Content-Type'] = 'video/mp4' + self.assertEqual(urlhandle_detect_ext(unittest.mock.Mock(headers=headers)), 'webm') + def test_try_call(self): def total(*x, **kwargs): return sum(x) + sum(kwargs.values()) diff --git a/yt_dlp/extractor/googledrive.py b/yt_dlp/extractor/googledrive.py index 91c9f60cd8..824aa86692 100644 --- a/yt_dlp/extractor/googledrive.py +++ b/yt_dlp/extractor/googledrive.py @@ -3,6 +3,7 @@ import re from .common import InfoExtractor from ..utils import ( determine_ext, + disposition_filename, extract_attributes, filter_dict, get_element_by_class, @@ -184,9 +185,7 @@ class GoogleDriveIE(InfoExtractor): def add_source_format(urlh): nonlocal title if not title: - title = self._search_regex( - r'\bfilename="([^"]+)"', urlh.headers.get('Content-Disposition'), - 'title', default=None) + title = disposition_filename(urlh.headers) or None formats.append({ # Use redirect URLs as download URLs in order to calculate # correct cookies in _calc_cookies. @@ -198,7 +197,7 @@ class GoogleDriveIE(InfoExtractor): 'format_id': 'source', 'quality': 1, }) - if urlh.headers.get('Content-Disposition'): + if urlh.headers.get_content_disposition(): add_source_format(urlh) else: confirmation_webpage = self._webpage_read_content( @@ -209,7 +208,7 @@ class GoogleDriveIE(InfoExtractor): get_element_html_by_id('download-form', confirmation_webpage) or '').get('action') if confirmed_source_url: urlh = request_source_file(confirmed_source_url, 'confirmed source', data=b'') - if urlh and urlh.headers.get('Content-Disposition'): + if urlh and urlh.headers.get_content_disposition(): add_source_format(urlh) else: self.report_warning( diff --git a/yt_dlp/utils/_utils.py b/yt_dlp/utils/_utils.py index b8a7bc447c..03c743a35e 100644 --- a/yt_dlp/utils/_utils.py +++ b/yt_dlp/utils/_utils.py @@ -3129,18 +3129,26 @@ def get_compatible_ext(*, vcodecs, acodecs, vexts, aexts, preferences=None): return 'mkv' if allow_mkv else preferences[-1] -def urlhandle_detect_ext(url_handle, default=NO_DEFAULT): - getheader = url_handle.headers.get +def disposition_filename(headers): + filename = None + for name, value in headers.get_params(header='Content-Disposition') or (): + if name.lower() != 'filename': + continue + filename = value + if isinstance(value, tuple): # Prefer filename* over filename + break - if cd := getheader('Content-Disposition'): - if m := re.match(r'attachment;\s*filename="(?P[^"]+)"', cd): - if ext := determine_ext(m.group('filename'), default_ext=None): - return ext + return email.utils.collapse_rfc2231_value(filename).strip() if filename is not None else None + + +def urlhandle_detect_ext(url_handle, default=NO_DEFAULT): + headers = url_handle.headers return ( - determine_ext(getheader('x-amz-meta-name'), default_ext=None) - or getheader('x-amz-meta-file-type') - or mimetype2ext(getheader('Content-Type'), default=default)) + determine_ext(disposition_filename(headers), default_ext=None) + or determine_ext(headers.get('x-amz-meta-name'), default_ext=None) + or headers.get('x-amz-meta-file-type') + or mimetype2ext(headers.get('Content-Type'), default=default)) def encode_data_uri(data, mime_type):