diff --git a/tests/core/test_playlist_encoding.py b/tests/core/test_playlist_encoding.py index 7003251f..c098101c 100644 --- a/tests/core/test_playlist_encoding.py +++ b/tests/core/test_playlist_encoding.py @@ -1,6 +1,7 @@ -"""Unit tests for M3U8 response decoding. Reading `.text` on a response with no encoding -makes requests run chardet over the whole body, so both playlist fetch paths set utf-8 -first. A charset the server did send must survive.""" +"""Unit tests for manifest response decoding. Reading `.text` on a response with no encoding +makes requests run chardet over the whole body, so the playlist fetch paths set utf-8 +first. A charset the server did send must survive. An MPD is parsed from bytes, so the +XML declaration gives the encoding.""" from __future__ import annotations @@ -9,7 +10,9 @@ from typing import Any, Optional import pytest import requests +from unshackle.core.manifests.dash import DASH from unshackle.core.manifests.hls import HLS +from unshackle.core.session import RnetResponse, utf8_text from unshackle.core.tracks import Video pytestmark = pytest.mark.unit @@ -24,6 +27,8 @@ def make_response(body: bytes, url: str, encoding: Optional[str]) -> requests.Re res._content = body res.url = url res.encoding = encoding + if encoding: + res.headers["Content-Type"] = f"application/vnd.apple.mpegurl; charset={encoding}" return res @@ -74,3 +79,61 @@ def test_load_drm_from_playlist_keeps_a_charset_the_server_sent() -> None: make_track(res).load_drm_from_playlist() assert res.encoding == "iso-8859-1" + + +MPD_LATIN1 = 'café'.encode("iso-8859-1") + + +class StubRnetResponse: + """The part of an rnet response that `RnetResponse` reads.""" + + ok = True + status_code = 200 + encoding = None + + def __init__(self, body: bytes, url: str) -> None: + self.body = body + self.url = url + + def bytes(self) -> bytes: + return self.body + + +def test_dash_from_url_follows_the_xml_declaration() -> None: + res = make_response(MPD_LATIN1, "https://example.com/manifest.mpd", None) + dash = DASH.from_url("https://example.com/manifest.mpd", FakeSession(res)) + + assert dash.manifest.findtext("Title") == "café" + + +def test_dash_from_url_reads_an_rnet_response() -> None: + res = RnetResponse(StubRnetResponse(MPD_LATIN1, "https://example.com/manifest.mpd")) + dash = DASH.from_url("https://example.com/manifest.mpd", FakeSession(res)) # type: ignore[arg-type] + + assert dash.manifest.findtext("Title") == "café" + + +def test_utf8_text_reads_an_rnet_response() -> None: + assert utf8_text(RnetResponse(StubRnetResponse(MEDIA, "https://example.com/video/720.m3u8"))) == MEDIA.decode() + + +def test_utf8_text_ignores_the_requests_default_for_a_text_type() -> None: + res = make_response('#EXTM3U\n#EXT-X-MEDIA:TYPE=AUDIO,NAME="中文"\n'.encode(), "https://example.com/a.m3u8", None) + res.headers["Content-Type"] = "text/plain" + res.encoding = "ISO-8859-1" + + assert "中文" in utf8_text(res) + + +def test_dash_from_url_uses_the_http_charset_without_an_xml_declaration() -> None: + res = make_response("中文".encode("gbk"), "https://example.com/manifest.mpd", "gbk") + dash = DASH.from_url("https://example.com/manifest.mpd", FakeSession(res)) + + assert dash.manifest.findtext("Title") == "中文" + + +def test_dash_from_url_rejects_an_empty_body() -> None: + res = make_response(b"", "https://example.com/manifest.mpd", None) + + with pytest.raises(ValueError, match="must be provided"): + DASH.from_url("https://example.com/manifest.mpd", FakeSession(res)) diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index bb288176..a70d4e80 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -31,6 +31,7 @@ from unshackle.core.constants import AUDIO_CODEC_MAP, DYNAMIC_RANGE_MAP, VIDEO_C from unshackle.core.providers.anilist import parse_anilist_ref from unshackle.core.proxies.resolve import initialize_proxy_providers, resolve_proxy from unshackle.core.services import Services +from unshackle.core.session import utf8_text from unshackle.core.titles import Episode, Movie, Song, Title_T from unshackle.core.tracks import Audio, Subtitle, Tracks, Video from unshackle.core.utilities import declared_kwargs, kind_only_conflict @@ -4327,7 +4328,7 @@ def ensure_track_drm(track: Any, session: Any = None, init_data: Optional[bytes] from unshackle.core.manifests import HLS http = getattr(track, "session", None) or session - playlist = m3u8.loads(http.get(track.url).text, uri=track.url) if http else m3u8.load(track.url) + playlist = m3u8.loads(utf8_text(http.get(track.url)), uri=track.url) if http else m3u8.load(track.url) keys = [k for k in (playlist.keys or []) + (playlist.session_keys or []) if k is not None] for key in keys: try: @@ -4807,14 +4808,11 @@ async def session_segment_filter_handler( def unwanted_segments() -> list[str]: import m3u8 - import requests response = session.service_instance.session.get(track.url) if not getattr(response, "ok", True): raise ValueError(f"playlist request failed: {response.status_code}") - if isinstance(response, requests.Response): - response.encoding = response.encoding or "utf-8" - playlist = m3u8.loads(response.text, uri=str(track.url)) + playlist = m3u8.loads(utf8_text(response), uri=str(track.url)) return [segment.absolute_uri for segment in playlist.segments if segment_filter(segment)] try: diff --git a/unshackle/core/drm/playready.py b/unshackle/core/drm/playready.py index f75d9d61..9cea94cd 100644 --- a/unshackle/core/drm/playready.py +++ b/unshackle/core/drm/playready.py @@ -25,6 +25,7 @@ from unshackle.core.config import config from unshackle.core.console import console from unshackle.core.constants import AnyTrack from unshackle.core.drm import key_args as drm_key_args +from unshackle.core.session import utf8_text from unshackle.core.utilities import get_boxes, log_event from unshackle.core.utils.subprocess import ffprobe @@ -180,7 +181,7 @@ class PlayReady: if track.descriptor == track.Descriptor.HLS: m3u_url = track.url - master = m3u8.loads(session.get(m3u_url).text, uri=m3u_url) + master = m3u8.loads(utf8_text(session.get(m3u_url)), uri=m3u_url) pssh_boxes.extend( Box.parse(base64.b64decode(x.uri.split(",")[-1])) for x in (master.session_keys or master.keys) diff --git a/unshackle/core/drm/widevine.py b/unshackle/core/drm/widevine.py index 377e2b9b..ced467aa 100644 --- a/unshackle/core/drm/widevine.py +++ b/unshackle/core/drm/widevine.py @@ -22,6 +22,7 @@ from unshackle.core.config import config from unshackle.core.console import console from unshackle.core.constants import AnyTrack from unshackle.core.drm import key_args as drm_key_args +from unshackle.core.session import utf8_text from unshackle.core.utilities import get_boxes, log_event from unshackle.core.utils.subprocess import ffprobe @@ -94,7 +95,7 @@ class Widevine: if track.descriptor == track.Descriptor.HLS: m3u_url = track.url - master = m3u8.loads(session.get(m3u_url).text, uri=m3u_url) + master = m3u8.loads(utf8_text(session.get(m3u_url)), uri=m3u_url) pssh_boxes.extend( Box.parse(base64.b64decode(x.uri.split(",")[-1])) for x in (master.session_keys or master.keys) diff --git a/unshackle/core/import_service.py b/unshackle/core/import_service.py index 8ba86456..35f39385 100644 --- a/unshackle/core/import_service.py +++ b/unshackle/core/import_service.py @@ -25,6 +25,7 @@ from unshackle.core.remote_service import ( resolve_proxy_arg, segment_uri_filter, ) +from unshackle.core.session import utf8_text from unshackle.core.titles import Episode, Movies, Series, Title_T, Titles_T, remap_titles from unshackle.core.tracks import Audio, Chapter, Chapters, Subtitle, Tracks, Video from unshackle.core.tracks.attachment import Attachment @@ -69,7 +70,7 @@ def hls_playlist_kids(track: Any, session: Optional[requests.Session] = None) -> """ http = session or requests.Session() try: - playlist = m3u8.loads(http.get(str(track.url)).text, uri=str(track.url)) + playlist = m3u8.loads(utf8_text(http.get(str(track.url))), uri=str(track.url)) keys = [k for k in playlist.keys if k is not None and k.method != "NONE"] except Exception as e: log.warning(f"Cannot read the media playlist of {track.id} for its KIDs: {e!r}") diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index edf8b395..0aef0528 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -19,7 +19,7 @@ from zlib import crc32 import requests from langcodes import Language, tag_is_valid -from lxml.etree import Element, ElementTree +from lxml.etree import Element, ElementTree, XMLSyntaxError from pyplayready.system.pssh import PSSH as PR_PSSH from pywidevine.cdm import Cdm as WidevineCdm from pywidevine.pssh import PSSH @@ -118,6 +118,7 @@ class DASH: @classmethod def from_url(cls, url: str, session: Optional[Union[Session, RnetSession]] = None, **args: Any) -> DASH: + """Fetch and parse an MPD. The XML declaration gives the encoding, then the HTTP charset.""" if not url: raise requests.URLRequired("DASH manifest URL must be provided for relative path computations.") if not isinstance(url, str): @@ -138,11 +139,15 @@ class DASH: log_event( "manifest_dash_fetch", level="DEBUG", - message=f"Fetched DASH manifest ({len(res.text)} bytes)", - context={"url": safe_display_url(url), "size": len(res.text)}, + message=f"Fetched DASH manifest ({len(res.content)} bytes)", + context={"url": safe_display_url(url), "size": len(res.content)}, ) - return DASH.from_text(res.text, url) + try: + return cls(load_xml(res.content), url) + except XMLSyntaxError: + # no declaration and not UTF-8: only the HTTP charset can name the encoding + return cls.from_text(res.text, url) @classmethod def from_text(cls, text: str, url: str) -> DASH: diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 9777c0bb..f1eb85c6 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -37,7 +37,7 @@ from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use from unshackle.core.drm.verify import decrypt_track from unshackle.core.events import events from unshackle.core.manifests.dash import RollingMerge -from unshackle.core.session import RnetResponse, RnetSession +from unshackle.core.session import RnetResponse, RnetSession, utf8_text from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume from unshackle.core.tracks.track import DRM_PREFERENCE_TYPES, assert_fragments_decrypted from unshackle.core.utilities import get_extension, is_close_match, log_event, try_ensure_utf8 @@ -81,17 +81,11 @@ class HLS: res = session.get(url, **args) - if isinstance(res, requests.Response): - if not res.ok: - raise requests.ConnectionError("Failed to request the M3U(8) document.", response=res) - res.encoding = res.encoding or "utf-8" - content = res.text - elif isinstance(res, RnetResponse): - if not res.ok: - raise requests.ConnectionError("Failed to request the M3U(8) document.", response=res) - content = res.text - else: + if not isinstance(res, (requests.Response, RnetResponse)): raise TypeError(f"Expected response to be a requests.Response or rnet.Response, not {type(res)}") + if not res.ok: + raise requests.ConnectionError("Failed to request the M3U(8) document.", response=res) + content = utf8_text(res) master = m3u8.loads(content, uri=url) @@ -397,9 +391,7 @@ class HLS: session = Session() res = session.get(variant_url) - if isinstance(res, requests.Response): - res.encoding = res.encoding or "utf-8" - variant = m3u8.loads(res.text, uri=variant_url) + variant = m3u8.loads(utf8_text(res), uri=variant_url) if not variant.segments: return None @@ -659,17 +651,13 @@ class HLS: master = m3u8.load(str(track.from_file)) else: response = session.get(track.url) - if isinstance(response, requests.Response) or isinstance(response, RnetResponse): - if not response.ok: - log.error(f"Failed to request the invariant M3U8 playlist: {response.status_code}") - sys.exit(1) - if isinstance(response, requests.Response): - response.encoding = response.encoding or "utf-8" - playlist_text = response.text - else: + if not isinstance(response, (requests.Response, RnetResponse)): raise TypeError(f"Expected response to be a requests.Response or rnet.Response, not {type(response)}") + if not response.ok: + log.error(f"Failed to request the invariant M3U8 playlist: {response.status_code}") + sys.exit(1) - master = m3u8.loads(playlist_text, uri=track.url) + master = m3u8.loads(utf8_text(response), uri=track.url) if not master.segments: log.error("Track's HLS playlist has no segments, expecting an invariant M3U8 playlist.") @@ -1385,8 +1373,7 @@ class HLS: if not value and data.uri: if not session: session = Session() - res = session.get(urljoin(manifest.base_uri or "", data.uri)) - value = res.text + value = utf8_text(session.get(urljoin(manifest.base_uri or "", data.uri))) if not value: continue diff --git a/unshackle/core/session.py b/unshackle/core/session.py index f1c8fff3..8257e31d 100644 --- a/unshackle/core/session.py +++ b/unshackle/core/session.py @@ -15,7 +15,7 @@ from typing import Any, Optional from urllib.parse import urlencode, urlparse, urlunparse import rnet -from requests import HTTPError, Request +from requests import HTTPError, Request, Response from requests.structures import CaseInsensitiveDict from unshackle.core.config import config @@ -227,6 +227,18 @@ class RnetResponse: pass +def utf8_text(res: Response | RnetResponse) -> str: + """Return the response text, decoded as UTF-8 unless the server sent a charset. + + With no charset, requests assumes ISO-8859-1 for a `text/*` type and runs charset detection + on the full body for any other type. Use this only for a format that is UTF-8 by + specification, such as M3U8. + """ + if isinstance(res, Response) and "charset" not in res.headers.get("content-type", "").lower(): + res.encoding = "utf-8" + return res.text + + class RnetSessionHeaders(CaseInsensitiveDict): """Dict-like headers that write to the rnet client through update().""" diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 2fcafe0f..890aa909 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -20,7 +20,7 @@ import click from construct import ConstructError from langcodes import Language from pymp4.parser import Box -from requests import Response, Session +from requests import Session from requests.adapters import HTTPAdapter, Retry from unshackle.core import binaries @@ -31,7 +31,7 @@ from unshackle.core.downloaders import requests from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine from unshackle.core.drm.verify import decrypt_track from unshackle.core.events import events -from unshackle.core.session import RnetSession +from unshackle.core.session import RnetSession, utf8_text from unshackle.core.tracks import resume from unshackle.core.utilities import get_boxes, log_event, try_ensure_utf8 from unshackle.core.utils.subprocess import ffprobe @@ -1141,9 +1141,7 @@ class Track: session = getattr(self, "session", None) or Session() response = session.get(self.url) - if isinstance(response, Response): - response.encoding = response.encoding or "utf-8" - playlist = m3u8.loads(response.text, self.url) + playlist = m3u8.loads(utf8_text(response), self.url) drm_list = [] diff --git a/unshackle/core/utils/bitrate.py b/unshackle/core/utils/bitrate.py index e518e09f..efdea62e 100644 --- a/unshackle/core/utils/bitrate.py +++ b/unshackle/core/utils/bitrate.py @@ -13,7 +13,7 @@ from urllib.parse import urljoin from requests import Session from unshackle.core.binaries import FFProbe -from unshackle.core.session import RnetSession +from unshackle.core.session import RnetSession, utf8_text from unshackle.core.utils.subprocess import log_tool_run if TYPE_CHECKING: @@ -320,7 +320,7 @@ def extract_hls(track: "Track", session: Union[Session, RnetSession]) -> list[Se playlist_url = track.url if isinstance(track.url, str) else track.url[0] res = session.get(playlist_url, timeout=PROBE_TIMEOUT) - playlist = m3u8.loads(res.text, uri=playlist_url) + playlist = m3u8.loads(utf8_text(res), uri=playlist_url) out: list[Segment] = [] for segment in playlist.segments: