From 811e0a19c00c2770c92c8f0f92e0ff7cd8388670 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sat, 17 Oct 2026 18:58:48 -0600 Subject: [PATCH] feat(tracks): probe HDR10 tracks for HDR10+ SEI before selection Relabel HEVC HDR10 tracks whose first fragment carries ST 2094-40 SEI as HDR10P, for dl, the REST API and remote sessions. Config key hdr10plus_probe (default true). --- docs/dev/creating-a-service.md | 17 +- docs/guide/downloading.md | 3 +- docs/reference/configuration/download.md | 31 ++++ docs/reference/configuration/index.md | 2 +- tests/core/test_hdr10plus_probe.py | 193 ++++++++++++++++++++ unshackle/commands/dl.py | 18 +- unshackle/core/api/handlers.py | 7 +- unshackle/core/config.py | 3 + unshackle/core/manifests/hls.py | 4 +- unshackle/core/utils/hdr10plus.py | 223 +++++++++++++++++++++++ unshackle/services/EXAMPLE/__init__.py | 5 - 11 files changed, 484 insertions(+), 22 deletions(-) create mode 100644 tests/core/test_hdr10plus_probe.py create mode 100644 unshackle/core/utils/hdr10plus.py diff --git a/docs/dev/creating-a-service.md b/docs/dev/creating-a-service.md index db24d192..2d5181f6 100644 --- a/docs/dev/creating-a-service.md +++ b/docs/dev/creating-a-service.md @@ -497,13 +497,16 @@ audio, add hand-built subtitles or an `Attachment` (e.g. cover art). Do *not* filter for resolution/bitrate. See `fetch_dash_manifest` in the EXAMPLE service for a thorough demonstration. -!!! note "Flip HDR10 → HDR10+ by hand when you know the platform embeds it" - HDR10+ is a **bitstream (SEI)** feature. The HLS parser intentionally does - not sniff it from the bitstream, so a manifest that embeds HDR10+ SEI but labels - the variant plain HDR10 will parse as HDR10. The label is the only signal - the parser has, and it is wrong. Only the service author knows the platform - embeds HDR10+, so it is on you to correct `video.range` to HDR10+ here for - those services. +!!! note "HDR10+ in tracks labelled HDR10" + HDR10+ is a **bitstream (SEI)** feature, and many manifests label an HDR10+ + variant as plain HDR10. Before unshackle lists or selects tracks, the + [`hdr10plus_probe`](../reference/configuration/download.md#hdr10plus_probe) + reads the first fragment of each HEVC HDR10 track and changes `video.range` to + HDR10+ when the fragment carries HDR10+ metadata. Set `video.range` yourself + only when the probe cannot read the fragment, for example an ISM track whose + SEI is encrypted, or when the platform gives a reliable signal of its own. + If the playback API has no HDR10+ value, request HDR10 for an HDR10+ request + and let the probe relabel the tracks. !!! warning "`get_tracks` is your only chance to capture license inputs" The license callbacks (`get_widevine_license` and the others) fire **much later**, in diff --git a/docs/guide/downloading.md b/docs/guide/downloading.md index 8d214157..1546a464 100644 --- a/docs/guide/downloading.md +++ b/docs/guide/downloading.md @@ -165,7 +165,8 @@ unshackle dl -q 2160 -r HDR10,DV EXAMPLE '...' When a title has HDR10+ but no DV track, `-r HYBRID --generate-dv` makes a Dolby Vision Profile 8.1 track from the HDR10+ metadata. It needs `hdr10plus_tool` as well as `dovi_tool`. When a title has a DV track, unshackle uses that track and ignores the flag. When it has none, -the flag also makes unshackle ask the service for its HDR10+ tracks. +the flag also makes unshackle ask the service for its HDR10+ tracks. A track that the service labels HDR10 but +that carries HDR10+ metadata also qualifies (see [`hdr10plus_probe`](../reference/configuration/download.md#hdr10plus_probe)). ```shell title="Generate Dolby Vision from HDR10+, and keep the HDR10+ file too" unshackle dl -q 2160 -r HYBRID,HDR10P --generate-dv EXAMPLE '...' diff --git a/docs/reference/configuration/download.md b/docs/reference/configuration/download.md index 88887209..47a25677 100644 --- a/docs/reference/configuration/download.md +++ b/docs/reference/configuration/download.md @@ -370,6 +370,37 @@ dual_multi_mode: strict Regional variants of the same base language (e.g. `en-US` and `en-GB`) collapse to a single language in both modes, so they never trigger `{dual}` on their own. +## `hdr10plus_probe` + +- **Type:** `bool`  ·  **Default:** `true` + +Find the HDR10+ tracks that a service labels as HDR10. Before unshackle lists or selects tracks, it +reads the first fragment of HEVC HDR10 tracks from DASH, HLS fMP4 and one-file sources. A track +whose first frame carries HDR10+ (SMPTE ST 2094-40) metadata becomes `HDR10P`. Then `-r HDR10P`, +`--list` and [`--generate-dv`](../../guide/downloading.md#dolby-vision-from-hdr10) use it. + +When two HDR10 tracks have the same resolution, they can be an HDR10 encode and an HDR10+ encode +side by side, so unshackle probes every track. When each track has its own resolution, unshackle +probes the two highest and the two lowest bitrates first: + +| Result | Action | +|--------|--------| +| All four carry HDR10+ | Every HDR10 track of the title becomes `HDR10P`. | +| Some carry HDR10+ | The title has HDR10 and HDR10+ encodes side by side, so unshackle probes every track. | +| None carry HDR10+ | Every track stays `HDR10`. | + +A probe that fails, or a fragment whose metadata is encrypted, keeps the label from the service. +unshackle does not probe ISM tracks. Each probe makes one to three small requests: the 64 KiB +fragment, plus the playlist of an HLS track or the header of a one-file track. The probes run in +parallel, with no more at a time than `--workers` (or `dl.workers`) sets for a track download. + +```yaml +hdr10plus_probe: true +``` + +A remote server probes its tracks before it sends them, so the setting on the server applies to +remote downloads and to the REST API. + ## `generate_dv_mode` - **Type:** `str`  ·  **Default:** `v29` diff --git a/docs/reference/configuration/index.md b/docs/reference/configuration/index.md index e8026089..5848264e 100644 --- a/docs/reference/configuration/index.md +++ b/docs/reference/configuration/index.md @@ -31,7 +31,7 @@ Where the file lives, how unshackle discovers it, and how to edit it: |------|------| | [Directories](directories.md) | `directories`, `filenames`, `services_repo_force` | | [Services & authentication](services.md) | `services`, `credentials`, `firefox_cookies`, `remote_services`, `serve` | -| [Download & processing](download.md) | `dl`, `continue_downloads`, `merge_segments`, `subtitle`, `audio`, `muxing`, `language_tags`, `dual_multi_mode`, `generate_dv_mode` | +| [Download & processing](download.md) | `dl`, `continue_downloads`, `merge_segments`, `subtitle`, `audio`, `muxing`, `language_tags`, `dual_multi_mode`, `generate_dv_mode`, `hdr10plus_probe` | | [Output & naming](output.md) | `output_template`, `tag`, `tag_rules`, `tag_group_name`, `tag_imdb_tmdb`, `chapter_fallback_name`, `unicode_filenames` | | [DRM & CDM](drm.md) | `cdm`, `remote_cdm`, `decryption` | | [Network & proxy](network.md) | `network`, `headers`, `proxy_providers` | diff --git a/tests/core/test_hdr10plus_probe.py b/tests/core/test_hdr10plus_probe.py new file mode 100644 index 00000000..e53cd6f5 --- /dev/null +++ b/tests/core/test_hdr10plus_probe.py @@ -0,0 +1,193 @@ +import logging +import struct +from types import SimpleNamespace +from typing import Any + +import pytest + +from unshackle.core.config import Config, config +from unshackle.core.remote_service import RemoteService +from unshackle.core.tracks import Video +from unshackle.core.utils import hdr10plus +from unshackle.core.utils.hdr10plus import HDR10PLUS_T35, fragment_has_hdr10plus + + +def box(kind: bytes, body: bytes) -> bytes: + return struct.pack(">I", 8 + len(body)) + kind + body + + +def nal(nal_type: int, payload: bytes) -> bytes: + data = bytes([nal_type << 1, 1]) + payload + return struct.pack(">I", len(data)) + data + + +def sei(payload_type: int, payload: bytes) -> bytes: + return nal(39, bytes([payload_type, len(payload)]) + payload + b"\x80") + + +def fragment(*nals: bytes) -> bytes: + return box(b"moof", b"\x00" * 16) + box(b"mdat", b"".join(nals)) + + +def test_hdr10plus_sei_before_the_slice_is_found() -> None: + data = fragment(nal(32, b"vps"), sei(137, b"\x00" * 24), sei(4, HDR10PLUS_T35 + b"\x00" * 8), nal(1, b"\xff" * 64)) + assert fragment_has_hdr10plus(data) is True + + +def test_static_metadata_only_is_plain_hdr10() -> None: + data = fragment(nal(32, b"vps"), sei(137, b"\x00" * 24), sei(4, b"\xb5\x00\x31GA94"), nal(1, b"\xff" * 64)) + assert fragment_has_hdr10plus(data) is False + + +def test_encrypted_or_truncated_bytes_prove_nothing() -> None: + assert fragment_has_hdr10plus(box(b"moof", b"") + box(b"mdat", b"\xff" * 64)) is None + assert fragment_has_hdr10plus(fragment(nal(32, b"vps"))) is None + assert fragment_has_hdr10plus(b"\x47" * 188) is None + + +def hevc(track_id: str, height: int, bitrate: int = 1) -> Any: + return SimpleNamespace( + id=track_id, + range=Video.Range.HDR10, + codec=Video.Codec.HEVC, + width=height * 16 // 9, + height=height, + bitrate=bitrate, + session=None, + ) + + +def ladder( + monkeypatch: pytest.MonkeyPatch, pattern: str, heights: list[int] | None = None +) -> tuple[list[Any], list[str]]: + """HEVC HDR10 tracks, highest bitrate first; '+' marks the ones whose fragment carries HDR10+.""" + probed: list[str] = [] + + def fake_fragment(video: Any, session: Any) -> bytes: + probed.append(video.id) + return video.id[-1].encode() + + monkeypatch.setattr(hdr10plus, "first_fragment", fake_fragment) + monkeypatch.setattr(hdr10plus, "fragment_has_hdr10plus", lambda data: data == b"+") + heights = heights or [2160 - 100 * i for i in range(len(pattern))] + videos = [hevc(f"v{i}{mark}", height, 100 - i) for i, (mark, height) in enumerate(zip(pattern, heights))] + hdr10plus.probe_hdr10plus(videos, None, log=logging.getLogger("test"), workers=4) # type: ignore[arg-type] + return videos, probed + + +def ranges(videos: list[Any]) -> str: + return "".join("+" if v.range == Video.Range.HDR10P else "-" for v in videos) + + +def test_hdr10plus_at_both_ends_relabels_the_whole_ladder(monkeypatch: pytest.MonkeyPatch) -> None: + videos, probed = ladder(monkeypatch, "++---++") + assert ranges(videos) == "+++++++" + assert len(probed) == 4 + + +def test_a_mixed_sample_probes_every_track(monkeypatch: pytest.MonkeyPatch) -> None: + videos, probed = ladder(monkeypatch, "+-+-+-+-") + assert ranges(videos) == "+-+-+-+-" + assert len(probed) == 8 + + +def test_plain_hdr10_ends_stop_after_the_sample(monkeypatch: pytest.MonkeyPatch) -> None: + videos, probed = ladder(monkeypatch, "--++--") + assert ranges(videos) == "------" + assert len(probed) == 4 + + +def test_two_encodes_at_one_resolution_are_all_probed(monkeypatch: pytest.MonkeyPatch) -> None: + videos, probed = ladder(monkeypatch, "++-++", heights=[2160, 1440, 1080, 1080, 720]) + assert ranges(videos) == "++-++" + assert len(probed) == 5 + + +def test_on_start_runs_only_when_there_is_a_track_to_probe(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(hdr10plus, "first_fragment", lambda video, session: None) + started: list[int] = [] + sdr = SimpleNamespace(id="s", range=Video.Range.SDR, codec=Video.Codec.AVC, bitrate=1, session=None) + for videos in ([sdr], [hevc("v", 2160)]): + hdr10plus.probe_hdr10plus( + videos, # type: ignore[arg-type] + None, # type: ignore[arg-type] + log=logging.getLogger("test"), + workers=4, + on_start=lambda: started.append(1), + ) + assert started == [1] + + +def test_first_fragment_skips_the_header_boxes_of_a_one_file_track(monkeypatch: pytest.MonkeyPatch) -> None: + media = fragment(nal(32, b"vps"), sei(4, HDR10PLUS_T35 + b"\x00" * 8), nal(1, b"\xff" * 64)) + data = box(b"ftyp", b"isom") + box(b"moov", b"\x00" * 100_000) + box(b"sidx", b"\x00" * 40) + media + + def fake_get(url: str, session: Any, byte_range: str) -> bytes: + start, end = (int(x) for x in byte_range.removeprefix("bytes=").split("-")) + return data[start : end + 1] + + monkeypatch.setattr(hdr10plus, "first_segment", lambda track, session: ("https://cdn/v.mp4", 0)) + monkeypatch.setattr(hdr10plus, "ranged_get", fake_get) + assert fragment_has_hdr10plus(hdr10plus.first_fragment(hevc("v", 2160), None) or b"") is True # type: ignore[arg-type] + + +def test_first_segment_of_an_hls_playlist_keeps_its_byte_offset() -> None: + playlist = '#EXTM3U\n#EXT-X-MAP:URI="init.mp4"\n#EXTINF:6.0,\n#EXT-X-BYTERANGE:5000@1200\nmedia.mp4\n' + session = SimpleNamespace(get=lambda url, timeout: SimpleNamespace(status_code=200, text=playlist)) + track = SimpleNamespace(url="https://cdn/v/index.m3u8", data={}, descriptor=Video.Descriptor.HLS) + assert hdr10plus.first_segment(track, session) == ("https://cdn/v/media.mp4", 1200) # type: ignore[arg-type] + + +def probe_calls(monkeypatch: pytest.MonkeyPatch, service: Any, **kwargs: Any) -> list[int]: + calls: list[int] = [] + monkeypatch.setattr( + hdr10plus, "probe_hdr10plus", lambda videos, session, *, log, workers, on_start: calls.append(workers) + ) + hdr10plus.probed_tracks(service, None, **kwargs) + return calls + + +def test_probed_tracks_follows_the_config_switch_and_the_workers_count(monkeypatch: pytest.MonkeyPatch) -> None: + service = SimpleNamespace(get_tracks=lambda title: SimpleNamespace(videos=[hevc("v", 2160)]), session=None) + assert probe_calls(monkeypatch, service, workers=3) == [3] + monkeypatch.setattr(config, "hdr10plus_probe", False) + assert probe_calls(monkeypatch, service, workers=3) == [] + + +def test_probed_tracks_leaves_a_remote_service_to_its_server(monkeypatch: pytest.MonkeyPatch) -> None: + remote = RemoteService.__new__(RemoteService) + remote.get_tracks = lambda title: SimpleNamespace(videos=[hevc("v", 2160)]) # type: ignore[method-assign] + assert probe_calls(monkeypatch, remote, workers=3) == [] + + +@pytest.mark.parametrize("workers", [-1, "8"]) +def test_a_bad_workers_count_is_an_error(monkeypatch: pytest.MonkeyPatch, workers: Any) -> None: + service = SimpleNamespace(get_tracks=lambda title: SimpleNamespace(videos=[hevc("v", 2160)]), session=None) + with pytest.raises(ValueError, match="workers"): + probe_calls(monkeypatch, service, workers=workers) + + +def test_a_quoted_hdr10plus_probe_value_is_an_error() -> None: + with pytest.raises(SystemExit, match="hdr10plus_probe"): + Config(hdr10plus_probe="false") + + +def test_a_fragment_inside_the_first_window_needs_one_request(monkeypatch: pytest.MonkeyPatch) -> None: + media = fragment(nal(32, b"vps"), sei(4, HDR10PLUS_T35 + b"\x00" * 8), nal(1, b"\xff" * 64)) + data = box(b"sidx", b"\x00" * 40) + media + requests: list[str] = [] + + def fake_get(url: str, session: Any, byte_range: str) -> bytes: + requests.append(byte_range) + return data + + monkeypatch.setattr(hdr10plus, "first_segment", lambda track, session: ("https://cdn/v.mp4", 0)) + monkeypatch.setattr(hdr10plus, "ranged_get", fake_get) + assert fragment_has_hdr10plus(hdr10plus.first_fragment(hevc("v", 2160), None) or b"") is True # type: ignore[arg-type] + assert len(requests) == 1 + + +def test_a_quoted_dl_workers_value_works_like_the_cli(monkeypatch: pytest.MonkeyPatch) -> None: + service = SimpleNamespace(get_tracks=lambda title: SimpleNamespace(videos=[hevc("v", 2160)]), session=None) + monkeypatch.setitem(config.dl, "workers", "8") + assert probe_calls(monkeypatch, service) == [8] diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 226792e9..2e46b0e5 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -117,6 +117,7 @@ from unshackle.core.utils.click_types import ( MultipleChoice, ) from unshackle.core.utils.collections import ci_get, merge_dict +from unshackle.core.utils.hdr10plus import probed_tracks from unshackle.core.utils.post_scripts import NO_POST_SCRIPTS, build_context, dispatch, season_context from unshackle.core.utils.redact import mask_proxy from unshackle.core.utils.selector import select_multiple @@ -369,9 +370,11 @@ def load_tracks( name: Optional[str] = None, year: Optional[int] = None, language: Optional[Language] = None, + workers: Optional[int] = None, + on_probe: Optional[Callable[[], None]] = None, ) -> None: """Load the tracks and chapters of a title, then apply the --enrich name, year and language again.""" - title.tracks.add(service.get_tracks(title), warn_only=True) + title.tracks.add(probed_tracks(service, title, workers=workers, on_probe=on_probe), warn_only=True) title.tracks.chapters = service.get_chapters(title) apply_enrich(title, name, year, language) @@ -2370,9 +2373,18 @@ class dl: events.subscribe(events.Types.TRACK_MULTIPLEX, service.on_track_multiplex) tracks_label = "Getting Remote Tracks..." if self.is_remote else "Getting Tracks..." - with console.status(tracks_label, spinner="dots"): + tracks_spinner = Spinner("dots", text=tracks_label) + with SyncLive(Padding(tracks_spinner, (0, 5)), console=console, refresh_per_second=12.5, transient=True): try: - load_tracks(service, title, name=enrich_title, year=enrich_year, language=enrich_lang) + load_tracks( + service, + title, + name=enrich_title, + year=enrich_year, + language=enrich_lang, + workers=workers, + on_probe=lambda: tracks_spinner.update(text="Probing HDR10 tracks for HDR10+..."), + ) except Exception as e: if self.debug_logger: self.debug_logger.log_error( diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 1133e524..40e4c3ad 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -43,6 +43,7 @@ from unshackle.core.tracks import Audio, Subtitle, Tracks, Video from unshackle.core.utilities import declared_kwargs, kind_only_conflict, original_tag from unshackle.core.utils.click_types import AUDIO_CODEC_LIST, SUBTITLE_CODEC, VIDEO_CODEC_LIST from unshackle.core.utils.collections import ci_get +from unshackle.core.utils.hdr10plus import probed_tracks from unshackle.core.utils.redact import REDACTED, URL_USERINFO_RE, redact_all, redact_secrets, redact_text log = logging.getLogger("api") @@ -2034,7 +2035,7 @@ async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Reques for title in sorted_titles: try: - tracks = await call(service_instance.get_tracks, title) + tracks = await call(probed_tracks, service_instance, title) video_tracks = sorted(tracks.videos, key=lambda t: t.bitrate or 0, reverse=True) audio_tracks = sorted(tracks.audio, key=lambda t: t.bitrate or 0, reverse=True) @@ -2081,7 +2082,7 @@ async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Reques else: first_title = titles - tracks = await call(service_instance.get_tracks, first_title) + tracks = await call(probed_tracks, service_instance, first_title) video_tracks = sorted(tracks.videos, key=lambda t: t.bitrate or 0, reverse=True) audio_tracks = sorted(tracks.audio, key=lambda t: t.bitrate or 0, reverse=True) @@ -4091,7 +4092,7 @@ async def session_tracks_handler( service_instance = session.service_instance async with session.lock: try: - tracks = await asyncio.to_thread(service_instance.get_tracks, title) + tracks = await asyncio.to_thread(probed_tracks, service_instance, title) finally: session.touch() diff --git a/unshackle/core/config.py b/unshackle/core/config.py index fc96173c..2bd70430 100644 --- a/unshackle/core/config.py +++ b/unshackle/core/config.py @@ -200,6 +200,9 @@ class Config: self.tag_rules: list = kwargs.get("tag_rules") or [] self.dual_multi_mode: str = (kwargs.get("dual_multi_mode") or "strict").lower() self.generate_dv_mode: str = str(kwargs.get("generate_dv_mode") or "v29").lower() + self.hdr10plus_probe: bool = kwargs.get("hdr10plus_probe", True) + if not isinstance(self.hdr10plus_probe, bool): + raise SystemExit(f"ERROR: hdr10plus_probe must be true or false, not {self.hdr10plus_probe!r}") self.output_template: dict = kwargs.get("output_template") or {} folder_cfg = self.output_template.pop("folder", "") self.folder_template: str = "" diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 54b2df50..0d503a22 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -250,8 +250,8 @@ class HLS: ) # DV-composite track: primary codec is plain HEVC but SUPPLEMENTAL-CODECS advertises # a DV codec. Range stays whatever VIDEO-RANGE signaled (HDR10/HLG/SDR); DVFixup will - # restore DV signaling post-download. Services that know their encoder embeds HDR10+ - # SEI must override `range` themselves. + # restore DV signaling post-download. HDR10+ SEI in an HDR10 track is found later by the + # core HDR10+ probe. dv_compatible_bitstream = primary_track_type is Video and not primary_has_dv and supp_dv_codec is not None tracks.add( diff --git a/unshackle/core/utils/hdr10plus.py b/unshackle/core/utils/hdr10plus.py new file mode 100644 index 00000000..991fedd6 --- /dev/null +++ b/unshackle/core/utils/hdr10plus.py @@ -0,0 +1,223 @@ +from __future__ import annotations + +import logging +import struct +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Callable, Optional, Union +from urllib.parse import urljoin + +from requests import Session + +from unshackle.core.manifests.ism_init import find_box, iter_boxes, remove_emulation_prevention +from unshackle.core.session import RnetSession, utf8_text +from unshackle.core.utils.bitrate import PROBE_TIMEOUT, ranged_get + +if TYPE_CHECKING: + from unshackle.core.service import Service + from unshackle.core.titles import Title_T + from unshackle.core.tracks import Tracks, Video + +PROBE_BYTES = 64 * 1024 +HDR10PLUS_T35 = b"\xb5\x00\x3c\x00\x01\x04" + + +def sei_number(rbsp: bytes, i: int) -> tuple[Optional[int], int]: + """Read the SEI payload type or size that starts at ``i``: 0xFF bytes add 255 each. Returns (value, next).""" + value = 0 + while i < len(rbsp) and rbsp[i] == 0xFF: + value += 255 + i += 1 + if i >= len(rbsp): + return None, i + return value + rbsp[i], i + 1 + + +def sei_has_hdr10plus(nal: bytes) -> bool: + """Whether an HEVC SEI NAL holds an ST 2094-40 message.""" + rbsp = remove_emulation_prevention(nal[2:]) + i = 0 + while i < len(rbsp) and rbsp[i] != 0x80: + payload_type, i = sei_number(rbsp, i) + payload_size, i = sei_number(rbsp, i) + if payload_type is None or payload_size is None: + return False + if payload_type == 4 and rbsp[i : i + len(HDR10PLUS_T35)] == HDR10PLUS_T35: + return True + i += payload_size + return False + + +def fragment_has_hdr10plus(data: bytes) -> Optional[bool]: + """Whether the first sample of an fMP4 fragment carries HDR10+ SEI, or None when the bytes cannot tell. + + The walk ends at the first VCL NAL, because a prefix SEI comes before it. CENC and cbcs keep + non-VCL NALs and every NAL length in the clear, so this needs no content key. A sample that is + encrypted whole can read as False as well as None; both leave the label alone. + """ + # 4-byte NAL lengths: reading lengthSizeMinusOne from hvcC would cost an init-segment fetch + mdat = find_box(data, 0, len(data), b"mdat") + if not mdat: + return None + pos = mdat[0] + while pos + 6 <= len(data): + size = struct.unpack(">I", data[pos : pos + 4])[0] + nal_type = (data[pos + 4] >> 1) & 0x3F + if size < 2 or nal_type >= 48: + return None # not a NAL: the sample is fully encrypted or the walk lost sync + if nal_type < 32: + return False + if nal_type == 39: + if pos + 4 + size > len(data): + return None + if sei_has_hdr10plus(data[pos + 4 : pos + 4 + size]): + return True + pos += 4 + size + return None + + +def first_segment(track: "Video", session: Union[Session, RnetSession]) -> Optional[tuple[str, int]]: + """The URL and start offset of the track's first media segment, looked up without a full parse.""" + from unshackle.core.manifests import DASH + from unshackle.core.tracks.track import Track + + url = track.url if isinstance(track.url, str) else track.url[0] + dash = track.data.get("dash") or {} + # a service can collapse a DASH representation to a one-file URL track and keep the manifest + if dash.get("manifest") is not None and dash.get("period") is not None: + _, segments, *_ = DASH.get_period_segments( + period=dash["period"], + adaptation_set=dash["adaptation_set"], + representation=dash["representation"], + manifest=dash["manifest"], + track=track, + track_url=url, + session=session, + probe_kid=False, + ) + if not segments: + return None + seg_url, byte_range = segments[0] + return seg_url, int((byte_range or "").partition("-")[0] or 0) + if track.descriptor == Track.Descriptor.HLS: + res = session.get(url, timeout=PROBE_TIMEOUT) + if getattr(res, "status_code", 0) != 200: + return None + offset = 0 + for line in utf8_text(res).splitlines(): + line = line.strip() + if line.startswith("#EXT-X-BYTERANGE:") and "@" in line: + offset = int(line.rpartition("@")[2]) + elif line and not line.startswith("#"): + return urljoin(url, line), offset + return None + if track.descriptor == Track.Descriptor.URL: + return url, 0 + return None + + +def first_fragment(track: "Video", session: Union[Session, RnetSession]) -> Optional[bytes]: + """The first PROBE_BYTES of the track's first media fragment.""" + segment = first_segment(track, session) + if not segment: + return None + url, start = segment + # A one-file track opens with ftyp, moov and sidx, which can outgrow PROBE_BYTES. Fetch again past + # them only when the fragment does not start inside this window. + for _ in range(4): + data = ranged_get(url, session, f"bytes={start}-{start + PROBE_BYTES - 1}") + if not data: + return None + skip = 0 + for kind, _, _, end in iter_boxes(data, 0, len(data)): + if kind in (b"styp", b"moof", b"mdat"): + return data + skip = end + if not skip: + return data + start += skip + return None + + +def probe_hdr10plus( + videos: list["Video"], + session: Union[Session, RnetSession], + *, + log: logging.Logger, + workers: int, + on_start: Optional[Callable[[], None]] = None, +) -> None: + """Relabel HEVC HDR10 tracks whose first sample carries HDR10+ SEI as HDR10P. + + A ladder with one track per resolution is sampled at both ends first. Two tracks at one resolution + can be an HDR10 and an HDR10+ encode side by side, so such a ladder is probed whole. The rules are + in the ``hdr10plus_probe`` section of docs/reference/configuration/download.md. A failed probe never + changes a label. ``workers`` caps the probes in flight, and ``on_start`` runs once before the first. + """ + from unshackle.core.tracks import Video + + candidates = sorted( + (v for v in videos if v.range == Video.Range.HDR10 and v.codec == Video.Codec.HEVC), + key=lambda v: v.bitrate or 0, + reverse=True, + ) + if candidates and on_start: + on_start() + if len(candidates) <= 4 or len({(v.width, v.height) for v in candidates}) < len(candidates): + samples, rest = candidates, [] + else: + samples, rest = candidates[:2] + candidates[-2:], candidates[2:-2] + + def probe(video: "Video") -> Optional[bool]: + try: + data = first_fragment(video, video.session or session) + found = fragment_has_hdr10plus(data) if data else None + except Exception as e: + log.debug(f"{video.id}: HDR10+ probe failed ({e})") + return None + log.debug(f"{video.id}: HDR10+ probe {found}") + return found + + def probe_all(tracks: list["Video"]) -> list[Optional[bool]]: + if not tracks: + return [] + with ThreadPoolExecutor(max_workers=min(workers, len(tracks))) as executor: + return list(executor.map(probe, tracks)) + + results = probe_all(samples) + if any(results) and not all(results): + samples, results = samples + rest, results + probe_all(rest) + hits = candidates if all(results) else [v for v, found in zip(samples, results) if found] + for video in hits: + video.range = Video.Range.HDR10P + + +def probed_tracks( + service: "Service", + title: "Title_T", + *, + workers: Optional[int] = None, + on_probe: Optional[Callable[[], None]] = None, +) -> "Tracks": + """The title's tracks from ``service.get_tracks``, with HDR10+ found in tracks labelled HDR10. + + ``workers`` defaults to ``dl.workers`` from the config. A remote service's tracks come back from a + server that already probed them. + """ + from unshackle.core.config import config + from unshackle.core.downloaders import default_max_workers + from unshackle.core.remote_service import RemoteService + + tracks = service.get_tracks(title) + if config.hdr10plus_probe and tracks.videos and not isinstance(service, RemoteService): + # int() as click does for --workers, so a quoted dl.workers works on the REST path too + workers = workers or int(config.dl.get("workers") or 0) or default_max_workers() + if not isinstance(workers, int) or isinstance(workers, bool) or workers < 1: + raise ValueError(f"workers must be a whole number of 1 or more, not {workers!r}") + probe_hdr10plus( + tracks.videos, + service.session, + log=logging.getLogger("hdr10plus"), + workers=workers, + on_start=on_probe, + ) + return tracks diff --git a/unshackle/services/EXAMPLE/__init__.py b/unshackle/services/EXAMPLE/__init__.py index 22f2f63c..086e23e9 100644 --- a/unshackle/services/EXAMPLE/__init__.py +++ b/unshackle/services/EXAMPLE/__init__.py @@ -376,11 +376,6 @@ class EXAMPLE(Service): # The manifest can't always be trusted for range - stamp what we asked for. video.range = range_ - # HDR10+ is a BITSTREAM feature the manifest often labels as plain HDR10. - # If this service is known to ship HDR10+ SEI, flip it so mediainfo agrees. - if range_ == Video.Range.HDR10P: - video.range = Video.Range.HDR10P - # DV-composite: a stream carrying DV RPU NALs inside a container that only # signals HEVC. Setting this flag makes DVFixup round-trip the bitstream # through dovi_tool so the muxed MKV is recognised as Dolby Vision.