fix(import): give each track only the keys it declares

The whole title pool stays the fallback for a track that declares no KID.

tests/core/test_import_track_keys.py

@@ -0,0 +1,145 @@
+"""An import gives each encrypted track only the exported keys for the KIDs it declares.
+
+A decrypter gives the all-zero KID the track's own content key only when it can tell which
+key that is. A track that holds the whole title pool makes that impossible, so the pool is
+the fallback for a track that declares no KID.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, cast
+from unittest.mock import patch
+from uuid import UUID
+
+import click
+from pywidevine.pssh import PSSH
+
+from unshackle.commands.dl import dl
+from unshackle.core.drm import Widevine
+from unshackle.core.drm.key_args import ZERO_KID
+from unshackle.core.import_service import ImportService
+from unshackle.core.manifests import DASH
+from unshackle.core.titles import Movie
+from unshackle.core.tracks import Video
+
+MPD_URL = "https://example.invalid/m.mpd"
+VIDEO_KID = UUID("00112233-4455-6677-8899-aabbccddeeff")
+AUDIO_KID = UUID("ffeeddcc-bbaa-9988-7766-554433221100")
+VIDEO_KEY = "0000000000000000000000000000000a"
+AUDIO_KEY = "0000000000000000000000000000000b"
+POOL = {VIDEO_KID: VIDEO_KEY, AUDIO_KID: AUDIO_KEY}
+
+
+def protection(kid: UUID | None) -> str:
+    default_kid = f' cenc:default_KID="{kid}"' if kid else ""
+    return f'<ContentProtection schemeIdUri="urn:mpeg:dash:mp4protection:2011" value="cenc"{default_kid}/>'
+
+
+MPD = f"""<?xml version="1.0"?>
+<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" xmlns:cenc="urn:mpeg:cenc:2013" type="static"
+     mediaPresentationDuration="PT10M">
+  <Period id="1">
+    <AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
+      {protection(VIDEO_KID)}
+      <Representation id="v1" codecs="avc1.640028" bandwidth="5000000" width="1920" height="1080">
+        <BaseURL>v.mp4</BaseURL><SegmentBase indexRange="0-1"/>
+      </Representation>
+    </AdaptationSet>
+    <AdaptationSet contentType="audio" mimeType="audio/mp4" lang="en">
+      {protection(AUDIO_KID)}
+      <Representation id="a1" codecs="mp4a.40.2" bandwidth="128000">
+        <BaseURL>a.mp4</BaseURL><SegmentBase indexRange="0-1"/>
+      </Representation>
+    </AdaptationSet>
+    <AdaptationSet contentType="audio" mimeType="audio/mp4" lang="de">
+      {protection(None)}
+      <Representation id="a2" codecs="mp4a.40.2" bandwidth="128000">
+        <BaseURL>a2.mp4</BaseURL><SegmentBase indexRange="0-1"/>
+      </Representation>
+    </AdaptationSet>
+  </Period>
+</MPD>
+"""
+
+
+class StubService:
+    """Stands in for the service class slot on Movie; never instantiated."""
+
+
+def widevine(kid: UUID, key: str) -> Widevine:
+    drm = Widevine(pssh=PSSH.new(system_id=PSSH.SystemId.Widevine, key_ids=[kid]), kid=kid)
+    drm.content_keys[kid] = key
+    return drm
+
+
+def write_export(tmp_path: Path) -> Path:
+    """Export the ladder as a finished download would: one licensed DRM per track."""
+    title = Movie(id_="movie-1", service=StubService, name="Example Movie", year=2024, language="en")
+    for track in DASH.from_text(MPD, MPD_URL).to_tracks(language="en"):
+        title.tracks.add(track)
+    title.tracks.manifest_url = MPD_URL
+    runner = dl.__new__(dl)
+    runner.service = "EXAMPLE"
+    runner.log = logging.getLogger("download")
+    export = tmp_path / "export.json"
+    drm = {"video": widevine(VIDEO_KID, VIDEO_KEY), "en": widevine(AUDIO_KID, AUDIO_KEY)}
+    for track in [*title.tracks.videos, *title.tracks.audio]:
+        runner.write_export(export, title, track, drm.get(label(track)))
+    return export
+
+
+def label(track: Any) -> str:
+    """The video track, or an audio track by its language: the manifest gives no stable ids."""
+    return "video" if isinstance(track, Video) else str(track.language)
+
+
+def import_service(export: Path) -> ImportService:
+    ctx = cast(click.Context, SimpleNamespace(parent=None, params={}))
+    return ImportService(ctx, "EXAMPLE", "movie-1", str(export))
+
+
+def keys_by_track(svc: ImportService) -> dict[str, dict[UUID, str]]:
+    movie = next(iter(svc.get_titles()))
+    with patch.object(DASH, "from_url", side_effect=lambda url, session=None, **kw: DASH.from_text(MPD, url)):
+        movie.tracks = svc.get_tracks(movie)
+    svc.resolve_server_keys(movie)
+    return {label(t): dict(t.drm[0].content_keys) for t in [*movie.tracks.videos, *movie.tracks.audio]}
+
+
+def test_each_track_gets_only_the_keys_for_its_own_kids(tmp_path: Path) -> None:
+    keys = keys_by_track(import_service(write_export(tmp_path)))
+
+    assert keys["video"] == {VIDEO_KID: VIDEO_KEY}
+    assert keys["en"] == {AUDIO_KID: AUDIO_KEY}
+
+
+def test_a_track_that_declares_no_kid_gets_the_whole_pool(tmp_path: Path) -> None:
+    keys = keys_by_track(import_service(write_export(tmp_path)))
+
+    assert keys["de"] == POOL
+
+
+def test_a_rebuilt_drm_holds_only_the_keys_for_the_kids_its_pssh_names(tmp_path: Path) -> None:
+    svc = import_service(write_export(tmp_path))
+    single: list[dict[str, Any]] = [widevine(VIDEO_KID, VIDEO_KEY).to_dict()]
+
+    with patch.object(svc, "title_drm_dicts", return_value=single):
+        rebuilt = svc.rebuild_drm({"id": "u1"}, "movie-1")
+
+    assert rebuilt and rebuilt[0].content_keys == {VIDEO_KID: VIDEO_KEY}
+
+
+def test_a_title_with_several_drm_entries_gives_no_track_a_guessed_own_key(tmp_path: Path) -> None:
+    """Which entry belongs to a direct-URL track is unknown, so the track gets a stub over the
+    pool, and the zero-KID fallback stays empty instead of taking the first entry's key."""
+    svc = import_service(write_export(tmp_path))
+
+    rebuilt = svc.rebuild_drm({"id": "u1"}, "movie-1")
+
+    assert rebuilt
+    drm = rebuilt[0]
+    assert drm.content_keys == POOL
+    assert not [arg for arg in drm.mp4decrypt_key_args() if arg.startswith(ZERO_KID)]

