diff --git a/docs/guide/downloading.md b/docs/guide/downloading.md index 408e8c15..ea014cf4 100644 --- a/docs/guide/downloading.md +++ b/docs/guide/downloading.md @@ -779,8 +779,9 @@ unshackle dl --skip-dl EXAMPLE 81234567 ``` `--export` writes a `mediaexport` JSON file into the configured exports directory. It holds -the manifest, the DRM init data, the content keys and the title metadata for each title. -`unshackle import` reads it to reconstruct a download later. The importer also reads the +the manifest, the DRM init data, the content keys, the KIDs of each track and the title +metadata for each title. `unshackle import` reads it to reconstruct a download later, and gives +each track only the content keys for its KIDs. The importer also reads the older unshackle `version: 2` files and unidl's own export files. ```shell title="Export track info and keys" diff --git a/tests/core/test_import_track_keys.py b/tests/core/test_import_track_keys.py index 21222cd1..65c870bb 100644 --- a/tests/core/test_import_track_keys.py +++ b/tests/core/test_import_track_keys.py @@ -7,6 +7,7 @@ the fallback for a track that declares no KID. from __future__ import annotations +import json import logging from pathlib import Path from types import SimpleNamespace @@ -75,8 +76,11 @@ def widevine(kid: UUID, key: str) -> Widevine: return drm -def write_export(tmp_path: Path) -> Path: - """Export the ladder as a finished download would: one licensed DRM per track.""" +def write_export(tmp_path: Path, kids: bool = True) -> Path: + """Export the ladder as a finished download would: one licensed DRM per track. + + ``kids=False`` removes the per-track KID lists, as in a file from before they existed. + """ 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) @@ -88,9 +92,20 @@ def write_export(tmp_path: Path) -> Path: 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))) + if not kids: + doc = json.loads(export.read_text(encoding="utf8")) + for row in doc["titles"][0]["tracks"]: + row.pop("kids", None) + export.write_text(json.dumps(doc), encoding="utf8") return export +def rows(export: Path) -> dict[str, dict[str, Any]]: + """The export's track rows, by the same labels as the tracks.""" + doc = json.loads(export.read_text(encoding="utf8")) + return {"video" if r["type"] == "video" else r["language"]: r for r in doc["titles"][0]["tracks"]} + + 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) @@ -109,21 +124,58 @@ def keys_by_track(svc: ImportService) -> dict[str, dict[UUID, str]]: return {label(t): dict(t.drm[0].content_keys) for t in [*movie.tracks.videos, *movie.tracks.audio]} +def test_the_export_lists_each_tracks_kids(tmp_path: Path) -> None: + exported = rows(write_export(tmp_path)) + + assert exported["video"]["kids"] == [VIDEO_KID.hex] + assert exported["en"]["kids"] == [AUDIO_KID.hex] + assert "kids" not in exported["de"], "a track exported without DRM names no KID" + + +def test_listed_kids_give_each_track_its_keys_without_a_probe(tmp_path: Path) -> None: + svc = import_service(write_export(tmp_path)) + probed: list[str] = [] + track_kids = svc.track_kids + + def spy(track: Any, session: Any = None) -> Any: + probed.append(label(track)) + return track_kids(track, session) + + with patch.object(svc, "track_kids", side_effect=spy): + keys = keys_by_track(svc) + + assert keys["video"] == {VIDEO_KID: VIDEO_KEY} + assert keys["en"] == {AUDIO_KID: AUDIO_KEY} + assert probed == ["de"] + + +def test_a_listed_kid_picks_the_drm_entry_of_a_direct_url_track(tmp_path: Path) -> None: + export = write_export(tmp_path) + svc = import_service(export) + + rebuilt = svc.rebuild_drm({"id": rows(export)["en"]["id"]}, "movie-1") + + assert rebuilt and isinstance(rebuilt[0], Widevine) + assert rebuilt[0].kids == [AUDIO_KID] + assert rebuilt[0].content_keys == {AUDIO_KID: AUDIO_KEY} + + 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))) + """An export from before the per-track KID lists falls back to the KIDs the manifest declares.""" + keys = keys_by_track(import_service(write_export(tmp_path, kids=False))) 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))) + keys = keys_by_track(import_service(write_export(tmp_path, kids=False))) 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)) + svc = import_service(write_export(tmp_path, kids=False)) single: list[dict[str, Any]] = [widevine(VIDEO_KID, VIDEO_KEY).to_dict()] with patch.object(svc, "title_drm_dicts", return_value=single): @@ -135,7 +187,7 @@ def test_a_rebuilt_drm_holds_only_the_keys_for_the_kids_its_pssh_names(tmp_path: 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)) + svc = import_service(write_export(tmp_path, kids=False)) rebuilt = svc.rebuild_drm({"id": "u1"}, "movie-1") diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 8c957353..9b22d34d 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -48,7 +48,7 @@ from unshackle.core.console import GradientPulseBarColumn, SyncLive, console, li from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings from unshackle.core.credential import Credential from unshackle.core.downloaders import default_max_workers, format_speed, parse_speed_limit, set_speed_limit -from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, verify +from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, own_kids, verify from unshackle.core.events import events from unshackle.core.providers.anilist import parse_anilist_ref from unshackle.core.providers.tvdb import SEASON_TYPES, parse_int @@ -3970,9 +3970,13 @@ class dl: tracks_map = entry.ext("unshackle").setdefault("tracks", {}) tracks_map.setdefault(str(track.id), track.to_dict()) + # 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: if row.get("id") == str(track.id): row["selected"] = True + if track_kids: + row["kids"] = list(dict.fromkeys([*(row.get("kids") or []), *track_kids])) if drm is not None: if hasattr(drm, "to_dict"): diff --git a/unshackle/core/drm/__init__.py b/unshackle/core/drm/__init__.py index bae6603a..2a56e7d0 100644 --- a/unshackle/core/drm/__init__.py +++ b/unshackle/core/drm/__init__.py @@ -48,4 +48,26 @@ def drm_from_dict(data: dict[str, Any]) -> Union[Widevine, PlayReady, ClearKeyCE return drm -__all__ = ("ClearKey", "ClearKeyCENC", "Widevine", "PlayReady", "MonaLisa", "DRM_T", "drm_from_dict") +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] + + +def own_kids(drm: Any) -> list[UUID]: + """Return the real KIDs that identify the track's own content key, as the DRM object knows them.""" + own = getattr(drm, "own_kids", None) + kids = own() if callable(own) else getattr(drm, "kids", None) or [] + return [kid for kid in kids if kid.int and kid not in Widevine.PLACEHOLDER_KIDS] + + +__all__ = ( + "ClearKey", + "ClearKeyCENC", + "Widevine", + "PlayReady", + "MonaLisa", + "DRM_T", + "drm_from_dict", + "real_kids", + "own_kids", +) diff --git a/unshackle/core/import_service.py b/unshackle/core/import_service.py index cbb0f4b8..2a4bebe8 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 Widevine, drm_from_dict +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.titles import Episode, Movies, Series, Title_T, Titles_T, remap_titles @@ -95,11 +95,6 @@ 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. @@ -417,34 +412,55 @@ class ImportService: out.append({"system": system, "pssh_b64": d.pssh, "kids": list(entry.keys)}) return out + def exported_kids(self, title_id: str, track_id: str) -> Optional[set[str]]: + """Return the KIDs (hex) the export's row for the track lists, or None when it lists none.""" + entry = self.doc.get(title_id) + row = next((r for r in (entry.tracks if entry else []) if str(r.get("id", "")) == track_id), None) + return set(row["kids"]) if row and row.get("kids") else None + + def exported_keys(self, title_id: str, track_id: str) -> dict[UUID, str]: + """Return the exported content keys for the KIDs the row of the track lists, else all of the title.""" + entry = self.doc.get(title_id) + return {UUID(hex=kid): key for kid, key in (entry.keys_for(track_id) if entry else {}).items()} + def rebuild_drm(self, track_dict: dict[str, Any], title_id: str) -> Optional[list[Any]]: """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. 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. + carries its own is honoured first. The track's row in the export can list its KIDs; + then the track gets only their keys, and the DRM entry that names them. Without that + list, 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. """ 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: + own_drm = track_dict.get("drm") + row_kids = None if own_keys else self.exported_kids(title_id, str(track_dict.get("id", ""))) + if own_keys: + keys = own_keys + else: + keys = {kid.hex: key for kid, key in self.exported_keys(title_id, str(track_dict.get("id", ""))).items()} + drm_dicts = own_drm or self.title_drm_dicts(title_id) + if row_kids and not own_drm: + drm_dicts = [d for d in drm_dicts if row_kids & {str(kid).lower() for kid in d.get("kids") or []}] + if not drm_dicts and not keys and not row_kids: return None drm_obj = None - if drm_dicts and (track_dict.get("drm") or len(drm_dicts) == 1): + if drm_dicts and (own_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: + if drm_obj is not None and not own_keys and not row_kids: 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: + stub_kids = list(keys) or sorted(row_kids or []) + if drm_obj is None and stub_kids: drm_type = (drm_dicts[0].get("system", "Widevine").lower()) if drm_dicts else "widevine" - drm_obj = RemoteService.create_drm_stub(drm_type, list(keys.keys())) + drm_obj = RemoteService.create_drm_stub(drm_type, stub_kids) if drm_obj is None: return None @@ -457,8 +473,10 @@ 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. 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) + tracks. A track gets only the content keys for the KIDs it declares: first the KIDs + its row in the export lists, which needs no probe, then the KIDs of its DRM, manifest + or init segment. It gets 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. @@ -468,14 +486,20 @@ class ImportService: return system = self.exported_drm_system() + title_id = str(title.id) for track in title.tracks: if not isinstance(track, (Video, Audio)): continue - 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 + row_kids = self.exported_kids(title_id, str(track.id)) + if row_kids: + declared: Optional[set[UUID]] = {UUID(hex=kid) for kid in row_kids} + keys = self.exported_keys(title_id, str(track.id)) + else: + 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: