diff --git a/tests/core/test_import_track_keys.py b/tests/core/test_import_track_keys.py new file mode 100644 index 00000000..21222cd1 --- /dev/null +++ b/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'' + + +MPD = f""" + + + + {protection(VIDEO_KID)} + + v.mp4 + + + + {protection(AUDIO_KID)} + + a.mp4 + + + + {protection(None)} + + a2.mp4 + + + + +""" + + +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)] diff --git a/unshackle/core/import_service.py b/unshackle/core/import_service.py index be83fae2..541732fc 100644 --- a/unshackle/core/import_service.py +++ b/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.""" diff --git a/unshackle/core/remote_service.py b/unshackle/core/remote_service.py index 5b750d2e..1c0089da 100644 --- a/unshackle/core/remote_service.py +++ b/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)