fix(import): scope the merge-back and detect ISM protection

Only merge back tracks from another manifest, so a drifted id no longer duplicates a rung or wipes the main manifest's own subtitles. track_is_encrypted now recognises an ISM ProtectionHeader.
2026-11-12 20:15:10 +00:00 · 2026-07-29 10:48:34 -06:00
parent cebfd7bc80
commit 42d7fd8303
3 changed files with 438 additions and 10 deletions
@@ -0,0 +1,184 @@
+"""Export/import round trips for the two manifest types that are not DASH.
+
+``get_tracks`` splits on manifest type, and the two halves need opposite things. ISM takes the
+re-parse path and needs ``track.data["ism"]["segments"]`` restored, because
+``ISM.download_track`` reads it. HLS takes the rebuild-from-dicts path and needs nothing
+restored, because ``HLS.download_track`` re-fetches the variant from ``track.url``. That makes
+the stored per-track URL load-bearing for HLS in a way it is not for ISM.
+
+These drive the real ``dl.write_export`` rather than hand-written export dicts, so a change to
+the export format fails here instead of silently invalidating the fixtures.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from unshackle.commands.dl import dl
+from unshackle.core.import_service import ImportService
+from unshackle.core.manifests import HLS, ISM
+
+ISM_URL = "https://cdn.example/x.ism/manifest"
+ISM_MANIFEST = """<?xml version="1.0"?>
+<SmoothStreamingMedia MajorVersion="2" MinorVersion="0" Duration="60000000" TimeScale="10000000">
+  <StreamIndex Type="video" Name="video" Language="en"
+               Url="QualityLevels({Bitrate})/Fragments(video={start_time})">
+    <QualityLevel Index="0" Bitrate="3000000" FourCC="H264" MaxWidth="1280" MaxHeight="720"
+                  CodecPrivateData="00000001674D401F"/>
+    <c t="0" d="20000000" r="3"/>
+  </StreamIndex>
+  <StreamIndex Type="audio" Name="audio" Language="en"
+               Url="QualityLevels({Bitrate})/Fragments(audio={start_time})">
+    <QualityLevel Index="0" Bitrate="128000" AudioTag="255" SamplingRate="48000"
+                  Channels="2" BitsPerSample="16" CodecPrivateData="1190"/>
+    <c t="0" d="20000000" r="3"/>
+  </StreamIndex>
+</SmoothStreamingMedia>
+"""
+
+HLS_URL = "https://cdn.example/master.m3u8"
+HLS_MANIFEST = """#EXTM3U
+#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID="aud",LANGUAGE="en",NAME="English",DEFAULT=YES,URI="audio_en.m3u8"
+#EXT-X-MEDIA:TYPE=SUBTITLES,GROUP-ID="sub",LANGUAGE="en",NAME="English",URI="sub_en.m3u8"
+#EXT-X-STREAM-INF:BANDWIDTH=3000000,RESOLUTION=1280x720,CODECS="avc1.4d401f,mp4a.40.2",AUDIO="aud",SUBTITLES="sub"
+720.m3u8
+#EXT-X-STREAM-INF:BANDWIDTH=6000000,RESOLUTION=1920x1080,CODECS="avc1.640028,mp4a.40.2",AUDIO="aud",SUBTITLES="sub"
+1080.m3u8
+"""
+
+
+def export_tracks(tmp_path: Path, tracks, manifest_url: str | None) -> Path:
+    """Write a real export sidecar for a parsed ladder, as a finished download would."""
+    from unshackle.core.titles import Movie
+
+    class StubService:
+        pass
+
+    title = Movie(id_="movie-1", service=StubService, name="Example Movie", year=2024, language="en")
+    for track in [*tracks.videos, *tracks.audio, *tracks.subtitles]:
+        title.tracks.add(track)
+    title.tracks.manifest_url = manifest_url
+
+    runner = dl.__new__(dl)
+    runner.service = "EXAMPLE"
+    export = tmp_path / "export.json"
+    for track in [*title.tracks.videos, *title.tracks.audio, *title.tracks.subtitles]:
+        runner.write_export(export, title, track)
+    return export
+
+
+def import_tracks(export: Path):
+    svc = ImportService(SimpleNamespace(parent=None, params={}), "EXAMPLE", "movie-1", str(export))
+    return svc, svc.get_tracks(next(iter(svc.get_titles())))
+
+
+@pytest.fixture
+def ism_offline(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(ISM, "from_url", lambda url, session=None, **kw: ISM.from_text(ISM_MANIFEST, url))
+
+
+@pytest.mark.parametrize("manifest_url", [ISM_URL, None], ids=["with-manifest-url", "without-manifest-url"])
+def test_ism_round_trip_restores_segment_urls(tmp_path: Path, ism_offline: None, manifest_url: str | None) -> None:
+    """ISM.download_track reads data["ism"]["segments"], which to_dict cannot carry across an export.
+
+    Without a manifest_url the export's own track URLs are the only way back to the manifest,
+    which is the case the re-parse branch alone never covered.
+    """
+    original = ISM.from_text(ISM_MANIFEST, ISM_URL).to_tracks(language="en")
+    expected = {t.id: t.data["ism"]["segments"] for t in [*original.videos, *original.audio]}
+
+    _, tracks = import_tracks(export_tracks(tmp_path, original, manifest_url))
+
+    rebuilt = {t.id: t.data.get("ism", {}).get("segments") for t in [*tracks.videos, *tracks.audio]}
+    assert rebuilt == expected, "every ISM track must come back with the fragment URLs it was exported with"
+
+
+ISM_ENCRYPTED = """<?xml version="1.0"?>
+<SmoothStreamingMedia MajorVersion="2" MinorVersion="0" Duration="60000000" TimeScale="10000000">
+  <Protection><ProtectionHeader SystemID="{9A04F079-9840-4286-AB92-E65BE0885F95}">dGVzdA==</ProtectionHeader></Protection>
+  <StreamIndex Type="video" Name="video" Language="en" Url="QualityLevels({Bitrate})/Fragments(video={start_time})">
+    <QualityLevel Index="0" Bitrate="3000000" FourCC="H264" MaxWidth="1280" MaxHeight="720"
+                  CodecPrivateData="00000001674D401F"/>
+    <QualityLevel Index="1" Bitrate="6000000" FourCC="H264" MaxWidth="1920" MaxHeight="1080"
+                  CodecPrivateData="00000001674D401F"/>
+    <c t="0" d="20000000" r="3"/>
+  </StreamIndex>
+</SmoothStreamingMedia>
+"""
+
+
+def test_encrypted_ism_rung_gets_the_exported_keys(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """resolve_server_keys must recognise an ISM manifest's protection, not only a DASH one.
+
+    ISM.download_track reads only ``track.drm``, so a rung that misses key injection downloads
+    encrypted and muxes without error. DASH survives the same omission because it re-derives its
+    DRM from the manifest elements at download time.
+    """
+    monkeypatch.setattr(ISM, "from_url", lambda url, session=None, **kw: ISM.from_text(ISM_ENCRYPTED, url))
+    parsed = ISM.from_text(ISM_ENCRYPTED, ISM_URL).to_tracks(language="en")
+
+    # export the licensed rung only, as a finished download would
+    licensed = max(parsed.videos, key=lambda t: t.bitrate or 0)
+    export = export_tracks(tmp_path, parsed, ISM_URL)
+    doc = json.loads(export.read_text(encoding="utf8"))
+    doc["titles"]["movie-1"]["tracks"][licensed.id]["keys"] = {"0" * 32: "ab" * 16}
+    export.write_text(json.dumps(doc), encoding="utf8")
+
+    svc, tracks = import_tracks(export)
+    movie = next(iter(svc.get_titles()))
+    movie.tracks = tracks
+    svc.resolve_server_keys(movie)
+
+    other = min(tracks.videos, key=lambda t: t.bitrate or 0)
+    assert svc.track_is_encrypted(other), "an ISM ProtectionHeader means the rung is encrypted"
+    assert other.drm and other.drm[0].content_keys, "the unlicensed rung never received the exported key"
+
+
+def test_hls_round_trip_keeps_every_variant_url(tmp_path: Path) -> None:
+    """HLS re-fetches each variant from track.url at download time, so the stored URL has to survive."""
+    original = HLS.from_text(HLS_MANIFEST, HLS_URL).to_tracks(language="en")
+    expected = {t.id: str(t.url) for t in [*original.videos, *original.audio, *original.subtitles]}
+
+    _, tracks = import_tracks(export_tracks(tmp_path, original, HLS_URL))
+
+    rebuilt = {t.id: str(t.url) for t in [*tracks.videos, *tracks.audio, *tracks.subtitles]}
+    assert rebuilt == expected, "the whole exported ladder must survive, each with its own variant URL"
+    assert {t.descriptor.name for t in tracks} == {"HLS"}
+
+
+def test_hls_import_never_refetches_the_master_playlist(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """A master playlist can hand out a fresh token per fetch, so import must not depend on it.
+
+    This is why HLS is absent from MANIFEST_DATA_KEYS. Re-fetching would add nothing and give a
+    long-expired export another way to fail.
+    """
+    original = HLS.from_text(HLS_MANIFEST, HLS_URL).to_tracks(language="en")
+    export = export_tracks(tmp_path, original, HLS_URL)
+
+    def explode(*args: object, **kwargs: object) -> None:
+        raise AssertionError("import re-fetched an HLS manifest")
+
+    monkeypatch.setattr(HLS, "from_url", explode)
+    monkeypatch.setattr(HLS, "from_text", explode)
+
+    _, tracks = import_tracks(export)
+    assert len(list(tracks)) == 4
+
+
+def test_hls_tracks_tolerate_the_empty_manifest_data(tmp_path: Path) -> None:
+    """HLS.download_track assigns into track.data["hls"] without creating it first.
+
+    Track.data is a defaultdict, so the assignment auto-vivifies rather than raising. Pin that,
+    because a plain dict here would turn every HLS import into a KeyError mid-download.
+    """
+    original = HLS.from_text(HLS_MANIFEST, HLS_URL).to_tracks(language="en")
+    _, tracks = import_tracks(export_tracks(tmp_path, original, HLS_URL))
+
+    track = tracks.videos[0]
+    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]
@@ -0,0 +1,226 @@
+"""Tests for the manifest data an import must rebuild before a DASH track can download.
+
+``Track.to_dict`` cannot serialise manifest XML, so a rebuilt DASH track carries no
+``track.data["dash"]``. ``DASH.download_track`` reads that mapping unconditionally, so any
+DASH track the import hands back without it raises on the first segment. Two exports reach
+that state: one with no top-level ``manifest_url``, and one whose ``manifest_url`` covers
+only part of the ladder (a separate MPD per codec).
+"""
+
+from __future__ import annotations
+
+import json
+from functools import partial
+from pathlib import Path
+from typing import Any, Optional
+
+import click
+import pytest
+
+from unshackle.core.import_service import ImportService
+from unshackle.core.manifests import DASH
+
+MPD_TEMPLATE = """<?xml version="1.0"?>
+<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT10M">
+  <Period id="1">
+    <AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
+      <Representation id="{rep_id}" codecs="{codecs}" bandwidth="{bandwidth}" width="{width}" height="{height}">
+        <BaseURL>v.mp4</BaseURL><SegmentBase indexRange="0-1"/>
+      </Representation>
+    </AdaptationSet>
+  </Period>
+</MPD>
+"""
+
+AVC_URL = "https://example.invalid/avc/dash/m.mpd"
+HEVC_URL = "https://example.invalid/hevc/dash/m.mpd"
+
+MANIFESTS = {
+    AVC_URL: MPD_TEMPLATE.format(rep_id="v-avc", codecs="avc1.640028", bandwidth=5000000, width=1920, height=1080),
+    HEVC_URL: MPD_TEMPLATE.format(
+        rep_id="v-hevc", codecs="hvc1.2.4.L153.B0", bandwidth=6000000, width=3840, height=2160
+    ),
+}
+
+
+def exported_video(track_id: str, url: str, codec: str, width: int, height: int) -> dict[str, Any]:
+    return {
+        "type": "Video",
+        "id": track_id,
+        "url": url,
+        "language": "en",
+        "descriptor": "DASH",
+        "codec": codec,
+        "bitrate": 5000000 if codec == "AVC" else 6000000,
+        "width": width,
+        "height": height,
+    }
+
+
+def make_service(tmp_path: Path, manifest_url: Optional[str], tracks: dict[str, Any]) -> ImportService:
+    entry: dict[str, Any] = {
+        "meta": {"type": "movie", "id": "movie-1", "name": "Example Movie", "year": 2024, "language": "en"},
+        "manifest_type": "DASH",
+        "tracks": tracks,
+    }
+    if manifest_url:
+        entry["manifest_url"] = manifest_url
+    export = {"version": 2, "service": "EXAMPLE", "titles": {"movie-1": entry}}
+    path = tmp_path / "export.json"
+    path.write_text(json.dumps(export), encoding="utf8")
+    return ImportService(click.Context(click.Command("dl")), "EXAMPLE", "movie-1", str(path))
+
+
+@pytest.fixture(autouse=True)
+def local_manifests(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Serve each MPD from memory, so a test never depends on the network."""
+    monkeypatch.setattr(DASH, "from_url", lambda url, session=None, **kw: DASH.from_text(MANIFESTS[url], url))
+
+
+def test_export_without_manifest_url_still_rebuilds_dash_manifest_data(tmp_path: Path) -> None:
+    """No top-level manifest_url means the export's own track URLs are the only MPDs to fetch."""
+    service = make_service(tmp_path, None, {"v-avc": exported_video("v-avc", AVC_URL, "AVC", 1920, 1080)})
+    tracks = service.get_tracks(next(iter(service.get_titles())))
+
+    assert [t.id for t in tracks.videos] == ["v-avc"]
+    # DASH.download_track reads all three of these; a missing one is a KeyError mid-download.
+    assert set(tracks.videos[0].data.get("dash", {})) >= {"manifest", "adaptation_set", "representation"}
+
+
+def test_exported_hevc_on_its_own_mpd_survives_the_reparse(tmp_path: Path) -> None:
+    """A ladder split across per-codec MPDs loses every codec but manifest_url's if only URL tracks merge back."""
+    service = make_service(
+        tmp_path,
+        AVC_URL,
+        {"v-hevc": exported_video("v-hevc", HEVC_URL, "HEVC", 3840, 2160)},
+    )
+    tracks = service.get_tracks(next(iter(service.get_titles())))
+
+    codecs = {t.codec.name for t in tracks.videos if t.codec}
+    assert "HEVC" in codecs, "the exported HEVC track was dropped, so -v h.265 has nothing to select"
+    hevc = next(t for t in tracks.videos if t.codec and t.codec.name == "HEVC")
+    assert set(hevc.data.get("dash", {})) >= {"manifest", "adaptation_set", "representation"}
+    # bound to its own manifest's representation, not just to any populated mapping
+    assert hevc.data["dash"]["representation"].get("id") == "v-hevc"
+
+
+SUB_MPD_URL = "https://example.invalid/subs/de.mpd"
+
+MAIN_WITH_SUB = """<?xml version="1.0"?>
+<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT10M">
+  <Period id="1">
+    <AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
+      <Representation id="v-avc" codecs="avc1.640028" bandwidth="5000000" width="1920" height="1080">
+        <BaseURL>v.mp4</BaseURL><SegmentBase indexRange="0-1"/></Representation>
+    </AdaptationSet>
+    <AdaptationSet contentType="text" mimeType="text/vtt" lang="en">
+      <Representation id="s-en" bandwidth="1000">
+        <BaseURL>en.vtt</BaseURL><SegmentBase indexRange="0-1"/></Representation>
+    </AdaptationSet>
+  </Period>
+</MPD>
+"""
+
+FOREIGN_SUB_MPD = """<?xml version="1.0"?>
+<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT10M">
+  <Period id="1">
+    <AdaptationSet contentType="text" mimeType="text/vtt" lang="de">
+      <Representation id="s-de" bandwidth="1000">
+        <BaseURL>de.vtt</BaseURL><SegmentBase indexRange="0-1"/></Representation>
+    </AdaptationSet>
+  </Period>
+</MPD>
+"""
+
+
+def test_foreign_subtitle_adds_to_the_manifests_own(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Only a side-loaded (URL) subtitle replaces the re-parsed set.
+
+    A subtitle that lives on its own MPD is an addition. Nothing about it says the service
+    considered the main manifest's copy the worse one, and clearing on it deletes languages the
+    export never meant to drop.
+    """
+    manifests = {AVC_URL: MAIN_WITH_SUB, SUB_MPD_URL: FOREIGN_SUB_MPD}
+    monkeypatch.setattr(DASH, "from_url", lambda url, session=None, **kw: DASH.from_text(manifests[url], url))
+    foreign = DASH.from_text(FOREIGN_SUB_MPD, SUB_MPD_URL).to_tracks(language="de").subtitles[0]
+
+    service = make_service(
+        tmp_path,
+        AVC_URL,
+        {
+            foreign.id: {
+                "type": "Subtitle",
+                "id": foreign.id,
+                "url": SUB_MPD_URL,
+                "language": "de",
+                "descriptor": "DASH",
+                "codec": "WebVTT",
+            }
+        },
+    )
+    tracks = service.get_tracks(next(iter(service.get_titles())))
+
+    assert sorted(str(t.language) for t in tracks.subtitles) == ["de", "en"], (
+        "the main manifest's own subtitles were wiped by a subtitle from another MPD"
+    )
+
+
+TWO_RUNG_MPD = """<?xml version="1.0"?>
+<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT10M">
+  <Period id="1">
+    <AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
+      <Representation id="lo" codecs="avc1.640028" bandwidth="3000000" width="1920" height="1080">
+        <BaseURL>lo.mp4</BaseURL><SegmentBase indexRange="0-1"/></Representation>
+      <Representation id="hi" codecs="avc1.640028" bandwidth="{top}" width="1920" height="1080">
+        <BaseURL>hi.mp4</BaseURL><SegmentBase indexRange="0-1"/></Representation>
+    </AdaptationSet>
+  </Period>
+</MPD>
+"""
+
+
+@pytest.mark.parametrize(
+    "exported_url",
+    [AVC_URL, AVC_URL + "?token=EXPIRED"],
+    ids=["same-url", "signed-url"],
+)
+def test_reencoded_top_rung_does_not_readd_the_exported_one(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, exported_url: str
+) -> None:
+    """Track ids hash the bitrate, so re-encoding a rung makes the exported id look like a foreign MPD's.
+
+    Merging it back on that basis duplicates the rung, and the attribute fallback then binds the
+    copy to the first representation of the same codec and resolution: the 3 Mbps one. The result
+    downloads `lo` under a 6 Mbps name.
+
+    The signed variant is the same manifest reached by two strings: a service may fetch a signed
+    URL while registering the bare one as Title.tracks.manifest_url, so comparing them whole
+    reopens the duplicate.
+    """
+    exported = DASH.from_text(TWO_RUNG_MPD.format(top=6000000), exported_url).to_tracks(language="en")
+    top = next(t for t in exported.videos if t.bitrate == 6000000)
+    monkeypatch.setattr(
+        DASH, "from_url", lambda url, session=None, **kw: DASH.from_text(TWO_RUNG_MPD.format(top=6500000), url)
+    )
+
+    service = make_service(tmp_path, AVC_URL, {top.id: exported_video(top.id, exported_url, "AVC", 1920, 1080)})
+    tracks = service.get_tracks(next(iter(service.get_titles())))
+
+    assert sorted(t.bitrate for t in tracks.videos) == [3000000, 6500000], (
+        "the re-parsed ladder is the only source for this manifest's rungs"
+    )
+    bound = {t.data["dash"]["representation"].get("id") for t in tracks.videos}
+    assert bound == {"lo", "hi"}, "two tracks were bound to the same representation"
+
+
+def test_clearkey_license_hook_is_bindable(tmp_path: Path) -> None:
+    """dl.py binds service.get_clearkey_license for every track, DRM or not, before downloading.
+
+    Bind it the way dl.py does rather than asking hasattr: a non-callable attribute satisfies
+    hasattr and still fails at the partial with a TypeError.
+    """
+    service = make_service(tmp_path, AVC_URL, {})
+    bound = partial(service.get_clearkey_license, title=None, track=None)
+    assert callable(bound)
+    with pytest.raises(RuntimeError, match="keys come from the export"):
+        bound(challenge=b"")
@@ -79,9 +79,7 @@ def _resolve_import_manifest_data(
                    f"No language for '{title}': the {manifest_type} manifest has none and the "
                    f"service did not set Title.language."
                )
-            raise click.ClickException(
-                f"Failed to parse the {manifest_type} manifest for '{title}'. ({e})"
-            )
+            raise click.ClickException(f"Failed to parse the {manifest_type} manifest for '{title}'. ({e})")
        except Exception as e:
            raise click.ClickException(
                f"Failed to fetch the {manifest_type} manifest for '{title}'. "
@@ -246,13 +244,22 @@ class ImportService:
            for track in parsed:
                tracks.add(track)
            parsed_ids = {str(t.id) for t in tracks}
-            track_dicts = [t for t in tracks_map.values() if t.get("descriptor", "URL") == "URL"]
-            track_dicts.extend(
+            url_dicts = [t for t in tracks_map.values() if t.get("descriptor", "URL") == "URL"]
+            # Merge back only tracks from another manifest, such as a separate HEVC one; the
+            # re-parse already covers manifest_url's own. Query stripped when comparing, since a
+            # service can sign the URL it fetches and register the bare one for the same manifest.
+            manifest_base = str(manifest_url).split("?")[0]
+            foreign_dicts = [
                t
                for t in tracks_map.values()
-                if t.get("descriptor") == manifest_type and str(t.get("id")) not in parsed_ids
-            )
-            if any(t.get("type") == "Subtitle" for t in track_dicts):
+                if t.get("descriptor") == manifest_type
+                and str(t.get("url")).split("?")[0] != manifest_base
+                and str(t.get("id")) not in parsed_ids
+            ]
+            track_dicts = [*url_dicts, *foreign_dicts]
+            # Only a side-loaded subtitle replaces the re-parsed set, since the service chose its
+            # own copy over the manifest's. A subtitle on another manifest is an addition.
+            if any(t.get("type") == "Subtitle" for t in url_dicts):
                tracks.subtitles.clear()
        else:
            track_dicts = list(tracks_map.values())
@@ -355,14 +362,25 @@ class ImportService:

    @staticmethod
    def track_is_encrypted(track: Any) -> bool:
-        """True if the track carries DRM or its DASH manifest declares ContentProtection."""
+        """True if the track carries DRM or its manifest declares protection.
+
+        ISM is checked as well as DASH because ISM.download_track reads only ``track.drm``. A rung
+        that misses key injection here downloads encrypted and muxes without error. DASH re-derives
+        its DRM from the manifest elements, so it survives the same omission.
+        """
        if track.drm:
            return True
-        dash = track.data.get("dash") if getattr(track, "data", None) else None
+        data = getattr(track, "data", None) or {}
+        dash = data.get("dash")
        if dash:
            for element in (dash.get("representation"), dash.get("adaptation_set")):
                if element is not None and element.findall("ContentProtection"):
                    return True
+        ism = data.get("ism")
+        if ism:
+            manifest = ism.get("manifest")
+            if manifest is not None and manifest.findall(".//ProtectionHeader"):
+                return True
        return False

    def exported_drm_system(self) -> str: