From c4677341c8a60c8c75214fffffdb826aca2d7f95 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Mon, 2 Nov 2026 21:13:23 -0700 Subject: [PATCH] refactor(drm): resolve content keys in one licence path behind ServerKeySource Drop null keys from every licence, server and export source, not only Widevine. dl now finds remote sessions and imports through ServerKeySource instead of private fields. --- docs/dev/service-architecture.md | 6 - tests/core/test_export.py | 2 +- tests/core/test_import_missing_key.py | 7 +- tests/core/test_licence_turns.py | 50 +- tests/core/test_vault_bad_keys.py | 9 +- tests/remote/key_source.py | 23 + .../test_remote_service_server_cdm_license.py | 10 +- .../unit/test_server_cdm_key_verification.py | 20 +- .../remote/unit/test_server_cdm_max_height.py | 14 +- .../unit/test_server_cdm_vault_write.py | 40 +- unshackle/commands/dl.py | 672 +++++++----------- unshackle/core/api/handlers.py | 7 +- unshackle/core/drm/__init__.py | 17 +- unshackle/core/import_service.py | 25 +- unshackle/core/remote_service.py | 52 +- 15 files changed, 463 insertions(+), 491 deletions(-) create mode 100644 tests/remote/key_source.py diff --git a/docs/dev/service-architecture.md b/docs/dev/service-architecture.md index e83ffdd7..ef207c33 100644 --- a/docs/dev/service-architecture.md +++ b/docs/dev/service-architecture.md @@ -98,12 +98,6 @@ These are abstract. A service will not load unless it defines all three: - **Event callbacks**: `on_segment_downloaded`, `on_track_downloaded`, `on_track_decrypted`, `on_track_repacked`, `on_track_multiplex` let a service react to pipeline events. -- **`resolve_server_keys(title)`**: an escape hatch for services whose decryption keys are - resolved server-side *ahead of* the normal per-track license flow. Unlike everything else - here it is **not declared on the `Service` base class at all**: `dl.py` probes for it with - `hasattr` and calls it (once, after track selection) only if your service defines it. Use - it when your service must fetch or inject keys in bulk before download, separate from - the per-track license callbacks below. ### What `__init__` wires up for you diff --git a/tests/core/test_export.py b/tests/core/test_export.py index e20918d0..3a5dc784 100644 --- a/tests/core/test_export.py +++ b/tests/core/test_export.py @@ -301,7 +301,7 @@ def test_import_takes_its_drm_type_from_the_export(tmp_path: Path, system: str, doc["titles"][0]["drm"] = [{"system": system, "pssh": "AAAA"}] export.write_text(json.dumps(doc), encoding="utf8") - assert ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export))._server_cdm_type == expected + assert ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export)).server_cdm_type == expected class KeyDRM: diff --git a/tests/core/test_import_missing_key.py b/tests/core/test_import_missing_key.py index e49b9143..10672619 100644 --- a/tests/core/test_import_missing_key.py +++ b/tests/core/test_import_missing_key.py @@ -18,6 +18,7 @@ import pytest from pywidevine.pssh import PSSH from rich.table import Table +from tests.remote.key_source import remote_keys from unshackle.commands import dl as dl_module from unshackle.commands.dl import dl from unshackle.core.drm import Widevine @@ -69,7 +70,7 @@ def make_cmd(monkeypatch: pytest.MonkeyPatch, vaults: Any) -> dl: cmd.service = "EXAMPLE" cmd.vaults = vaults cmd.vault_cache_tally = None - cmd._remote_service = SimpleNamespace(_server_cdm=True, _server_cdm_type="widevine") + cmd._remote_service = remote_keys() dl.LICENSE_KEY_CACHE.clear() dl.DRM_LOCKS.clear() return cmd @@ -111,7 +112,7 @@ def test_the_error_names_the_kids_without_a_content_key() -> None: def test_the_server_cdm_branch_lets_the_import_error_through(monkeypatch: pytest.MonkeyPatch) -> None: drm = FakeDRM({}) - track = SimpleNamespace(drm=[drm]) + track = SimpleNamespace(id="vid", drm=[drm]) cmd = make_cmd(monkeypatch, EmptyVaults()) with pytest.raises(click.ClickException, match=KID.hex): @@ -130,7 +131,7 @@ def test_the_retry_after_a_bad_vault_key_ends_in_the_import_error( assert vaults.get_key(KID) == (BAD, poisoned) cmd = make_cmd(monkeypatch, vaults) drm = FakeDRM({KID: BAD, OTHER: OTHER_KEY}) - track = SimpleNamespace(drm=[drm]) + track = SimpleNamespace(id="vid", drm=[drm]) calls: list[str] = [] path = tmp_path / "video.mp4" path.write_bytes(b"cipher") diff --git a/tests/core/test_licence_turns.py b/tests/core/test_licence_turns.py index 1b2be4b7..ff7443d9 100644 --- a/tests/core/test_licence_turns.py +++ b/tests/core/test_licence_turns.py @@ -6,7 +6,7 @@ import logging import threading import time from types import SimpleNamespace -from typing import Any +from typing import Any, Callable, Optional from uuid import UUID import pytest @@ -14,10 +14,11 @@ from pywidevine.pssh import PSSH from rich.table import Table from rich.tree import Tree +from tests.core.test_all_drm import playready from unshackle.commands.dl import dl from unshackle.core.cdm.detect import cdm_type_stub from unshackle.core.constants import DOWNLOAD_CANCELLED, DownloadCancelled -from unshackle.core.drm import Widevine +from unshackle.core.drm import ClearKeyCENC, PlayReady, Widevine pytestmark = pytest.mark.unit @@ -151,3 +152,48 @@ def test_a_key_label_is_added_once_with_its_marker_and_source() -> None: f"[text2]{KID_A.hex}:{KEYS[KID_A]}* from cache", f"[text2]{KID_B.hex}:{KEYS[KID_B]}", ] + + +@pytest.mark.parametrize( + "drm_class, build, cdm", + [ + (Widevine, lambda: widevine(KID_A), "widevine"), + (PlayReady, lambda: playready(KID_A), "playready"), + (ClearKeyCENC, lambda: ClearKeyCENC(kids=[KID_A]), None), + ], + ids=["widevine", "playready", "clearkey"], +) +def test_an_all_zero_licence_key_is_dropped( + cmd: dl, monkeypatch: pytest.MonkeyPatch, drm_class: type, build: Callable[[], Any], cdm: Optional[str] +) -> None: + def get_content_keys(self: Any, **_: Any) -> None: + self.content_keys.update({KID_A: KEYS[KID_A], KID_C: "0" * 32}) + + monkeypatch.setattr(drm_class, "get_content_keys", get_content_keys) + monkeypatch.setattr(dl, "LICENSE_KEY_CACHE", {}) + monkeypatch.setattr(dl, "DRM_LOCKS", {}) + drm = build() + if cdm: + cmd.cdm = cdm_type_stub(cdm) + + prepare(cmd, drm, KID_A, new_table()) + + assert drm.content_keys == {KID_A: KEYS[KID_A]} + assert dl.LICENSE_KEY_CACHE == {KID_A: KEYS[KID_A]} + + +def test_vaults_only_stops_at_the_first_kid_no_vault_holds(cmd: dl, licences: list[str]) -> None: + asked: list[UUID] = [] + + class EmptyVaults(NoVaults): + def get_key(self, kid: UUID) -> tuple[None, None]: + asked.append(kid) + return None, None + + cmd.vaults = EmptyVaults() + + with pytest.raises(Widevine.Exceptions.CEKNotFound, match=KID_A.hex): + prepare(cmd, widevine(KID_A, KID_B), KID_A, new_table(), vaults_only=True) + + assert asked == [KID_A] + assert licences == [] diff --git a/tests/core/test_vault_bad_keys.py b/tests/core/test_vault_bad_keys.py index 1359af4d..77e2fdf0 100644 --- a/tests/core/test_vault_bad_keys.py +++ b/tests/core/test_vault_bad_keys.py @@ -8,6 +8,7 @@ from uuid import UUID import pytest +from tests.remote.key_source import remote_keys from unshackle.commands import dl as dl_module from unshackle.commands.dl import dl from unshackle.core.vault import RateLimited, vault_session @@ -831,10 +832,11 @@ def test_prefer_vault_keys_queries_each_named_kid_once() -> None: asked: list[UUID] = [] cmd = dl.__new__(dl) cmd.vaults = SimpleNamespace(get_key=lambda kid: asked.append(kid) or (None, None)) - cmd._remote_service = SimpleNamespace(server_vault_keys={}) video, audio = (SimpleNamespace(kids=[KID], content_keys={KID: GOOD, SIBLING: BAD}) for _ in range(2)) - cmd.prefer_vault_keys(SimpleNamespace(tracks=[SimpleNamespace(drm=[video]), SimpleNamespace(drm=[audio])])) + cmd.prefer_vault_keys( + remote_keys(), SimpleNamespace(tracks=[SimpleNamespace(drm=[video]), SimpleNamespace(drm=[audio])]) + ) assert asked == [KID] @@ -845,13 +847,12 @@ def test_cache_resolved_keys_skips_a_key_a_local_vault_holds(tmp_path: Path) -> cmd = dl.__new__(dl) cmd.vaults = Vaults("SVC") cmd.vaults.vaults = [local] - cmd._remote_service = SimpleNamespace(server_vault_keys={}) pushed: list[dict] = [] cmd.flush_vault_writes = lambda writes: pushed.extend(w.args[0] for w in writes) dl.LICENSE_KEY_CACHE.clear() drm = SimpleNamespace(kids=[KID], content_keys={KID: GOOD, SIBLING: BAD}) - cmd.cache_resolved_keys(SimpleNamespace(tracks=[SimpleNamespace(drm=[drm])])) + cmd.cache_resolved_keys(remote_keys(), SimpleNamespace(tracks=[SimpleNamespace(drm=[drm])])) assert pushed == [{KID: GOOD}] diff --git a/tests/remote/key_source.py b/tests/remote/key_source.py new file mode 100644 index 00000000..396728eb --- /dev/null +++ b/tests/remote/key_source.py @@ -0,0 +1,23 @@ +from types import SimpleNamespace +from typing import Any, Iterable, Optional +from uuid import UUID + +from unshackle.core.remote_service import RemoteService, ServerVault + + +def remote_keys( + server_cdm: Optional[bool] = True, + server_cdm_type: str = "widevine", + client_licensed: Iterable[str] = (), + local_cdm: Any = None, + server_vault_keys: Optional[dict[UUID, str]] = None, +) -> RemoteService: + """A RemoteService with only the state dl reads through ServerKeySource, and no server.""" + svc = RemoteService.__new__(RemoteService) + svc._server_cdm = server_cdm + svc.server_cdm_type = server_cdm_type + svc.client_licensed = set(client_licensed) + svc.server_vault_keys = dict(server_vault_keys or {}) + svc.server_vault = ServerVault(svc) + svc.ctx = SimpleNamespace(obj=SimpleNamespace(cdm=local_cdm)) + return svc diff --git a/tests/remote/unit/test_remote_service_server_cdm_license.py b/tests/remote/unit/test_remote_service_server_cdm_license.py index 6401512a..890b3ab2 100644 --- a/tests/remote/unit/test_remote_service_server_cdm_license.py +++ b/tests/remote/unit/test_remote_service_server_cdm_license.py @@ -11,7 +11,7 @@ from uuid import UUID import pytest -from unshackle.core.remote_service import RemoteService +from tests.remote.key_source import remote_keys pytestmark = pytest.mark.unit @@ -46,14 +46,10 @@ class _Client: def _service(server_type, answer): - svc = RemoteService.__new__(RemoteService) + svc = remote_keys(server_cdm_type=server_type) svc._track_alias = {} svc._title_alias = {} - svc._server_cdm = True - svc._server_cdm_type = server_type - svc.client_licensed = set() svc._session_id = "sess" - svc.server_vault_keys = {} svc.log = logging.getLogger("test-remote-license") svc.client = _Client(answer) return svc @@ -86,7 +82,7 @@ def test_response_type_refreshes_the_server_type(): svc.proxy_license(b"", track, "widevine") - assert svc._server_cdm_type == "playready" + assert svc.server_cdm_type == "playready" def test_proxy_mode_keeps_the_callers_type(): diff --git a/tests/remote/unit/test_server_cdm_key_verification.py b/tests/remote/unit/test_server_cdm_key_verification.py index 9718f78b..3c2a6e6e 100644 --- a/tests/remote/unit/test_server_cdm_key_verification.py +++ b/tests/remote/unit/test_server_cdm_key_verification.py @@ -12,11 +12,12 @@ from uuid import UUID import pytest +from tests.remote.key_source import remote_keys from unshackle.core.api import handlers from unshackle.core.api.errors import APIError from unshackle.core.api.session_store import SessionEntry from unshackle.core.drm import BlowfishStripe -from unshackle.core.remote_service import RemoteService, ServerVault +from unshackle.core.remote_service import RemoteService pytestmark = pytest.mark.unit @@ -25,13 +26,10 @@ CDM_KID = UUID("43de497c6d1945f49acc2951696ec628") def remote_service() -> RemoteService: - svc = RemoteService.__new__(RemoteService) + svc = remote_keys(server_cdm=None) svc._track_alias = {} svc._title_alias = {} - svc.server_vault_keys = {} - svc.server_vault = ServerVault(svc) svc._session_id = "sess" - svc.client_licensed = set() svc.log = SimpleNamespace(warning=lambda *a, **k: None) return svc @@ -277,3 +275,15 @@ async def test_a_stripe_cipher_track_is_clear_and_never_licensed(monkeypatch): assert json.loads(resp.body)["clear_tracks"] == ["aud"] assert isinstance(track.drm[0], BlowfishStripe) + + +def test_the_server_licence_drops_a_null_key(): + kid, null = UUID(int=1), UUID(int=2) + + class Licensed: + content_keys: dict = {} + + def get_content_keys(self, **kwargs): + self.content_keys = {kid: "ab" * 16, null: "0" * 32} + + assert handlers.license_with_track_fallback(Licensed(), SimpleNamespace(drm=None)) == {kid.hex: "ab" * 16} diff --git a/tests/remote/unit/test_server_cdm_max_height.py b/tests/remote/unit/test_server_cdm_max_height.py index 4c6970eb..eff0829a 100644 --- a/tests/remote/unit/test_server_cdm_max_height.py +++ b/tests/remote/unit/test_server_cdm_max_height.py @@ -13,6 +13,7 @@ import pywidevine.pssh as wv_pssh_mod from rich.table import Table import unshackle.core.drm as drm_mod +from tests.remote.key_source import remote_keys from unshackle.commands.dl import dl from unshackle.core.api import handlers from unshackle.core.api.errors import APIError, APIErrorCode @@ -203,16 +204,11 @@ _WV.__name__ = "Widevine" def remote(local_cdm, answer): - svc = RemoteService.__new__(RemoteService) + svc = remote_keys(server_cdm_type="playready", local_cdm=local_cdm) svc._track_alias = {} svc._title_alias = {} - svc._server_cdm = True - svc._server_cdm_type = "playready" svc._session_id = "sess" - svc.server_vault_keys = {} - svc.client_licensed = set() svc.log = logging.getLogger("test-server-cdm-cap") - svc.ctx = SimpleNamespace(obj=SimpleNamespace(cdm=local_cdm)) posts: list = [] svc.client = SimpleNamespace(post=lambda endpoint, payload, expect=(): posts.append(payload) or answer) svc.drain_server_logs = lambda: None @@ -557,7 +553,7 @@ def prepare(svc, licence): def test_dl_licenses_a_client_licensed_track_with_the_local_device(): - svc = SimpleNamespace(_server_cdm=True, client_licensed={"vid"}, local_cdm="local device") + svc = remote_keys(client_licensed={"vid"}, local_cdm="local device") server_calls: list = [] @@ -566,7 +562,7 @@ def test_dl_licenses_a_client_licensed_track_with_the_local_device(): def test_dl_takes_the_local_path_when_the_server_hands_the_track_over(): - svc = SimpleNamespace(_server_cdm=True, client_licensed=set(), local_cdm="local device") + svc = remote_keys(local_cdm="local device") assert prepare(svc, lambda **kwargs: svc.client_licensed.add("vid")) == "local device" @@ -587,7 +583,7 @@ class PlayReady: def test_a_mid_download_handover_puts_the_keys_on_the_callers_drm(monkeypatch): """The caller decrypts with the DRM object it passed in, so the local licence must fill that one.""" monkeypatch.setattr(dl, "LICENSE_KEY_CACHE", {}) - svc = SimpleNamespace(_server_cdm=True, client_licensed=set(), local_cdm="local device", server_vault_keys={}) + svc = remote_keys(local_cdm="local device") runner = dl.__new__(dl) runner._remote_service = svc runner.cdm = "server stub" diff --git a/tests/remote/unit/test_server_cdm_vault_write.py b/tests/remote/unit/test_server_cdm_vault_write.py index 7f77bb50..db799a11 100644 --- a/tests/remote/unit/test_server_cdm_vault_write.py +++ b/tests/remote/unit/test_server_cdm_vault_write.py @@ -13,6 +13,7 @@ from uuid import UUID import pytest from rich.table import Table +from tests.remote.key_source import remote_keys from unshackle.commands.dl import dl pytestmark = pytest.mark.unit @@ -68,7 +69,7 @@ def make_cmd(vaults): cmd.vaults = vaults cmd.vault_cache_tally = None cmd.debug_logger = None - cmd._remote_service = SimpleNamespace(_server_cdm=True, _server_cdm_type="widevine") + cmd._remote_service = remote_keys() dl.LICENSE_KEY_CACHE.clear() dl.DRM_LOCKS.clear() return cmd @@ -78,7 +79,7 @@ def prepare(cmd, drm, licence, track_kid, **kwargs): table = Table.grid() table.add_row("") cmd.prepare_drm( - track=SimpleNamespace(drm=[drm]), + track=SimpleNamespace(id="vid", drm=[drm]), title=None, drm=drm, licence=licence, @@ -128,7 +129,7 @@ def test_batch_resolved_keys_are_cached_before_download(): drm.content_keys = {NEW: "k_new"} title = SimpleNamespace(tracks=[SimpleNamespace(drm=[drm]), SimpleNamespace(drm=None)]) - cmd.cache_resolved_keys(title) + cmd.cache_resolved_keys(remote_keys(), title) cmd.wait_vault_writes() assert vaults.added == {NEW: "k_new"} @@ -168,7 +169,7 @@ def test_server_vault_keys_are_held_back_until_verified(): drm2 = FakeWidevine([CACHED]) drm2.content_keys = {CACHED: "k_cached"} dl.LICENSE_KEY_CACHE.clear() - cmd.cache_resolved_keys(SimpleNamespace(tracks=[SimpleNamespace(drm=[drm2])])) + cmd.cache_resolved_keys(remote, SimpleNamespace(tracks=[SimpleNamespace(drm=[drm2])])) cmd.wait_vault_writes() assert vaults.added == {NEW: "k_new"} @@ -242,7 +243,7 @@ def test_no_cdm_and_no_server_cdm_names_the_missing_cdm(system): from unshackle.core.drm import PlayReady, Widevine cmd = make_cmd(FakeVaults()) - cmd._remote_service = SimpleNamespace(_server_cdm=False) + cmd._remote_service = remote_keys(server_cdm=False) cmd.cdm = None cmd.cdm_for = lambda system, quality=None: None drm_class = PlayReady if system == "playready" else Widevine @@ -250,3 +251,32 @@ def test_no_cdm_and_no_server_cdm_names_the_missing_cdm(system): drm.content_keys = {} with pytest.raises(click.ClickException, match="No CDM configured for TEST"): prepare(cmd, drm, lambda **kwargs: None, NEW) + + +NULL = UUID("0b3e2d1c4a5f46e7b8c9d0e1f2a3b4c5") + + +def test_a_null_key_from_the_server_licence_is_dropped(): + vaults = FakeVaults() + drm = run_prepare(vaults, [NEW], {NEW: "k_new", NULL: "0" * 32}, track_kid=NEW) + assert drm.content_keys == {NEW: "k_new"} + assert vaults.added == {NEW: "k_new"} + assert NULL not in dl.LICENSE_KEY_CACHE + + +def test_a_null_key_from_the_batch_licence_is_dropped(monkeypatch): + vaults = FakeVaults() + cmd = make_cmd(vaults) + drm = FakeWidevine([NEW]) + title = SimpleNamespace(tracks=[SimpleNamespace(drm=[drm])]) + source = cmd._remote_service + monkeypatch.setattr( + source, "resolve_server_keys", lambda t: drm.content_keys.update({NEW: "k_new", NULL: "0" * 32}) + ) + + cmd.take_server_keys(source, title, cdm_only=False) + cmd.wait_vault_writes() + + assert drm.content_keys == {NEW: "k_new"} + assert vaults.added == {NEW: "k_new"} + assert NULL not in dl.LICENSE_KEY_CACHE diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 5df0b800..82878758 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -23,7 +23,7 @@ from io import StringIO from itertools import product from pathlib import Path from threading import Lock -from typing import Any, Callable, Collection, Optional, Sequence, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Callable, Collection, Optional, Sequence, TypedDict, Union, cast from uuid import UUID import click @@ -70,6 +70,8 @@ from unshackle.core.drm import ( PlayReady, Widevine, cbcs, + drop_null_keys, + is_null_key, own_kids, real_kids, verify, @@ -148,6 +150,9 @@ from unshackle.core.utils.subprocess import ffmpeg_decodes, ffprobe, run_step from unshackle.core.vault import Vault from unshackle.core.vaults import Vaults +if TYPE_CHECKING: + from unshackle.core.remote_service import ServerKeySource + class SkippedSubtitle(TypedDict): """A subtitle skipped under ``--skip-subtitle-errors``. Accumulated as ``dl.skipped_subtitles``, @@ -1251,6 +1256,7 @@ class dl: export_snapshots: dict[str, Optional[dict[str, Any]]] LICENSE_KEY_CACHE: dict[UUID, str] = {} ALL_DRM_LOCK = Lock() + _remote_service: Optional["ServerKeySource"] = None def __init__( self, @@ -1724,13 +1730,12 @@ class dl: self.search_source = None self.service_anime = bool(getattr(service, "ANIME", False)) self.service_daily = bool(getattr(service, "DAILY", False)) - self.server_cdm = getattr(service, "_server_cdm", False) - from unshackle.core.import_service import ImportService + from unshackle.core.import_service import SERVER_KEY_SOURCES, ImportService # a remote session sends --cdm to the server, which can hand it the tracks above its cap if self.cdm_override and isinstance(service, ImportService): self.log.warning("--cdm is ignored: an import uses the keys in the export") - self._remote_service: Any = service if hasattr(service, "_server_cdm") else None + self._remote_service = service if isinstance(service, SERVER_KEY_SOURCES) else None start_time = time.time() lang, lang_excl = partition_exclusions(lang) @@ -3093,8 +3098,8 @@ class dl: return prepare server_cdm_type = None - if getattr(self._remote_service, "_server_cdm", False): - server_cdm_type = getattr(self._remote_service, "_server_cdm_type", "widevine") + if self._remote_service and self._remote_service.server_cdm: + server_cdm_type = self._remote_service.server_cdm_type self.cdm = cdm_type_stub(server_cdm_type) from unshackle.core.import_service import ImportService @@ -3145,11 +3150,8 @@ class dl: self.all_drm_warned.add(system) self.log.warning(f"The {system} CDM is remote and can answer with cached content keys") - if hasattr(service, "resolve_server_keys"): - service.resolve_server_keys(title) - if not cdm_only: - self.prefer_vault_keys(title) - self.cache_resolved_keys(title) + if self._remote_service: + self.take_server_keys(self._remote_service, title, bool(cdm_only)) if sample and not skip_dl: whole_files = [ @@ -4607,7 +4609,7 @@ class dl: def passes_left(done: int) -> bool: """One pass per source that can still answer: every loaded vault, the server's vaults as one source that walks them itself over --remote, and one for the CDM.""" - remote = 1 if getattr(getattr(self, "_remote_service", None), "_server_cdm", False) else 0 + remote = 1 if self._remote_service and self._remote_service.server_cdm else 0 return done < len(self.vaults) + remote + 1 try: @@ -4668,7 +4670,17 @@ class dl: self.vault_cache_tally = None self.log.info(f"Cached {keys} Key{'' if keys == 1 else 's'} to {successful_caches}/{len(self.vaults)} Vaults") - def prefer_vault_keys(self, title: Title_T) -> None: + def take_server_keys(self, source: "ServerKeySource", title: Title_T, cdm_only: bool) -> None: + """Put the content keys of ``source`` on the tracks of ``title``, drop the null keys, then cache the rest.""" + source.resolve_server_keys(title) + for track in title.tracks: + for drm in track.drm or []: + drop_null_keys(drm) + if not cdm_only: + self.prefer_vault_keys(source, title) + self.cache_resolved_keys(source, title) + + def prefer_vault_keys(self, source: "ServerKeySource", title: Title_T) -> None: """Let the client vaults answer before a content key the server batch licence returned. The client only learns the KIDs from the batch response, so the server has already @@ -4679,7 +4691,6 @@ class dl: waits in ``Vaults.candidates`` and gets the same check when its turn comes. Only the KIDs the DRM names are looked up, each once, so sibling keys add no remote vault requests. """ - server_vault_keys = getattr(self._remote_service, "server_vault_keys", {}) looked_up: dict[UUID, tuple[Optional[str], Any]] = {} for track in title.tracks: for drm in getattr(track, "drm", None) or []: @@ -4694,9 +4705,9 @@ class dl: if not vault_key or vault_key == server_key: continue drm.content_keys[kid] = vault_key - if server_vault_keys.get(kid) == server_key: - source = self._remote_service.server_vault - self.vaults.candidates.setdefault(kid, []).append((server_key, source)) + server_vault = source.vault_source(kid, server_key) + if server_vault is not None: + self.vaults.candidates.setdefault(kid, []).append((server_key, server_vault)) elif kid not in self.LICENSE_KEY_CACHE: self.LICENSE_KEY_CACHE[kid] = server_key self.flush_vault_writes([partial(self.cache_keys_to_vaults, {kid: server_key})]) @@ -4704,20 +4715,19 @@ class dl: f"{vault.name} holds a different key for {kid.hex} than the server, testing it first" ) - def cache_resolved_keys(self, title: Title_T) -> None: + def cache_resolved_keys(self, source: "ServerKeySource", title: Title_T) -> None: """Cache the keys a batch licence or an import put on the tracks. Both fill track.drm before any prepare_drm call, so this is the only point where those keys can still reach the local vaults. A content key a local vault already holds is not pushed again. """ - server_vault_keys = getattr(getattr(self, "_remote_service", None), "server_vault_keys", {}) resolved = { kid: key for track in title.tracks for drm in (getattr(track, "drm", None) or []) for kid, key in getattr(drm, "content_keys", {}).items() if kid not in self.LICENSE_KEY_CACHE - and server_vault_keys.get(kid) != key + and source.vault_source(kid, key) is None and self.vaults.sources.get(kid, (None,))[0] != key } keys = {kid: key for kid, key in resolved.items() if not self.vaults.held_locally(kid, key)} @@ -4830,6 +4840,189 @@ class dl: ) return True + @staticmethod + def with_track_kid(drm: Any, track_kid: Optional[UUID]) -> list[UUID]: + """Return the KIDs the DRM names, plus the track's own KID when the DRM leaves it out.""" + kids = list(getattr(drm, "kids", None) or []) + if track_kid and track_kid not in kids: + kids.append(track_kid) + return kids + + def take_known_keys( + self, + drm: Union[ClearKeyCENC, Widevine, PlayReady, MonaLisa], + kids: list[UUID], + track: AnyTrack, + track_kid: Optional[UUID], + cek_tree: Optional[Tree], + use_cache: bool = True, + use_vaults: bool = True, + stop_at_miss: bool = False, + ) -> list[UUID]: + """Give the DRM the content keys of ``kids`` the run cache or a key vault holds, and return the KIDs still missing. + + ``stop_at_miss`` returns at the first miss, so no more vaults are asked for a run that fails anyway. + """ + missing = [] + for kid in kids: + if kid in drm.content_keys: + if cek_tree is not None: + self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) + continue + if use_cache and self.take_cached_key(drm, kid, track, track_kid, cek_tree): + continue + if not (use_vaults and self.take_vault_key(drm, kid, track, track_kid, cek_tree)): + missing.append(kid) + if stop_at_miss: + break + return missing + + def license_drm( + self, + drm: Union[ClearKeyCENC, Widevine, PlayReady], + get_keys: Callable[[], None], + track: AnyTrack, + title: Title_T, + track_kid: Optional[UUID], + table: Table, + cdm_only: bool, + vaults_only: bool, + force: bool, + ) -> None: + """Give a Widevine, PlayReady or ClearKey DRM its content keys: run cache, then key vaults, then ``get_keys``. + + ``force`` skips the run cache, so the licence runs again for each KID no vault answers. + A licence can answer a KID with a blank or null key; a vault key for that KID wins. + """ + pssh: Optional[str] = None + if isinstance(drm, ClearKeyCENC): + name = "ClearKey" + else: + name = type(drm).__name__ + pssh = drm.pssh.dumps() if isinstance(drm, Widevine) else drm.pssh_b64 or "" + if pssh is not None and self.debug_logger: + self.debug_logger.log_drm_operation( + drm_type=name, + operation="prepare_drm", + service=self.service, + context={ + "track": str(track), + "title": str(title), + "pssh": pssh, + "kids": [k.hex for k in drm.kids], + "track_kid": track_kid.hex if track_kid else None, + }, + ) + pending_vault_writes: list[Callable[[], Any]] = [] + with self.drm_lock(drm), ExitStack() as licence_turn: + label = [(f"({self.truncate_pssh_for_display(pssh, name)})", "text")] if pssh is not None else [] + cek_tree = Tree(Text.assemble((name, "cyan"), *label, overflow="fold")) + pre_existing_tree = next( + (x for x in table.columns[0].cells if isinstance(x, Tree) and x.label == cek_tree.label), None + ) + if pre_existing_tree: + cek_tree = pre_existing_tree + + def show_error(msg: str) -> None: + cek_tree.add(f"[logging.level.error]{msg}") + if not pre_existing_tree: + table.add_row(cek_tree) + self.flush_vault_writes(pending_vault_writes) + + all_kids = self.with_track_kid(drm, track_kid) + stop_at_miss = vaults_only and not cdm_only + missing = self.take_known_keys( + drm, all_kids, track, track_kid, cek_tree, not force, not cdm_only, stop_at_miss + ) + if missing and stop_at_miss: + msg = f"No Vault has a Key for {missing[0].hex} and --vaults-only was used" + log_event( + "vault_key_not_found", + level="ERROR", + service=self.service, + message=msg, + context={"kid": missing[0].hex, "track": str(track), "drm_type": type(drm).__name__}, + ) + show_error(msg) + raise drm.Exceptions.CEKNotFound(msg) + + need_license = bool(missing) + if need_license: + # force licenses again, so the turn takes no cached keys + self.take_licence_turn(licence_turn, drm, [] if force else all_kids, track, track_kid, cek_tree) + need_license = not all(kid in drm.content_keys for kid in all_kids) + + if need_license and not vaults_only: + from_vaults = drm.content_keys.copy() + log_event( + "get_license", + level="INFO", + service=self.service, + message=f"Requesting {name} license from service", + context={ + "track": str(track), + "drm_type": type(drm).__name__, + "kids_needed": [k.hex for k in all_kids if k not in drm.content_keys], + }, + ) + try: + get_keys() + except Exception as e: + if drm.content_keys: + self.log.debug(f"License call failed but keys already in content_keys: {e}") + else: + if isinstance(e, (drm.Exceptions.EmptyLicense, drm.Exceptions.CEKNotFound)): + msg = str(e) + else: + msg = f"An exception occurred in the Service's license function: {e}" + if self.debug_logger: + self.debug_logger.log_error( + "get_license" if name == "Widevine" else f"get_license_{name.lower()}", + e, + service=self.service, + context={ + "track": str(track), + "exception_type": type(e).__name__, + "drm_type": type(drm).__name__, + }, + ) + show_error(msg) + raise + + log_event( + "license_keys_retrieved", + level="INFO", + service=self.service, + context={ + "track": str(track), + "drm_type": type(drm).__name__, + "keys_count": len(drm.content_keys), + "kids": [k.hex for k in drm.content_keys.keys()], + }, + ) + for kid, key in drm.content_keys.items(): + self.add_key_label(cek_tree, kid, f"[red]{key}[/]" if is_null_key(key) else key, track_kid) + drop_null_keys(drm) + drm.content_keys.update(from_vaults) + self.LICENSE_KEY_CACHE.update(drm.content_keys) + pending_vault_writes.append( + partial( + self.cache_keys_to_vaults, + {k: v for k, v in drm.content_keys.items() if k not in from_vaults}, + ) + ) + + if track_kid and track_kid not in drm.content_keys: + msg = f"No Content Key for KID {track_kid.hex} was returned in the License" + show_error(msg) + raise drm.Exceptions.CEKNotFound(msg) + + if cek_tree.children and not pre_existing_tree: + table.add_row() + table.add_row(cek_tree) + + self.flush_vault_writes(pending_vault_writes) + @staticmethod def add_key_label(tree: Tree, kid: UUID, key: str, track_kid: Optional[UUID] = None, source: str = "") -> None: """Add ``kid:key`` to a content key tree once, with ``*`` on the track's own KID and ``source`` after it.""" @@ -4980,22 +5173,16 @@ class dl: return assert table is not None - svc_for_cdm = getattr(self, "_remote_service", None) - server_cdm = getattr(svc_for_cdm, "_server_cdm", getattr(self, "server_cdm", False)) - client_licensed = getattr(svc_for_cdm, "client_licensed", ()) - licensed_locally = bool(client_licensed) and str(track.id) in client_licensed + remote = self._remote_service + server_cdm = bool(remote and remote.server_cdm) + licensed_locally = bool(remote and remote.licenses_locally(track)) if server_cdm and not licensed_locally: + assert remote is not None with self.drm_lock(drm), ExitStack() as licence_turn: pending_vault_writes: list[Callable[[], Any]] = [] - vault_kids = list(getattr(drm, "kids", None) or []) - if track_kid and track_kid not in vault_kids: - vault_kids.append(track_kid) - for kid in vault_kids: - if kid in drm.content_keys or self.take_cached_key(drm, kid, track, track_kid, None): - continue - if not cdm_only: - self.take_vault_key(drm, kid, track, track_kid, None) + vault_kids = self.with_track_kid(drm, track_kid) + self.take_known_keys(drm, vault_kids, track, track_kid, None, use_vaults=not cdm_only) def missing_track_key() -> bool: return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys) @@ -5016,24 +5203,22 @@ class dl: raise except (Exception, RemoteError) as e: self.log.debug(f"Server CDM licence with an empty challenge failed: {e!r}") + drop_null_keys(drm) - server_vault_keys = {} if cdm_only else getattr(svc_for_cdm, "server_vault_keys", {}) + new_keys = {} for kid, key in drm.content_keys.items(): - if server_vault_keys.get(kid) == key: - assert svc_for_cdm is not None - self.vaults.sources[kid] = (key, svc_for_cdm.server_vault) - new_keys = { - kid: key - for kid, key in drm.content_keys.items() - if kid not in known_keys and server_vault_keys.get(kid) != key - } + source = None if cdm_only else remote.vault_source(kid, key) + if source is not None: + self.vaults.sources[kid] = (key, source) + elif kid not in known_keys: + new_keys[kid] = key if new_keys: self.LICENSE_KEY_CACHE.update(new_keys) pending_vault_writes.append(partial(self.cache_keys_to_vaults, new_keys)) self.flush_vault_writes(pending_vault_writes) # the server can refuse the licence above its cap and hand the track to this machine - licensed_locally = bool(client_licensed) and str(track.id) in client_licensed + licensed_locally = bool(remote and remote.licenses_locally(track)) if server_cdm and not licensed_locally: if not drm.content_keys: self.log.warning("Server CDM did not resolve any keys for this track") @@ -5043,11 +5228,9 @@ class dl: if isinstance(drm, PlayReady): raise PlayReady.Exceptions.CEKNotFound(msg) raise Widevine.Exceptions.CEKNotFound(msg) - svc = getattr(self, "_remote_service", None) - server_drm_type = getattr(svc, "_server_cdm_type", None) if svc else None - drm_name = {"widevine": "Widevine", "playready": "PlayReady"}.get( - server_drm_type or "", drm.__class__.__name__ - ) + assert remote is not None + server_drm_type = remote.server_cdm_type + drm_name = {"widevine": "Widevine", "playready": "PlayReady"}.get(server_drm_type, drm.__class__.__name__) log_event( "drm_content_keys", level="INFO", @@ -5074,9 +5257,7 @@ class dl: cek_tree = Tree(Text.assemble((drm_name, "cyan"), (f"({pssh_str})", "text"), overflow="fold")) else: cek_tree = Tree(Text.assemble((drm_name, "cyan"), overflow="fold")) - all_kids = list(getattr(drm, "kids", [])) - if track_kid and track_kid not in all_kids: - all_kids.append(track_kid) + all_kids = self.with_track_kid(drm, track_kid) for kid in all_kids: if kid in drm.content_keys: self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) @@ -5097,7 +5278,8 @@ class dl: track_cdm = self.cdm if licensed_locally: - track_cdm = getattr(svc_for_cdm, "local_cdm", None) + assert remote is not None + track_cdm = remote.local_cdm local_drm = track.get_drm_for_cdm(track_cdm) if local_drm is not None and local_drm is not drm: assert not isinstance(local_drm, ClearKey) @@ -5137,388 +5319,14 @@ class dl: else: raise ValueError(f"Title needs a {name} CDM but {self.service} is configured with {other}.") - if isinstance(drm, Widevine): - if self.debug_logger: - self.debug_logger.log_drm_operation( - drm_type="Widevine", - operation="prepare_drm", - service=self.service, - context={ - "track": str(track), - "title": str(title), - "pssh": drm.pssh.dumps() if drm.pssh else None, - "kids": [k.hex for k in drm.kids], - "track_kid": track_kid.hex if track_kid else None, - }, + if isinstance(drm, (Widevine, PlayReady, ClearKeyCENC)): + if isinstance(drm, ClearKeyCENC): + get_keys = partial( + drm.get_content_keys, licence=clearkey_licence or (lambda **_: None), session=service_session ) - - pending_vault_writes = [] - with self.drm_lock(drm), ExitStack() as licence_turn: - pssh_display = self.truncate_pssh_for_display(drm.pssh.dumps(), "Widevine") - cek_tree = Tree(Text.assemble(("Widevine", "cyan"), (f"({pssh_display})", "text"), overflow="fold")) - pre_existing_tree = next( - (x for x in table.columns[0].cells if isinstance(x, Tree) and x.label == cek_tree.label), None - ) - if pre_existing_tree: - cek_tree = pre_existing_tree - - need_license = False - all_kids = list(drm.kids) - if track_kid and track_kid not in all_kids: - all_kids.append(track_kid) - - for kid in all_kids: - if kid in drm.content_keys: - self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) - continue - - if not force and self.take_cached_key(drm, kid, track, track_kid, cek_tree): - continue - - if cdm_only: - need_license = True - elif not self.take_vault_key(drm, kid, track, track_kid, cek_tree): - if vaults_only: - msg = f"No Vault has a Key for {kid.hex} and --vaults-only was used" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - log_event( - "vault_key_not_found", - level="ERROR", - service=self.service, - message=msg, - context={"kid": kid.hex, "track": str(track)}, - ) - self.flush_vault_writes(pending_vault_writes) - raise Widevine.Exceptions.CEKNotFound(msg) - need_license = True - - if need_license: - # force licenses again, so the turn takes no cached keys - self.take_licence_turn(licence_turn, drm, [] if force else all_kids, track, track_kid, cek_tree) - need_license = not all(kid in drm.content_keys for kid in all_kids) - - if need_license and not vaults_only: - from_vaults = drm.content_keys.copy() - - log_event( - "get_license", - level="INFO", - service=self.service, - message="Requesting Widevine license from service", - context={ - "track": str(track), - "kids_needed": [k.hex for k in all_kids if k not in drm.content_keys], - }, - ) - - try: - drm.get_content_keys(cdm=track_cdm, licence=licence, certificate=certificate) - except Exception as e: - if drm.content_keys: - self.log.debug(f"License call failed but keys already in content_keys: {e}") - else: - if isinstance(e, (Widevine.Exceptions.EmptyLicense, Widevine.Exceptions.CEKNotFound)): - msg = str(e) - else: - msg = f"An exception occurred in the Service's license function: {e}" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - if self.debug_logger: - self.debug_logger.log_error( - "get_license", - e, - service=self.service, - context={"track": str(track), "exception_type": type(e).__name__}, - ) - self.flush_vault_writes(pending_vault_writes) - raise e - - log_event( - "license_keys_retrieved", - level="INFO", - service=self.service, - context={ - "track": str(track), - "keys_count": len(drm.content_keys), - "kids": [k.hex for k in drm.content_keys.keys()], - }, - ) - - for kid_, key in drm.content_keys.items(): - if key == "0" * 32: - key = f"[red]{key}[/]" - self.add_key_label(cek_tree, kid_, key, track_kid) - - drm.content_keys = { - kid_: key for kid_, key in drm.content_keys.items() if key and key.count("0") != len(key) - } - - # The CDM keys may have returned blank content keys for KIDs we got from vaults. - # So we re-add the keys from vaults earlier overwriting blanks or removed KIDs data. - drm.content_keys.update(from_vaults) - - self.LICENSE_KEY_CACHE.update(drm.content_keys) - - pending_vault_writes.append( - partial( - self.cache_keys_to_vaults, - {k: v for k, v in drm.content_keys.items() if k not in from_vaults}, - ) - ) - - if track_kid and track_kid not in drm.content_keys: - msg = f"No Content Key for KID {track_kid.hex} was returned in the License" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - self.flush_vault_writes(pending_vault_writes) - raise Widevine.Exceptions.CEKNotFound(msg) - - if cek_tree.children and not pre_existing_tree: - table.add_row() - table.add_row(cek_tree) - - self.flush_vault_writes(pending_vault_writes) - if export: - self.write_export(export, title, track, drm) - - elif isinstance(drm, PlayReady): - if self.debug_logger: - self.debug_logger.log_drm_operation( - drm_type="PlayReady", - operation="prepare_drm", - service=self.service, - context={ - "track": str(track), - "title": str(title), - "pssh": drm.pssh_b64 or "", - "kids": [k.hex for k in drm.kids], - "track_kid": track_kid.hex if track_kid else None, - }, - ) - - pending_vault_writes = [] - with self.drm_lock(drm), ExitStack() as licence_turn: - pssh_display = self.truncate_pssh_for_display(drm.pssh_b64 or "", "PlayReady") - cek_tree = Tree( - Text.assemble( - ("PlayReady", "cyan"), - (f"({pssh_display})", "text"), - overflow="fold", - ) - ) - pre_existing_tree = next( - (x for x in table.columns[0].cells if isinstance(x, Tree) and x.label == cek_tree.label), None - ) - if pre_existing_tree: - cek_tree = pre_existing_tree - - need_license = False - all_kids = list(drm.kids) - if track_kid and track_kid not in all_kids: - all_kids.append(track_kid) - - for kid in all_kids: - if kid in drm.content_keys: - self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) - continue - - if not force and self.take_cached_key(drm, kid, track, track_kid, cek_tree): - continue - - if cdm_only: - need_license = True - elif not self.take_vault_key(drm, kid, track, track_kid, cek_tree): - if vaults_only: - msg = f"No Vault has a Key for {kid.hex} and --vaults-only was used" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - log_event( - "vault_key_not_found", - level="ERROR", - service=self.service, - message=msg, - context={"kid": kid.hex, "track": str(track), "drm_type": "PlayReady"}, - ) - self.flush_vault_writes(pending_vault_writes) - raise PlayReady.Exceptions.CEKNotFound(msg) - need_license = True - - if need_license: - # force licenses again, so the turn takes no cached keys - self.take_licence_turn(licence_turn, drm, [] if force else all_kids, track, track_kid, cek_tree) - need_license = not all(kid in drm.content_keys for kid in all_kids) - - if need_license and not vaults_only: - from_vaults = drm.content_keys.copy() - - try: - drm.get_content_keys(cdm=track_cdm, licence=licence, certificate=certificate) - except Exception as e: - if drm.content_keys: - self.log.debug(f"License call failed but keys already in content_keys: {e}") - else: - if isinstance(e, (PlayReady.Exceptions.EmptyLicense, PlayReady.Exceptions.CEKNotFound)): - msg = str(e) - else: - msg = f"An exception occurred in the Service's license function: {e}" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - if self.debug_logger: - self.debug_logger.log_error( - "get_license_playready", - e, - service=self.service, - context={ - "track": str(track), - "exception_type": type(e).__name__, - "drm_type": "PlayReady", - }, - ) - self.flush_vault_writes(pending_vault_writes) - raise e - - for kid_, key in drm.content_keys.items(): - self.add_key_label(cek_tree, kid_, key, track_kid) - - drm.content_keys.update(from_vaults) - - self.LICENSE_KEY_CACHE.update(drm.content_keys) - - pending_vault_writes.append( - partial( - self.cache_keys_to_vaults, - {k: v for k, v in drm.content_keys.items() if k not in from_vaults}, - ) - ) - - if track_kid and track_kid not in drm.content_keys: - msg = f"No Content Key for KID {track_kid.hex} was returned in the License" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - self.flush_vault_writes(pending_vault_writes) - raise PlayReady.Exceptions.CEKNotFound(msg) - - if cek_tree.children and not pre_existing_tree: - table.add_row() - table.add_row(cek_tree) - - self.flush_vault_writes(pending_vault_writes) - if export: - self.write_export(export, title, track, drm) - - elif isinstance(drm, ClearKeyCENC): - pending_vault_writes = [] - with self.drm_lock(drm), ExitStack() as licence_turn: - cek_tree = Tree(Text.assemble(("ClearKey", "cyan"), overflow="fold")) - pre_existing_tree = next( - (x for x in table.columns[0].cells if isinstance(x, Tree) and x.label == cek_tree.label), None - ) - if pre_existing_tree: - cek_tree = pre_existing_tree - - need_license = False - all_kids = list(drm.kids) - if track_kid and track_kid not in all_kids: - all_kids.append(track_kid) - - for kid in all_kids: - if kid in drm.content_keys: - self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) - continue - - if self.take_cached_key(drm, kid, track, track_kid, cek_tree): - continue - - if cdm_only: - need_license = True - elif not self.take_vault_key(drm, kid, track, track_kid, cek_tree): - if vaults_only: - msg = f"No Vault has a Key for {kid.hex} and --vaults-only was used" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - log_event( - "vault_key_not_found", - level="ERROR", - service=self.service, - message=msg, - context={"kid": kid.hex, "track": str(track), "drm_type": "ClearKeyCENC"}, - ) - self.flush_vault_writes(pending_vault_writes) - raise ClearKeyCENC.Exceptions.CEKNotFound(msg) - need_license = True - - if need_license: - # force licenses again, so the turn takes no cached keys - self.take_licence_turn(licence_turn, drm, [] if force else all_kids, track, track_kid, cek_tree) - need_license = not all(kid in drm.content_keys for kid in all_kids) - - if need_license and not vaults_only: - from_vaults = drm.content_keys.copy() - - try: - drm.get_content_keys(licence=clearkey_licence or (lambda **_: None), session=service_session) - except Exception as e: - if drm.content_keys: - self.log.debug(f"License call failed but keys already in content_keys: {e}") - else: - if isinstance( - e, (ClearKeyCENC.Exceptions.EmptyLicense, ClearKeyCENC.Exceptions.CEKNotFound) - ): - msg = str(e) - else: - msg = f"An exception occurred in the Service's license function: {e}" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - if self.debug_logger: - self.debug_logger.log_error( - "get_license_clearkey", - e, - service=self.service, - context={ - "track": str(track), - "exception_type": type(e).__name__, - "drm_type": "ClearKeyCENC", - }, - ) - self.flush_vault_writes(pending_vault_writes) - raise e - - for kid_, key in drm.content_keys.items(): - self.add_key_label(cek_tree, kid_, key, track_kid) - - drm.content_keys.update(from_vaults) - - self.LICENSE_KEY_CACHE.update(drm.content_keys) - - pending_vault_writes.append( - partial( - self.cache_keys_to_vaults, - {k: v for k, v in drm.content_keys.items() if k not in from_vaults}, - ) - ) - - if track_kid and track_kid not in drm.content_keys: - msg = f"No Content Key for KID {track_kid.hex} was returned in the License" - cek_tree.add(f"[logging.level.error]{msg}") - if not pre_existing_tree: - table.add_row(cek_tree) - self.flush_vault_writes(pending_vault_writes) - raise ClearKeyCENC.Exceptions.CEKNotFound(msg) - - if cek_tree.children and not pre_existing_tree: - table.add_row() - table.add_row(cek_tree) - - self.flush_vault_writes(pending_vault_writes) + else: + get_keys = partial(drm.get_content_keys, cdm=track_cdm, licence=licence, certificate=certificate) + self.license_drm(drm, get_keys, track, title, track_kid, table, cdm_only, vaults_only, force) if export: self.write_export(export, title, track, drm) diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 572c32a2..a8de1cf6 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -35,7 +35,7 @@ from unshackle.core.cacher import Cacher from unshackle.core.cdm.detect import cdm_type_stub from unshackle.core.config import config, load_yaml from unshackle.core.constants import AUDIO_CODEC_MAP, DYNAMIC_RANGE_MAP, VIDEO_CODEC_MAP, AnyTrack -from unshackle.core.drm import BlowfishStripe, cbcs +from unshackle.core.drm import BlowfishStripe, cbcs, is_null_key from unshackle.core.manifests.snapshot import capture from unshackle.core.providers.anilist import parse_anilist_ref from unshackle.core.proxies.resolve import initialize_proxy_providers, is_loopback, resolve_proxy @@ -4909,16 +4909,17 @@ def track_injected_keys(track: Any, drm_class: str) -> Dict[str, str]: def license_with_track_fallback(drm: Any, track: Any, **kwargs: Any) -> Dict[str, str]: - """Run `drm.get_content_keys` and merge in the keys the service put on the track itself.""" + """Run `drm.get_content_keys`, merge in the keys the service put on the track itself, and drop the null keys.""" try: drm.get_content_keys(**kwargs) except Exception: if not track_injected_keys(track, drm.__class__.__name__): raise - return { + keys = { **track_injected_keys(track, drm.__class__.__name__), **{kid.hex: key for kid, key in drm.content_keys.items()}, } + return {kid: key for kid, key in keys.items() if not is_null_key(key)} def pssh_kids(pssh_b64: str, drm_type: str) -> set: diff --git a/unshackle/core/drm/__init__.py b/unshackle/core/drm/__init__.py index 3792a13c..8196cd49 100644 --- a/unshackle/core/drm/__init__.py +++ b/unshackle/core/drm/__init__.py @@ -1,5 +1,5 @@ import base64 -from typing import Any, Union +from typing import Any, Optional, Union from uuid import UUID from unshackle.core.drm.blowfish_stripe import BlowfishStripe @@ -50,6 +50,19 @@ def drm_from_dict(data: dict[str, Any]) -> Union[Widevine, PlayReady, ClearKeyCE return drm +def is_null_key(key: Optional[str]) -> bool: + """True for a blank or all-zero content key, the null key that no vault stores.""" + return not key or key.count("0") == len(key) + + +def drop_null_keys(drm: Any) -> None: + """Remove the null keys from a DRM object's content keys. A read-only key view stays as it is.""" + content_keys = getattr(drm, "content_keys", None) + if isinstance(content_keys, dict): + for kid in [kid for kid, key in content_keys.items() if is_null_key(key)]: + del content_keys[kid] + + 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] @@ -71,6 +84,8 @@ __all__ = ( "MonaLisa", "DRM_T", "drm_from_dict", + "drop_null_keys", + "is_null_key", "real_kids", "own_kids", ) diff --git a/unshackle/core/import_service.py b/unshackle/core/import_service.py index d3312969..c4587ead 100644 --- a/unshackle/core/import_service.py +++ b/unshackle/core/import_service.py @@ -20,6 +20,7 @@ from unshackle.core.drm import ClearKey, MonaLisa, drm_from_dict, real_kids from unshackle.core.manifests import DASH, HLS, ISM, MissingLanguageError, snapshot from unshackle.core.remote_service import ( RemoteService, + ServerKeySource, build_title, resolve_proxy_arg, segment_uri_filter, @@ -207,15 +208,16 @@ class ImportService: """Reconstructs a download from an export JSON. ImportService does not authenticate and does not license. It rebuilds the tracks from - the export and injects the keys directly. ``_server_cdm``/``_server_cdm_type`` keep their - underscores: dl.py reads them through getattr as the server-CDM contract that skips - client licensing. + the export and injects the keys directly. It is a ``ServerKeySource``, so dl takes the keys + from the tracks and licenses nothing itself. """ ALIASES: tuple[str, ...] = () GEOFENCE: tuple[str, ...] = () GEOBLOCK: tuple[str, ...] = () NO_SUBTITLES: bool = False + server_cdm = True + local_cdm = None def __init__(self, ctx: click.Context, service_tag: str, title: str, import_file: Optional[str]) -> None: self.__class__.__name__ = service_tag @@ -247,8 +249,7 @@ class ImportService: self.titles: Optional[Titles_T] = None self.tracks_by_title: dict[str, Tracks] = {} - self._server_cdm = True - self._server_cdm_type = "playready" if self.exported_drm_system() == "playready" else "widevine" + self.server_cdm_type = "playready" if self.exported_drm_system() == "playready" else "widevine" first = self.doc.titles[0] # a title with no manifest keeps its request headers on its file rows @@ -675,6 +676,14 @@ class ImportService: drm_obj.content_keys[UUID(hex=kid_hex)] = key return [drm_obj] + def licenses_locally(self, track: AnyTrack) -> bool: + """Always False: an import licenses nothing.""" + return False + + def vault_source(self, kid: UUID, key: str) -> None: + """Always None: no server vault serves an import.""" + return None + def resolve_server_keys(self, title: Title_T) -> None: """Inject exported keys into the selected encrypted tracks by KID (no network). @@ -716,7 +725,7 @@ class ImportService: 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() + self.server_cdm_type = drm_obj.__class__.__name__.lower() @classmethod def track_is_encrypted(cls, track: Any, session: Optional[requests.Session] = None) -> bool: @@ -850,3 +859,7 @@ class ImportService: self.session.close() except Exception: pass + + +SERVER_KEY_SOURCES: tuple[type[ServerKeySource], ...] = (RemoteService, ImportService) +"""The services dl takes content keys from before ``prepare_drm``; mypy checks each against ServerKeySource.""" diff --git a/unshackle/core/remote_service.py b/unshackle/core/remote_service.py index 806d6045..5d798d6b 100644 --- a/unshackle/core/remote_service.py +++ b/unshackle/core/remote_service.py @@ -24,7 +24,7 @@ from enum import Enum from http.cookiejar import CookieJar from pathlib import Path from threading import Event, Lock, Thread, local -from typing import Any, Callable, Dict, Iterable, Iterator, Optional, TypeVar, Union +from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Protocol, TypeVar, Union from urllib.parse import urlencode from uuid import UUID @@ -916,6 +916,31 @@ def cache_key_filter(service_tag: str, profile: Optional[str]) -> Callable[[str] return relevant +class ServerKeySource(Protocol): + """A service whose content keys dl gets before ``prepare_drm``: a remote session or an import.""" + + @property + def server_cdm(self) -> bool: + """True when the keys come from the server CDM or the export, so dl licenses with no CDM of its own.""" + + @property + def server_cdm_type(self) -> str: + """``widevine`` or ``playready``: the DRM system the keys came from.""" + + @property + def local_cdm(self) -> Any: + """The CDM of this machine for a track ``licenses_locally`` names, or None.""" + + def resolve_server_keys(self, title: Title_T) -> None: + """Put the content keys the server or the export holds on the tracks of ``title``.""" + + def licenses_locally(self, track: AnyTrack) -> bool: + """True when this machine's own CDM licenses ``track`` instead of the server.""" + + def vault_source(self, kid: UUID, key: str) -> Optional[Vault]: + """The server vault when it served ``key`` for ``kid``, else None. That content key must pass a decrypt check.""" + + class ServerVault(Vault): """The server's vaults as one source seen from the client: it can only be reported against. @@ -963,6 +988,7 @@ class RemoteService: NO_SUBTITLES: bool = False ANIME: bool = False DAILY: bool = False + server_cdm_type: str = "widevine" def __init__( self, @@ -992,7 +1018,6 @@ class RemoteService: self._tracks_by_title: Dict[str, Tracks] = {} self._chapters_by_title: Dict[str, list] = {} self._session_id: Optional[str] = None - self._server_cdm_type: str = "widevine" 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) @@ -1701,7 +1726,7 @@ class RemoteService: elif self._server_cdm and not server_cdm: self._server_cdm = False self.report_local_fallback() - self._server_cdm_type = result.get("server_cdm_type", "widevine") + self.server_cdm_type = result.get("server_cdm_type", "widevine") self._sample_decrypt_tracks.update(result.get("sample_decrypt_tracks") or []) for track in [*tracks.videos, *tracks.audio, *tracks.subtitles]: @@ -1778,7 +1803,7 @@ class RemoteService: if not keyed: return - drm_type = getattr(self, "_server_cdm_type", "widevine") + drm_type = self.server_cdm_type if getattr(self, "_server_vault", False): from unshackle.core.cdm.detect import is_playready_cdm @@ -1808,7 +1833,7 @@ class RemoteService: server_drm_type = resp.get("drm_type", drm_type) drm_types_by_track = resp.get("drm_types", {}) capped = resp.get("capped_tracks") or {} - self._server_cdm_type = server_drm_type + self.server_cdm_type = server_drm_type self.log.debug(f"Server responded with drm_type={server_drm_type}, keys for {len(keys_by_track)} track(s)") for track in title.tracks: @@ -1877,6 +1902,19 @@ class RemoteService: self.client_licensed.add(str(track.id)) self.log.info(f"{why}; licensing {track.id} with your own {drm_class} device") + @property + def server_cdm(self) -> bool: + """True when the server CDM licenses this remote session, so dl licenses with no CDM of its own.""" + return bool(self._server_cdm) + + def licenses_locally(self, track: AnyTrack) -> bool: + """True when the server handed ``track`` to this machine's own CDM to license.""" + return str(track.id) in self.client_licensed + + def vault_source(self, kid: UUID, key: str) -> Optional[Vault]: + """The server vault when it served ``key`` for ``kid``, else None.""" + return self.server_vault if self.server_vault_keys.get(kid) == key else None + def note_vault_keys(self, keys: Dict[str, str], vault_kids: set[str]) -> None: """Remember which of the served keys a server vault supplied, and forget the rest.""" for kid_hex, key in keys.items(): @@ -1991,7 +2029,7 @@ class RemoteService: if self._server_cdm and str(track.id) not in self.client_licensed: from uuid import UUID - server_type = self._server_cdm_type + server_type = self.server_cdm_type other_type = "playready" if server_type == "widevine" else "widevine" drm_type, pssh_b64 = server_type, track_pssh(track, server_type) if not pssh_b64: @@ -2009,7 +2047,7 @@ class RemoteService: ) if resp.get("error_code") == "SERVER_CDM_CAPPED": refusal = resp.get("details") or {} - self._server_cdm_type = resp.get("drm_type", server_type) + self.server_cdm_type = resp.get("drm_type", server_type) keys = resp.get("keys", {}) self.note_vault_keys(keys, set(resp.get("vault_keys", []))) if keys and track.drm: