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.

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"

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"

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:

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]:

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 = []

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