From 355df2e54fdbe3300b2451d12d8736537604891c Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Thu, 8 Oct 2026 10:38:49 -0600 Subject: [PATCH] fix(import): carry the HLS segment filter in an export An import kept the segments a service drops and failed on their key. A licence-only run now also licenses a key that starts part-way. --- .../test_hls_licence_only_key_rotation.py | 186 ++++++++++++++++++ tests/core/test_import_hls_ism.py | 21 ++ unshackle/commands/dl.py | 7 +- unshackle/core/import_service.py | 12 +- unshackle/core/manifests/hls.py | 33 +++- unshackle/core/remote_service.py | 29 ++- 6 files changed, 267 insertions(+), 21 deletions(-) create mode 100644 tests/core/test_hls_licence_only_key_rotation.py diff --git a/tests/core/test_hls_licence_only_key_rotation.py b/tests/core/test_hls_licence_only_key_rotation.py new file mode 100644 index 00000000..503cc2ae --- /dev/null +++ b/tests/core/test_hls_licence_only_key_rotation.py @@ -0,0 +1,186 @@ +"""Regression test for a licence-only HLS run on a playlist that changes its key part-way. + +``HLS.download_track`` licensed only the first key before it returned, so an export made +with no download held no content key for the later KID, and the import failed there. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from uuid import UUID + +from pywidevine.pssh import PSSH + +from unshackle.core.constants import DOWNLOAD_LICENCE_ONLY +from unshackle.core.manifests.hls import HLS +from unshackle.core.remote_service import segment_uri_filter +from unshackle.core.tracks import Video +from unshackle.core.tracks.track import DownloadContext, Track + +FIRST_KID = UUID("11111111-1111-1111-1111-111111111111") +LATER_KID = UUID("22222222-2222-2222-2222-222222222222") + + +def key_tag(kid: UUID) -> str: + pssh = PSSH.new(key_ids=[kid], system_id=PSSH.SystemId.Widevine).dumps() + return ( + '#EXT-X-KEY:METHOD=SAMPLE-AES,URI="data:text/plain;base64,' + f'{pssh}",KEYFORMAT="urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed",KEYFORMATVERSIONS="1"' + ) + + +PLAYLIST = f"""#EXTM3U +#EXT-X-VERSION:6 +#EXT-X-TARGETDURATION:6 +{key_tag(FIRST_KID)} +#EXTINF:6.0, +seg0.m4s +#EXT-X-DISCONTINUITY +{key_tag(LATER_KID)} +#EXTINF:6.0, +seg1.m4s +#EXT-X-ENDLIST +""" + + +def test_licence_only_licenses_a_key_that_starts_part_way(tmp_path: Path, monkeypatch: Any) -> None: + playlist = tmp_path / "media.m3u8" + playlist.write_text(PLAYLIST, encoding="utf-8") + track = Video( + url="https://example.test/media.m3u8", + language="en", + descriptor=Track.Descriptor.HLS, + codec=Video.Codec.HEVC, + ) + track.from_file = playlist + # the KID normally comes from the init segment; this test is offline + monkeypatch.setattr(HLS, "get_track_kid_from_init", staticmethod(lambda *args, **kwargs: None)) + + licensed: list[UUID] = [] + ctx = DownloadContext( + save_path=tmp_path / "video.mp4", + save_dir=tmp_path / "video_segments", + progress=lambda **_: None, + license_widevine=lambda drm, track_kid=None: licensed.append(track_kid), + cdm=None, + ) + + DOWNLOAD_LICENCE_ONLY.set() + try: + HLS.download_track(track=track, ctx=ctx) + finally: + DOWNLOAD_LICENCE_ONLY.clear() + + assert licensed == [FIRST_KID, LATER_KID] + assert track.data["hls"]["unwanted_segments"] == [] + + +def test_licence_only_skips_the_key_of_a_filtered_segment(tmp_path: Path, monkeypatch: Any) -> None: + """The filtered segment URIs go to the track data, and an import filter built from them drops the same segment.""" + playlist = tmp_path / "media.m3u8" + playlist.write_text(PLAYLIST, encoding="utf-8") + monkeypatch.setattr(HLS, "get_track_kid_from_init", staticmethod(lambda *args, **kwargs: None)) + + def licence_only(segment_filter: Any) -> tuple[list[UUID], list[str]]: + track = Video( + url="https://example.test/media.m3u8", + language="en", + descriptor=Track.Descriptor.HLS, + codec=Video.Codec.HEVC, + ) + track.from_file = playlist + track.OnSegmentFilter = segment_filter + licensed: list[UUID] = [] + ctx = DownloadContext( + save_path=tmp_path / "video.mp4", + save_dir=tmp_path / "video_segments", + progress=lambda **_: None, + license_widevine=lambda drm, track_kid=None: licensed.append(track_kid), + cdm=None, + ) + DOWNLOAD_LICENCE_ONLY.set() + try: + HLS.download_track(track=track, ctx=ctx) + finally: + DOWNLOAD_LICENCE_ONLY.clear() + return licensed, track.data["hls"]["unwanted_segments"] + + licensed, unwanted = licence_only(lambda segment: segment.uri == "seg1.m4s") + assert licensed == [FIRST_KID] + assert len(unwanted) == 1 and unwanted[0].endswith("seg1.m4s") + + assert licence_only(segment_uri_filter(unwanted)) == (licensed, unwanted) + + +def test_licence_only_warns_when_no_dropped_segment_is_in_the_playlist( + tmp_path: Path, monkeypatch: Any, caplog: Any +) -> None: + """A changed host or path makes an exported list match nothing, and the download then keeps every segment.""" + playlist = tmp_path / "media.m3u8" + playlist.write_text(PLAYLIST, encoding="utf-8") + monkeypatch.setattr(HLS, "get_track_kid_from_init", staticmethod(lambda *args, **kwargs: None)) + track = Video( + url="https://example.test/media.m3u8", + language="en", + descriptor=Track.Descriptor.HLS, + codec=Video.Codec.HEVC, + ) + track.from_file = playlist + stale = ["https://other.example/moved.m4s"] + track.OnSegmentFilter = segment_uri_filter(stale) + track.data["hls"]["unwanted_segments"] = stale + ctx = DownloadContext( + save_path=tmp_path / "video.mp4", + save_dir=tmp_path / "video_segments", + progress=lambda **_: None, + license_widevine=lambda drm, track_kid=None: None, + cdm=None, + ) + + DOWNLOAD_LICENCE_ONLY.set() + try: + with caplog.at_level("WARNING", logger="HLS"): + HLS.download_track(track=track, ctx=ctx) + finally: + DOWNLOAD_LICENCE_ONLY.clear() + + assert "keeps every segment" in caplog.text + + +def test_licence_only_fetches_no_aes_128_key(tmp_path: Path, monkeypatch: Any) -> None: + """An AES-128 key with an IV on each segment is a new key each time, and licence-only must not fetch them.""" + playlist = tmp_path / "media.m3u8" + playlist.write_text( + "#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-TARGETDURATION:6\n" + + "".join( + f'#EXT-X-KEY:METHOD=AES-128,URI="https://example.test/key",IV=0x{i:032x}\n#EXTINF:6.0,\nseg{i}.ts\n' + for i in range(1, 4) + ) + + "#EXT-X-ENDLIST\n", + encoding="utf-8", + ) + track = Video( + url="https://example.test/media.m3u8", + language="en", + descriptor=Track.Descriptor.HLS, + codec=Video.Codec.AVC, + ) + track.from_file = playlist + fetched: list[Any] = [] + monkeypatch.setattr(HLS, "get_drm", staticmethod(lambda key, session=None: fetched.append(key))) + ctx = DownloadContext( + save_path=tmp_path / "video.mp4", + save_dir=tmp_path / "video_segments", + progress=lambda **_: None, + license_widevine=lambda drm, track_kid=None: None, + cdm=None, + ) + + DOWNLOAD_LICENCE_ONLY.set() + try: + HLS.download_track(track=track, ctx=ctx) + finally: + DOWNLOAD_LICENCE_ONLY.clear() + + assert len(fetched) <= 1, "only the first key of the playlist, as before" diff --git a/tests/core/test_import_hls_ism.py b/tests/core/test_import_hls_ism.py index e919be75..ce3f5b0a 100644 --- a/tests/core/test_import_hls_ism.py +++ b/tests/core/test_import_hls_ism.py @@ -181,3 +181,24 @@ def test_hls_tracks_tolerate_the_empty_manifest_data(tmp_path: Path) -> None: assert not track.data.get("hls"), "to_dict cannot carry the m3u8 objects, and does not need to" track.data["hls"]["segment_durations"] = [1, 2, 3] # hls.py, in download_track assert track.data["hls"]["segment_durations"] == [1, 2, 3] + + +def test_hls_round_trip_keeps_the_dropped_segments(tmp_path: Path) -> None: + """An import has no service to run the segment filter, so the export carries the dropped URIs.""" + original = HLS.from_text(HLS_MANIFEST, HLS_URL).to_tracks(language="en") + video = original.videos[0] + dropped = "https://cdn.example/bumper_1.m4s?token=a" + video.data["hls"]["unwanted_segments"] = [dropped] # hls.py, in download_track + + export = export_tracks(tmp_path, original, HLS_URL) + _, tracks = import_tracks(export) + + rebuilt = next(t for t in tracks.videos if t.id == video.id) + assert rebuilt.OnSegmentFilter(SimpleNamespace(absolute_uri="https://cdn.example/bumper_1.m4s?token=b")) + assert not rebuilt.OnSegmentFilter(SimpleNamespace(absolute_uri="https://cdn.example/main_1.m4s")) + assert rebuilt.data["hls"]["unwanted_segments"] == [dropped] + assert all(t.OnSegmentFilter is None for t in tracks if t.id != video.id) + + video.data["hls"]["unwanted_segments"] = [] + _, tracks = import_tracks(export_tracks(tmp_path, original, HLS_URL)) + assert next(t for t in tracks.videos if t.id == video.id).OnSegmentFilter is None, "a stale list must go" diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 4cd959cc..05c21734 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -4064,7 +4064,12 @@ class dl: doc.add(entry) tracks_map = entry.ext("unshackle").setdefault("tracks", {}) - tracks_map.setdefault(str(track.id), track.to_dict()) + track_dict = tracks_map.setdefault(str(track.id), track.to_dict()) + unwanted = (track.data.get("hls") or {}).get("unwanted_segments") + if unwanted: + track_dict["unwanted_segments"] = unwanted + elif unwanted is not None: + track_dict.pop("unwanted_segments", None) # the row names the track's KIDs, so an import gives the track only their keys track_kids = [kid.hex for kid in own_kids(drm)] if drm is not None else [] for row in entry.tracks: diff --git a/unshackle/core/import_service.py b/unshackle/core/import_service.py index efba9b18..3a57052d 100644 --- a/unshackle/core/import_service.py +++ b/unshackle/core/import_service.py @@ -18,7 +18,13 @@ from unshackle.core.constants import AnyTrack from unshackle.core.credential import Credential from unshackle.core.drm import drm_from_dict, real_kids from unshackle.core.manifests import DASH, HLS, ISM -from unshackle.core.remote_service import RemoteService, build_title, match_track, resolve_proxy_arg +from unshackle.core.remote_service import ( + RemoteService, + build_title, + match_track, + resolve_proxy_arg, + segment_uri_filter, +) 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 @@ -409,6 +415,10 @@ class ImportService: drm = self.rebuild_drm(track_dict, title_id) if drm: track.drm = drm + if track_dict.get("unwanted_segments"): + track.OnSegmentFilter = segment_uri_filter(track_dict["unwanted_segments"]) + # HLS.download_track warns when none of these is in the playlist + track.data["hls"]["unwanted_segments"] = list(track_dict["unwanted_segments"]) tracks.add(track, warn_only=True) if manifest_url and parser is None and not [*tracks.videos, *tracks.audio, *tracks.subtitles]: diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 118fb6d5..2d2e1ea7 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -674,6 +674,17 @@ class HLS: log.error("Track's HLS playlist has no segments, expecting an invariant M3U8 playlist.") sys.exit(1) + unwanted_segments = [ + segment for segment in master.segments if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment) + ] + hls_data = track.data.setdefault("hls", {}) + if hls_data.get("unwanted_segments") and not unwanted_segments: + log.warning( + f"None of the {len(hls_data['unwanted_segments'])} segments that the export drops is in the " + f"playlist of {track.id}. The download keeps every segment." + ) + hls_data["unwanted_segments"] = [s.absolute_uri for s in unwanted_segments] + # Get session DRM as fallback but prefer media playlist keys for accurate KID matching if track.drm: session_drm = track.get_drm_for_cdm(cdm) @@ -745,13 +756,21 @@ class HLS: raise if DOWNLOAD_LICENCE_ONLY.is_set(): + # a playlist can change its key part-way, and only the download loop licenses the later keys + licensed_keys = [initial_drm_key] + for segment in master.segments: + if not segment.keys or segment in unwanted_segments: + continue + key = HLS.resolve_segment_key(segment.keys, cdm, track.drm_preference) + if key is None or key.method == "AES-128" or key in licensed_keys: + continue + licensed_keys.append(key) + later_drm = HLS.get_drm(key, session) + if isinstance(later_drm, (Widevine, PlayReady)) and license_widevine: + license_widevine(later_drm, track_kid=later_drm.kid) progress(downloaded="[yellow]SKIPPED") return - unwanted_segments = [ - segment for segment in master.segments if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment) - ] - # Downloaded segment files are named by post-filter index; map that back to the wanted # segment so the IV uses the absolute media sequence number, not the download index. wanted_segments = [seg for seg in master.segments if seg not in unwanted_segments] @@ -1152,9 +1171,6 @@ class HLS: drm.content_keys.setdefault(prev_kid, prev_key) encryption_data = (key, drm) - if DOWNLOAD_LICENCE_ONLY.is_set(): - continue - if is_last_segment: # required as it won't end with EXT-X-DISCONTINUITY nor a new key if encryption_data: @@ -1165,9 +1181,6 @@ class HLS: progress(advance=1) - if DOWNLOAD_LICENCE_ONLY.is_set(): - return - def find_segments_recursively(directory: Path) -> list[Path]: """Find all segment files recursively in any directory structure created by downloaders.""" segments = [] diff --git a/unshackle/core/remote_service.py b/unshackle/core/remote_service.py index d50fe5f9..42c655c6 100644 --- a/unshackle/core/remote_service.py +++ b/unshackle/core/remote_service.py @@ -22,7 +22,7 @@ from enum import Enum from http.cookiejar import CookieJar from pathlib import Path from threading import Event, Lock, Thread -from typing import Any, Callable, Dict, Iterator, Optional, Union +from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Union from uuid import UUID import click @@ -465,6 +465,22 @@ def resolve_manifest_data(tracks: Tracks, manifests: list) -> None: remote_track.drm = matched.drm +def segment_uri_filter(uris: Iterable[Any]) -> Callable[[Any], bool]: + """Return an HLS ``OnSegmentFilter`` that drops the segments with these URIs. + + A URI also matches without its query string, because a later fetch of the playlist can + sign the same segment with a new one. + """ + exact = {str(u) for u in uris} + paths = {u.split("?")[0] for u in exact} + + def segment_filter(segment: Any) -> bool: + uri = str(getattr(segment, "absolute_uri", "") or "") + return bool(exact) and (uri in exact or uri.split("?")[0] in paths) + + return segment_filter + + def same_bitrate(local: Optional[int], remote: Optional[int]) -> bool: """Equal to the nearest kb/s. A server that predates ``bitrate_bps`` rounds to kb/s in transit.""" if not local or not remote: @@ -895,7 +911,7 @@ class RemoteService: self.client_licensed: set[str] = set() # track ids this machine's device licenses by challenge relay self.server_vault_keys: Dict[UUID, str] = {} self.server_vault = ServerVault(self) - self._segment_filters: Dict[str, tuple[set[str], set[str]]] = {} + self._segment_filters: Dict[str, Callable[[Any], bool]] = {} self._log_seq = 0 self._log_drain_lock = Lock() self._keepalive_stop = Event() @@ -1407,14 +1423,9 @@ class RemoteService: result = self.client.post_optional( f"/api/session/{self._session_id}/segment_filter", {"track_id": track_id} ) - uris = [str(u) for u in (result.get("unwanted") or [])] - unwanted = (set(uris), {u.split("?")[0] for u in uris}) + unwanted = segment_uri_filter(result.get("unwanted") or []) self._segment_filters[track_id] = unwanted - exact, paths = unwanted - if not exact: - return False - uri = str(getattr(segment, "absolute_uri", "") or "") - return uri in exact or uri.split("?")[0] in paths + return unwanted(segment) return segment_filter