From ad49f2cf87b39d1a26d405f86bc45f1bfcbccc0d Mon Sep 17 00:00:00 2001 From: pawank925 <229111348+pawank925@users.noreply.github.com> Date: Fri, 16 Oct 2026 16:13:41 +0530 Subject: [PATCH] Set a Wikimedia-compliant User-Agent in stream_file --- CHANGELOG.md | 1 + src/zimscraperlib/constants.py | 1 + src/zimscraperlib/download.py | 21 +++++++- tests/download/test_download.py | 96 ++++++++++++++++++++++++++++++++- 4 files changed, 117 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e1f874c5..b260efc1 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/zimscraperlib/constants.py b/src/zimscraperlib/constants.py index 485342b3..e07165f5 100644 --- a/src/zimscraperlib/constants.py +++ b/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" diff --git a/src/zimscraperlib/download.py b/src/zimscraperlib/download.py index 63c1ed06..64bccf9d 100644 --- a/src/zimscraperlib/download.py +++ b/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, diff --git a/tests/download/test_download.py b/tests/download/test_download.py index 3a1bf8e7..fe4b001a 100644 --- a/tests/download/test_download.py +++ b/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"