[utils] Improve Content-Disposition filename parsing (#17727)

Authored by: doe1080

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())

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(

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<filename>[^"]+)"', 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):