From c6c03310070b486a8e7c07cdc5218e8717d1a045 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sat, 31 Oct 2026 13:28:29 -0600 Subject: [PATCH] refactor(dl): share the cache, vault and key-label steps in prepare_drm take_cached_key, take_vault_key and add_key_label replace the code copied across the server-CDM, Widevine, PlayReady, ClearKey and MonaLisa branches. --- tests/core/test_licence_turns.py | 19 +- .../unit/test_server_cdm_vault_write.py | 2 + unshackle/commands/dl.py | 279 +++++++----------- 3 files changed, 117 insertions(+), 183 deletions(-) diff --git a/tests/core/test_licence_turns.py b/tests/core/test_licence_turns.py index 2913e8d1..1b2be4b7 100644 --- a/tests/core/test_licence_turns.py +++ b/tests/core/test_licence_turns.py @@ -12,6 +12,7 @@ from uuid import UUID import pytest from pywidevine.pssh import PSSH from rich.table import Table +from rich.tree import Tree from unshackle.commands.dl import dl from unshackle.core.cdm.detect import cdm_type_stub @@ -107,7 +108,7 @@ def run_together(cmd: dl, *jobs: tuple[Widevine, UUID]) -> tuple[Table, list[Bas def tree_labels(table: Table, kid: UUID) -> list[str]: - trees = [cell for cell in table.columns[0].cells if hasattr(cell, "children")] + trees = [cell for cell in table.columns[0].cells if isinstance(cell, Tree)] return [str(child.label) for tree in trees for child in tree.children if kid.hex in str(child.label)] @@ -119,8 +120,7 @@ def test_one_licence_for_two_pssh_with_different_kids(cmd: dl, licences: list[st assert not errors assert len(licences) == 1 assert first.content_keys[KID_A] == KEYS[KID_A] and second.content_keys[KID_C] == KEYS[KID_C] - from_cache = [label for kid in (KID_A, KID_C) for label in tree_labels(table, kid) if label.endswith(" from cache")] - assert from_cache + assert any(label.endswith(" from cache") for kid in (KID_A, KID_C) for label in tree_labels(table, kid)) def test_a_track_waiting_for_its_turn_sends_no_licence_after_a_failure(cmd: dl, licences: list[str]) -> None: @@ -138,3 +138,16 @@ def test_force_licenses_again_although_the_keys_are_cached(cmd: dl, licences: li prepare(cmd, widevine(KID_C), KID_C, new_table(), force=True) assert len(licences) == 1 + + +def test_a_key_label_is_added_once_with_its_marker_and_source() -> None: + tree = Tree("Widevine") + + dl.add_key_label(tree, KID_A, KEYS[KID_A], KID_A, "from cache") + dl.add_key_label(tree, KID_A, KEYS[KID_A]) + dl.add_key_label(tree, KID_B, KEYS[KID_B], KID_A) + + assert [str(child.label) for child in tree.children] == [ + f"[text2]{KID_A.hex}:{KEYS[KID_A]}* from cache", + f"[text2]{KID_B.hex}:{KEYS[KID_B]}", + ] diff --git a/tests/remote/unit/test_server_cdm_vault_write.py b/tests/remote/unit/test_server_cdm_vault_write.py index 4773f175..7f77bb50 100644 --- a/tests/remote/unit/test_server_cdm_vault_write.py +++ b/tests/remote/unit/test_server_cdm_vault_write.py @@ -67,6 +67,7 @@ def make_cmd(vaults): cmd.service = "TEST" cmd.vaults = vaults cmd.vault_cache_tally = None + cmd.debug_logger = None cmd._remote_service = SimpleNamespace(_server_cdm=True, _server_cdm_type="widevine") dl.LICENSE_KEY_CACHE.clear() dl.DRM_LOCKS.clear() @@ -120,6 +121,7 @@ def test_batch_resolved_keys_are_cached_before_download(): cmd.service = "TEST" cmd.vaults = vaults cmd.vault_cache_tally = None + cmd.debug_logger = None dl.LICENSE_KEY_CACHE.clear() drm = FakeWidevine([NEW]) diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 56222b37..eaf7487c 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -4785,18 +4785,74 @@ class dl: if DOWNLOAD_CANCELLED.is_set(): raise DownloadCancelled() for kid in kids: - key = self.LICENSE_KEY_CACHE.get(kid) - if not key or kid in drm.content_keys: - continue - drm.content_keys[kid] = key - if cek_tree is not None and not any(f"{kid.hex}:{key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(f"[text2]{kid.hex}:{key}{['', '*'][kid == track_kid]} from cache") - log_event( - "license_cache_hit", - level="INFO", + if kid not in drm.content_keys: + self.take_cached_key(drm, kid, track, track_kid, cek_tree) + + def take_cached_key( + self, + drm: Union[ClearKeyCENC, Widevine, PlayReady, MonaLisa], + kid: UUID, + track: AnyTrack, + track_kid: Optional[UUID], + cek_tree: Optional[Tree], + ) -> bool: + """Give the DRM the content key LICENSE_KEY_CACHE holds for ``kid``, and return whether it held one.""" + key = self.LICENSE_KEY_CACHE.get(kid) + if not key: + return False + drm.content_keys[kid] = key + if cek_tree is not None: + self.add_key_label(cek_tree, kid, key, track_kid, "from cache") + log_event( + "license_cache_hit", + level="INFO", + service=self.service, + context={"kid": kid.hex, "content_key": key, "track": str(track), "drm_type": type(drm).__name__}, + ) + return True + + def take_vault_key( + self, + drm: Union[ClearKeyCENC, Widevine, PlayReady, MonaLisa], + kid: UUID, + track: AnyTrack, + track_kid: Optional[UUID], + cek_tree: Optional[Tree], + ) -> bool: + """Give the DRM the content key a key vault holds for ``kid``, cache it for the run, and return whether one did.""" + key, vault = self.vaults.get_key(kid) + if not key: + return False + drm.content_keys[kid] = key + self.LICENSE_KEY_CACHE[kid] = key + if cek_tree is not None: + self.add_key_label(cek_tree, kid, key, track_kid, f"from {vault}") + if self.debug_logger: + self.debug_logger.log_vault_query( + vault_name=str(vault), + operation="get_key_success", service=self.service, - context={"kid": kid.hex, "content_key": key, "track": str(track), "drm_type": type(drm).__name__}, + context={ + "kid": kid.hex, + "content_key": key, + "track": str(track), + "from_cache": True, + "drm_type": type(drm).__name__, + }, ) + return True + + @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.""" + if any(f"{kid.hex}:{key}" in str(x.label) for x in tree.children): + return + label = f"[text2]{kid.hex}:{key}" + if kid == track_kid: + label += "*" + if source: + label += f" {source}" + tree.add(label) def service_licence(self, service: Service, drm_system: Optional[str] = None, **kwargs: Any) -> Any: """Call the service's licence function for the track's DRM system. @@ -4948,14 +5004,10 @@ class dl: 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: + if kid in drm.content_keys or self.take_cached_key(drm, kid, track, track_kid, None): continue - content_key = self.LICENSE_KEY_CACHE.get(kid) - if not content_key and not cdm_only: - content_key, vault_used = self.vaults.get_key(kid) - if content_key: - drm.content_keys[kid] = content_key - self.LICENSE_KEY_CACHE[kid] = content_key + if not cdm_only: + self.take_vault_key(drm, kid, track, track_kid, None) def missing_track_key() -> bool: return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys) @@ -5039,12 +5091,10 @@ class dl: all_kids.append(track_kid) for kid in all_kids: if kid in drm.content_keys: - is_track_kid = ["", "*"][kid == track_kid] - key = drm.content_keys[kid] - cek_tree.add(f"[text2]{kid.hex}:{key}{is_track_kid}") + self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) for kid, key in drm.content_keys.items(): if kid not in all_kids: - cek_tree.add(f"[text2]{kid.hex}:{key}") + self.add_key_label(cek_tree, kid, key) if not any(isinstance(x, Tree) and x.label == cek_tree.label for x in table.columns[0].cells): table.add_row(cek_tree) if export: @@ -5131,56 +5181,16 @@ class dl: for kid in all_kids: if kid in drm.content_keys: - is_track_kid = ["", "*"][kid == track_kid] - key = drm.content_keys[kid] - label = f"[text2]{kid.hex}:{key}{is_track_kid}" - if not any(f"{kid.hex}:{key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) + self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) continue - is_track_kid = ["", "*"][kid == track_kid] - - cached_key = None if force else self.LICENSE_KEY_CACHE.get(kid) - if cached_key: - drm.content_keys[kid] = cached_key - label = f"[text2]{kid.hex}:{cached_key}{is_track_kid} from cache" - if not any(f"{kid.hex}:{cached_key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) - log_event( - "license_cache_hit", - level="INFO", - service=self.service, - context={ - "kid": kid.hex, - "content_key": cached_key, - "track": str(track), - "drm_type": "Widevine", - }, - ) + if not force and self.take_cached_key(drm, kid, track, track_kid, cek_tree): continue - if not cdm_only: - content_key, vault_used = self.vaults.get_key(kid) - if content_key: - drm.content_keys[kid] = content_key - label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" - if not any(f"{kid.hex}:{content_key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) - self.LICENSE_KEY_CACHE[kid] = content_key - - if self.debug_logger: - self.debug_logger.log_vault_query( - vault_name=str(vault_used), - operation="get_key_success", - service=self.service, - context={ - "kid": kid.hex, - "content_key": content_key, - "track": str(track), - "from_cache": True, - }, - ) - elif vaults_only: + 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: @@ -5194,16 +5204,12 @@ class dl: ) self.flush_vault_writes(pending_vault_writes) raise Widevine.Exceptions.CEKNotFound(msg) - else: - need_license = True - - if kid not in drm.content_keys and cdm_only: 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) - if need_license and all(kid in drm.content_keys for kid in all_kids): - need_license = False + 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() @@ -5256,10 +5262,7 @@ class dl: for kid_, key in drm.content_keys.items(): if key == "0" * 32: key = f"[red]{key}[/]" - is_track_kid_marker = ["", "*"][kid_ == track_kid] - label = f"[text2]{kid_.hex}:{key}{is_track_kid_marker}" - if not any(f"{kid_.hex}:{key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) + 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) @@ -5332,57 +5335,16 @@ class dl: for kid in all_kids: if kid in drm.content_keys: - is_track_kid = ["", "*"][kid == track_kid] - key = drm.content_keys[kid] - label = f"[text2]{kid.hex}:{key}{is_track_kid}" - if not any(f"{kid.hex}:{key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) + self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) continue - is_track_kid = ["", "*"][kid == track_kid] - - cached_key = None if force else self.LICENSE_KEY_CACHE.get(kid) - if cached_key: - drm.content_keys[kid] = cached_key - label = f"[text2]{kid.hex}:{cached_key}{is_track_kid} from cache" - if not any(f"{kid.hex}:{cached_key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) - log_event( - "license_cache_hit", - level="INFO", - service=self.service, - context={ - "kid": kid.hex, - "content_key": cached_key, - "track": str(track), - "drm_type": "PlayReady", - }, - ) + if not force and self.take_cached_key(drm, kid, track, track_kid, cek_tree): continue - if not cdm_only: - content_key, vault_used = self.vaults.get_key(kid) - if content_key: - drm.content_keys[kid] = content_key - label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" - if not any(f"{kid.hex}:{content_key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) - self.LICENSE_KEY_CACHE[kid] = content_key - - if self.debug_logger: - self.debug_logger.log_vault_query( - vault_name=str(vault_used), - operation="get_key_success", - service=self.service, - context={ - "kid": kid.hex, - "content_key": content_key, - "track": str(track), - "from_cache": True, - "drm_type": "PlayReady", - }, - ) - elif vaults_only: + 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: @@ -5396,16 +5358,12 @@ class dl: ) self.flush_vault_writes(pending_vault_writes) raise PlayReady.Exceptions.CEKNotFound(msg) - else: - need_license = True - - if kid not in drm.content_keys and cdm_only: 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) - if need_license and all(kid in drm.content_keys for kid in all_kids): - need_license = False + 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() @@ -5438,10 +5396,7 @@ class dl: raise e for kid_, key in drm.content_keys.items(): - is_track_kid_marker = ["", "*"][kid_ == track_kid] - label = f"[text2]{kid_.hex}:{key}{is_track_kid_marker}" - if not any(f"{kid_.hex}:{key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) + self.add_key_label(cek_tree, kid_, key, track_kid) drm.content_keys.update(from_vaults) @@ -5487,43 +5442,16 @@ class dl: for kid in all_kids: if kid in drm.content_keys: - is_track_kid = ["", "*"][kid == track_kid] - key = drm.content_keys[kid] - label = f"[text2]{kid.hex}:{key}{is_track_kid}" - if not any(f"{kid.hex}:{key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) + self.add_key_label(cek_tree, kid, drm.content_keys[kid], track_kid) continue - is_track_kid = ["", "*"][kid == track_kid] - - cached_key = self.LICENSE_KEY_CACHE.get(kid) - if cached_key: - drm.content_keys[kid] = cached_key - label = f"[text2]{kid.hex}:{cached_key}{is_track_kid} from cache" - if not any(f"{kid.hex}:{cached_key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) - log_event( - "license_cache_hit", - level="INFO", - service=self.service, - context={ - "kid": kid.hex, - "content_key": cached_key, - "track": str(track), - "drm_type": "ClearKeyCENC", - }, - ) + if self.take_cached_key(drm, kid, track, track_kid, cek_tree): continue - if not cdm_only: - content_key, vault_used = self.vaults.get_key(kid) - if content_key: - drm.content_keys[kid] = content_key - label = f"[text2]{kid.hex}:{content_key}{is_track_kid} from {vault_used}" - if not any(f"{kid.hex}:{content_key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) - self.LICENSE_KEY_CACHE[kid] = content_key - elif vaults_only: + 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: @@ -5537,16 +5465,12 @@ class dl: ) self.flush_vault_writes(pending_vault_writes) raise ClearKeyCENC.Exceptions.CEKNotFound(msg) - else: - need_license = True - - if kid not in drm.content_keys and cdm_only: 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) - if need_license and all(kid in drm.content_keys for kid in all_kids): - need_license = False + 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() @@ -5581,10 +5505,7 @@ class dl: raise e for kid_, key in drm.content_keys.items(): - is_track_kid_marker = ["", "*"][kid_ == track_kid] - label = f"[text2]{kid_.hex}:{key}{is_track_kid_marker}" - if not any(f"{kid_.hex}:{key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) + self.add_key_label(cek_tree, kid_, key, track_kid) drm.content_keys.update(from_vaults) @@ -5625,9 +5546,7 @@ class dl: cek_tree = pre_existing_tree for kid_, key in drm.content_keys.items(): - label = f"[text2]{kid_.hex}:{key}*" - if not any(f"{kid_.hex}:{key}" in str(x.label) for x in cek_tree.children): - cek_tree.add(label) + self.add_key_label(cek_tree, kid_, key, track_kid=kid_) # every MonaLisa key is the track's own if cek_tree.children and not pre_existing_tree: table.add_row()