diff --git a/plugins/tools/_download_runner.py b/plugins/tools/_download_runner.py index 3e8ca44..e9f59d1 100644 --- a/plugins/tools/_download_runner.py +++ b/plugins/tools/_download_runner.py @@ -5,6 +5,7 @@ import argparse import contextlib import email.message +import email.utils import fcntl import hashlib import http.client @@ -56,10 +57,6 @@ "image/", "text/", ) -_CONTENT_DISPOSITION_NAME_RE = re.compile( - r"""filename\*?=(?:UTF-8''|["']?)([^;"']+)""", - re.IGNORECASE, -) class DownloadError(RuntimeError): @@ -270,11 +267,21 @@ def _preflight( def _header_filename(headers: email.message.Message) -> str: - disposition = headers.get("Content-Disposition") or "" - match = _CONTENT_DISPOSITION_NAME_RE.search(disposition) - if not match: - return "" - return urllib.parse.unquote(match.group(1)).strip() + extended = "" + fallback = "" + for name, value in headers.get_params(header="content-disposition", failobj=[]): + if name != "filename": + continue + # The parser represents filename* as an RFC 2231 tuple. Prefer it + # over the ordinary filename, regardless of parameter order, but only + # when it decodes to something usable: an empty filename* must not + # discard a valid ordinary filename. + if isinstance(value, tuple): + if not extended: + extended = email.utils.collapse_rfc2231_value(value).strip() + elif not fallback: + fallback = value.strip() + return extended or fallback def _safe_filename(candidate: str) -> str: diff --git a/tests/test_download_filenames.py b/tests/test_download_filenames.py new file mode 100644 index 0000000..bbedf59 --- /dev/null +++ b/tests/test_download_filenames.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from email.message import Message + +import pytest + +from plugins.tools._download_runner import _header_filename, _resolve_basename + + +@pytest.mark.parametrize( + ("disposition", "expected"), + [ + ('attachment; filename="report.pdf"', "report.pdf"), + ('attachment; filename="report;final.pdf"', "report;final.pdf"), + ('attachment; filename="O\'Reilly.pdf"', "O'Reilly.pdf"), + ('attachment; filename = "report.pdf"', "report.pdf"), + ('attachment; filename="report%20draft.pdf"', "report%20draft.pdf"), + ("attachment; filename*=UTF-8''report%20final.pdf", "report final.pdf"), + ("attachment; filename*=UTF-8'en'report%20final.pdf", "report final.pdf"), + ("attachment; filename*=UTF-8''caf%C3%A9.pdf", "café.pdf"), + ("attachment; filename*=ISO-8859-1''caf%E9.pdf", "café.pdf"), + ( + 'attachment; filename="fallback.txt"; filename*=UTF-8\'\'actual.pdf', + "actual.pdf", + ), + ( + 'attachment; filename*=UTF-8\'\'actual.pdf; filename="fallback.txt"', + "actual.pdf", + ), + ('attachment; filename="paper.pdf"; filename*=UTF-8\'\'', "paper.pdf"), + ("attachment; filename*=UTF-8''; filename=\"paper.pdf\"", "paper.pdf"), + ("attachment; filename*=UTF-8''%20%20; filename=\"paper.pdf\"", "paper.pdf"), + ("attachment", ""), + ], +) +def test_response_filename(disposition: str, expected: str) -> None: + headers = Message() + headers["Content-Disposition"] = disposition + + assert _header_filename(headers) == expected + + +@pytest.mark.parametrize( + ("path_arg", "disposition", "expected"), + [ + ("", 'attachment; filename="report;final.pdf"', "report_final.pdf"), + ("", 'attachment; filename="O\'Reilly.pdf"', "O_Reilly.pdf"), + ( + "", + 'attachment; filename="fallback.txt"; filename*=UTF-8\'\'actual.pdf', + "actual.pdf", + ), + ("", 'attachment; filename="../../outside.pdf"', "outside.pdf"), + ("", "attachment; filename*=UTF-8''..%2Foutside.pdf", "outside.pdf"), + ("", "attachment", "url name.pdf"), + ("", 'attachment; filename="paper.pdf"; filename*=UTF-8\'\'', "paper.pdf"), + ("", "", "url name.pdf"), + ("chosen.pdf", 'attachment; filename="report;final.pdf"', "chosen.pdf"), + ], +) +def test_download_basename(path_arg: str, disposition: str, expected: str) -> None: + headers = Message() + if disposition: + headers["Content-Disposition"] = disposition + + assert _resolve_basename(path_arg, "https://example.com/url%20name.pdf", headers) == expected