Set a Wikimedia-compliant User-Agent in stream_file

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Index ePub documents for full-text search, with a new `get_epub_index_data` helper and automatic indexing of ePub items, like PDFs already are (#333)
+- Send a Wikimedia-compliant default `User-Agent` in `stream_file` when downloading from `upload.wikimedia.org` (#214)
 
 ### Fixed
 

src/zimscraperlib/constants.py

@@ -9,6 +9,7 @@
 VERSION = __version__
 CONTACT = "dev@openzim.org"
 DEFAULT_USER_AGENT = f"{NAME}/{__version__} ({CONTACT})"
+WIKIMEDIA_UPLOAD_HOST = "upload.wikimedia.org"
 
 UTF8 = "UTF-8"
 

src/zimscraperlib/download.py

@@ -2,6 +2,7 @@
 import subprocess
 from concurrent.futures import Future, ThreadPoolExecutor
 from typing import Any, ClassVar
+from urllib.parse import urlparse
 
 import requests
 import requests.adapters
@@ -10,7 +11,11 @@
 import yt_dlp as youtube_dl  # pyright: ignore[reportMissingTypeStubs]
 
 from zimscraperlib import logger
-from zimscraperlib.constants import DEFAULT_WEB_REQUESTS_TIMEOUT
+from zimscraperlib.constants import (
+    DEFAULT_USER_AGENT,
+    DEFAULT_WEB_REQUESTS_TIMEOUT,
+    WIKIMEDIA_UPLOAD_HOST,
+)
 from zimscraperlib.typing import SupportsSeekableWrite, SupportsWrite
 
 
@@ -181,6 +186,16 @@ def get_session(max_retries: int | None = 5) -> requests.Session:
     return session
 
 
+def is_wikimedia_upload_url(url: str) -> bool:
+    """Whether the URL points to a file hosted on Wikimedia upload servers"""
+    return (urlparse(url).hostname or "").lower() == WIKIMEDIA_UPLOAD_HOST
+
+
+def has_user_agent_header(headers: dict[str, str] | None) -> bool:
+    """Whether a User-Agent header is already set, whatever its casing"""
+    return any(key.lower() == "user-agent" for key in (headers or {}))
+
+
 def stream_file(  # noqa: PLR0917
     url: str,
     fpath: pathlib.Path | None = None,
@@ -213,6 +228,10 @@ def stream_file(  # noqa: PLR0917
 
     if not session:
         session = get_session(max_retries)
+
+    if is_wikimedia_upload_url(url) and not has_user_agent_header(headers):
+        headers = {**(headers or {}), "User-Agent": DEFAULT_USER_AGENT}
+
     resp = session.get(
         url,
         stream=True,

tests/download/test_download.py

@@ -10,11 +10,13 @@
 import requests.structures
 from yt_dlp.utils import DownloadError
 
-from zimscraperlib.constants import DEFAULT_WEB_REQUESTS_TIMEOUT
+from zimscraperlib.constants import DEFAULT_USER_AGENT, DEFAULT_WEB_REQUESTS_TIMEOUT
 from zimscraperlib.download import (
     BestMp4,
     BestWebm,
     YoutubeDownloader,
+    has_user_agent_header,
+    is_wikimedia_upload_url,
     save_large_file,
     stream_file,
 )
@@ -117,6 +119,98 @@ def test_first_block_download_custom_session(mocker: Mock, valid_http_url: str):
     requests.Session.assert_not_called()  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
 
 
+@pytest.mark.parametrize(
+    "url,expected",
+    [
+        ("https://upload.wikimedia.org/wikipedia/commons/4/4a/Commons-logo.svg", True),
+        ("http://upload.wikimedia.org/wikipedia/commons/a/b.png", True),
+        ("https://UPLOAD.WIKIMEDIA.ORG/wikipedia/commons/a/b.png", True),
+        ("https://commons.wikimedia.org/wiki/File:Commons-logo.svg", False),
+        ("https://wikimedia.org/", False),
+        ("https://example.com/favicon.ico", False),
+        ("not-an-url", False),
+    ],
+)
+def test_is_wikimedia_upload_url(url: str, *, expected: bool):
+    assert is_wikimedia_upload_url(url) is expected
+
+
+@pytest.mark.parametrize(
+    "headers,expected",
+    [
+        (None, False),
+        ({}, False),
+        ({"Accept": "image/png"}, False),
+        ({"User-Agent": "zimscraperlib-test"}, True),
+        ({"user-agent": "zimscraperlib-test"}, True),
+    ],
+)
+def test_has_user_agent_header(headers: dict[str, str] | None, *, expected: bool):
+    assert has_user_agent_header(headers) is expected
+
+
+@pytest.mark.parametrize(
+    "headers",
+    [None, {}, {"Accept": "image/png"}],
+)
+def test_wikimedia_upload_default_user_agent(
+    mocker: Mock, headers: dict[str, str] | None
+):
+    url = "https://upload.wikimedia.org/wikipedia/commons/4/4a/Commons-logo.svg"
+    byte_stream = io.BytesIO()
+    custom_session = mocker.Mock(spec=requests.Session)
+
+    expected_response = requests.Response()
+    expected_response.status_code = 200
+    expected_response.raw = io.BytesIO(b"Whatever\n")
+    custom_session.get.return_value = expected_response
+
+    mocker.patch("requests.Session")
+    stream_file(
+        url=url,
+        byte_stream=byte_stream,
+        headers=headers,
+        only_first_block=True,
+        session=custom_session,
+    )
+    custom_session.get.assert_called_once_with(
+        url,
+        stream=True,
+        proxies=None,
+        headers={**(headers or {}), "User-Agent": DEFAULT_USER_AGENT},
+        timeout=DEFAULT_WEB_REQUESTS_TIMEOUT,
+    )
+    requests.Session.assert_not_called()  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
+
+
+def test_wikimedia_upload_custom_user_agent(mocker: Mock):
+    url = "https://upload.wikimedia.org/wikipedia/commons/4/4a/Commons-logo.svg"
+    byte_stream = io.BytesIO()
+    custom_session = mocker.Mock(spec=requests.Session)
+
+    expected_response = requests.Response()
+    expected_response.status_code = 200
+    expected_response.raw = io.BytesIO(b"Whatever\n")
+    custom_session.get.return_value = expected_response
+
+    mocker.patch("requests.Session")
+    stream_file(
+        url=url,
+        byte_stream=byte_stream,
+        headers={"User-Agent": "myscraper/1.0 (me@example.org)"},
+        only_first_block=True,
+        session=custom_session,
+    )
+    custom_session.get.assert_called_once_with(
+        url,
+        stream=True,
+        proxies=None,
+        headers={"User-Agent": "myscraper/1.0 (me@example.org)"},
+        timeout=DEFAULT_WEB_REQUESTS_TIMEOUT,
+    )
+    requests.Session.assert_not_called()  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
+
+
 @pytest.mark.slow
 def test_user_agent():
     ua = "zimscraperlib-test"