unshackle/core/import_service.py

@@ -14,7 +14,7 @@ from langcodes import tag_is_valid
 from unshackle.core.config import config
 from unshackle.core.constants import AnyTrack
 from unshackle.core.credential import Credential
-from unshackle.core.drm import drm_from_dict
+from unshackle.core.drm import Widevine, drm_from_dict
 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.titles import Episode, Movies, Series, Title_T, Titles_T, remap_titles
@@ -95,6 +95,11 @@ def resolve_import_manifest_data(
                 track.data.update(matched.data)
 
 
+def real_kids(drm: Any) -> list[UUID]:
+    """Return the KIDs a DRM object names, without the all-zero and test-pattern placeholders."""
+    return [kid for kid in getattr(drm, "kids", None) or [] if kid.int and kid not in Widevine.PLACEHOLDER_KIDS]
+
+
 class ImportService:
     """Reconstructs a download from an export JSON.
 
@@ -416,19 +421,26 @@ class ImportService:
         """Rebuild a DRM object (from stored PSSH, falling back to a stub) with the exported keys.
 
         Keys and DRM are title-level in the shared format; a legacy v2 track dict that still
-        carries its own is honoured first.
+        carries its own is honoured first. A title with more than one DRM entry does not say
+        which entry is this track's, so the track gets a stub that names every exported KID
+        instead of a guess. A rebuilt PSSH holds only the keys for the KIDs it names.
         """
-        keys = track_dict.get("keys") or {kid.hex: key for kid, key in self.key_pool(title_id).items()}
+        own_keys = track_dict.get("keys")
+        keys = own_keys or {kid.hex: key for kid, key in self.key_pool(title_id).items()}
         drm_dicts = track_dict.get("drm") or self.title_drm_dicts(title_id)
         if not drm_dicts and not keys:
             return None
 
         drm_obj = None
-        if drm_dicts:
+        if drm_dicts and (track_dict.get("drm") or len(drm_dicts) == 1):
             try:
                 drm_obj = drm_from_dict(drm_dicts[0])
             except Exception as e:
                 self.log.debug(f"Falling back to DRM stub (PSSH rebuild failed: {e})")
+            if drm_obj is not None and not own_keys:
+                declared = {kid.hex for kid in real_kids(drm_obj)}
+                if declared:
+                    keys = {kid_hex: key for kid_hex, key in keys.items() if kid_hex in declared}
 
         if drm_obj is None and keys:
             drm_type = (drm_dicts[0].get("system", "Widevine").lower()) if drm_dicts else "widevine"
@@ -445,54 +457,84 @@ class ImportService:
         """Inject exported keys into the selected encrypted tracks by KID (no network).
 
         dl.py calls this method after selection. It touches only encrypted video and audio
-        tracks. Encrypted DASH tracks (no DRM at parse time) get a stub holding the keys,
-        which DASH.download_track preserves. decrypt() applies the content key whose KID
-        matches the media.
+        tracks. A track gets only the content keys for the KIDs it declares, and the whole
+        title pool only when it declares no KID. Encrypted DASH tracks (no DRM at parse time)
+        get a stub holding the content keys, which DASH.download_track preserves. The
+        decrypters then match each content key to the media by KID, and give the all-zero KID
+        only the content key of the track.
         """
         pool = self.key_pool(str(title.id)) or self.key_pool()
         if not pool:
             return
 
         system = self.exported_drm_system()
-        kid_hexes = [kid.hex for kid in pool]
 
         for track in title.tracks:
-            if not isinstance(track, (Video, Audio)) or not self.track_is_encrypted(track, self.session):
+            if not isinstance(track, (Video, Audio)):
                 continue
-            drm_obj = track.drm[0] if track.drm else RemoteService.create_drm_stub(system, kid_hexes)
-            for kid, key in pool.items():
+            declared = self.track_kids(track, self.session)
+            if declared is None:
+                continue
+            keys = {kid: key for kid, key in pool.items() if kid in declared} if declared else pool
+            if track.drm:
+                drm_obj = track.drm[0]
+            else:
+                drm_obj = RemoteService.create_drm_stub(system, [kid.hex for kid in (declared or pool)])
+            for kid, key in keys.items():
                 drm_obj.content_keys[kid] = key
             track.drm = [drm_obj]
             self._server_cdm_type = drm_obj.__class__.__name__.lower()
 
-    @staticmethod
-    def track_is_encrypted(track: Any, session: Optional[requests.Session] = None) -> bool:
-        """True if the track carries DRM, its manifest declares protection, or its init segment does.
+    @classmethod
+    def track_is_encrypted(cls, track: Any, session: Optional[requests.Session] = None) -> bool:
+        """True if the track carries DRM, its manifest declares protection, or its init segment does."""
+        return cls.track_kids(track, session) is not None
 
-        This method examines ISM as well as DASH, because ISM.download_track reads only
-        ``track.drm``. A rung that misses content key injection here downloads encrypted and
-        muxes without error. A DASH manifest with no ContentProtection can still describe an
-        encrypted track, with the PSSH only in the init segment, so that case probes the init
-        segment with `session`.
+    @staticmethod
+    def track_kids(track: Any, session: Optional[requests.Session] = None) -> Optional[set[UUID]]:
+        """Return the KIDs an encrypted track declares, or None when the track is not encrypted.
+
+        An empty set means that the track is encrypted but declares no KID. The KIDs come from
+        the track's DRM, then from the DASH ContentProtection default_KID, then from the DRM in
+        the init segment. This method examines ISM as well as DASH, because ISM.download_track
+        reads only ``track.drm``. A rung that misses content key injection here downloads
+        encrypted and muxes without error. A DASH manifest with no ContentProtection can still
+        describe an encrypted track, with the PSSH only in the init segment, so that case
+        probes the init segment with `session`.
         """
         if track.drm:
-            return True
+            return {kid for drm in track.drm for kid in real_kids(drm)}
         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
+            protections = [
+                cp
+                for element in (dash.get("representation"), dash.get("adaptation_set"))
+                if element is not None
+                for cp in element.findall("ContentProtection")
+            ]
+            if protections:
+                kids: set[UUID] = set()
+                for cp in protections:
+                    value = cp.get("{urn:mpeg:cenc:2013}default_KID") or cp.get("default_KID")
+                    try:
+                        kid = UUID(value) if value else None
+                    except ValueError:
+                        kid = None
+                    if kid and kid.int:
+                        kids.add(kid)
+                return kids
             from unshackle.core.api.handlers import drm_from_init_segment
 
-            if drm_from_init_segment(track, session):
-                return True
+            init_drm = drm_from_init_segment(track, session)
+            if init_drm:
+                return {kid for drm in init_drm for kid in real_kids(drm)}
         ism = data.get("ism")
         if ism:
             manifest = ism.get("manifest")
             if manifest is not None and manifest.findall(".//ProtectionHeader"):
-                return True
-        return False
+                return set()
+        return None
 
     def exported_drm_system(self) -> str:
         """The DRM system the exporter licensed (e.g. 'playready'), defaulting to widevine."""

unshackle/core/remote_service.py

@@ -1345,7 +1345,8 @@ class RemoteService:
             kid_uuids = [UUID(hex=k) for k in kid_hexes]
             WIDEVINE_SYSTEM_ID = UUID("edef8ba9-79d6-4ace-a3c8-27dcd51d21ed")
             dummy_pssh = WvPSSH.new(system_id=WIDEVINE_SYSTEM_ID, key_ids=kid_uuids)
-            return Widevine(pssh=dummy_pssh, kid=kid_hexes[0])
+            # no kid: the stub names several KIDs, and the first is not the track's own
+            return Widevine(pssh=dummy_pssh)
 
     def get_chapters(self, title: Title_T) -> Chapters:
         title_id = str(title.id)