diff --git a/docs/guide/vaults.md b/docs/guide/vaults.md index 3a2968c4..294e2439 100644 --- a/docs/guide/vaults.md +++ b/docs/guide/vaults.md @@ -30,6 +30,8 @@ When unshackle recovers new keys from a license, it pushes them to **all** confi - unshackle skips keys that already exist in a vault (it never overwrites or deletes existing data). - A vault marked `no_push: true` receives lookups, but unshackle never writes to it. This is useful for read-only or shared upstream vaults you do not own. +unshackle pushes keys from every source, not only a local CDM. A run against a remote server that licenses with its own CDM (`--remote`) writes the keys the server returns, and `unshackle import` writes the keys it reads from an export file. + !!! tip "A good two-vault setup" A common arrangement is a fast local **SQLite** vault plus a shared remote vault (**MySQL**, **HTTP**, or **API**). The local vault answers instantly and works offline; the remote vault lets you share keys across machines or with a group. diff --git a/tests/remote/unit/test_server_cdm_vault_write.py b/tests/remote/unit/test_server_cdm_vault_write.py new file mode 100644 index 00000000..1e8a108c --- /dev/null +++ b/tests/remote/unit/test_server_cdm_vault_write.py @@ -0,0 +1,173 @@ +"""Keys the server CDM returns must land in the local vaults. + +The server-CDM branch of prepare_drm reads the local vaults but used to return before +the vault writes, so a remote run never kept the keys it paid for. Keys that came from a +vault in the first place must not be pushed back. +""" + +import threading +import time +from types import SimpleNamespace +from uuid import UUID + +import pytest +from rich.table import Table + +from unshackle.commands.dl import dl + +pytestmark = pytest.mark.unit + +NEW = UUID("161219ec3df64e0eadbb3827c0ccc46d") +CACHED = UUID("75982fac2cd8412a990a59c0ef2d4bb1") + + +class FakeVaults: + def __init__(self, held=None, vault_used=None): + self.held = held or {} + self.vault_used = vault_used + self.added: dict = {} + self.pushes: list = [] + self.replicated: list = [] + + def __len__(self): + return 1 + + def get_key(self, kid): + return self.held.get(kid), self.vault_used + + def add_keys(self, kid_keys): + self.added.update(kid_keys) + self.pushes.append(dict(kid_keys)) + return 1 + + def add_key(self, kid, key, excluding=None): + self.replicated.append((kid, key, excluding)) + return 1 + + +class FakeWidevine: + def __init__(self, kids): + self.kids = list(kids) + self.content_keys: dict = {} + + +def make_cmd(vaults): + cmd = dl.__new__(dl) + cmd.log = SimpleNamespace(warning=print, info=print, debug=print) + cmd.service = "TEST" + cmd.vaults = vaults + cmd.vault_cache_tally = None + cmd._remote_service = SimpleNamespace(_server_cdm=True, _server_cdm_type="widevine") + dl.LICENSE_KEY_CACHE.clear() + dl.DRM_LOCKS.clear() + return cmd + + +def prepare(cmd, drm, licence, track_kid, **kwargs): + table = Table.grid() + table.add_row("") + cmd.prepare_drm( + track=SimpleNamespace(drm=[drm]), + title=None, + drm=drm, + licence=licence, + certificate=None, + track_kid=track_kid, + table=table, + **kwargs, + ) + + +def run_prepare(vaults, kids, licence_keys, track_kid): + cmd = make_cmd(vaults) + drm = FakeWidevine(kids) + + def licence(**kwargs): + drm.content_keys.update(licence_keys) + + prepare(cmd, drm, licence, track_kid) + cmd.wait_vault_writes() + return drm + + +def test_server_cdm_keys_are_cached_to_local_vaults(): + vaults = FakeVaults() + run_prepare(vaults, [NEW], {NEW: "k_new"}, track_kid=NEW) + assert vaults.added == {NEW: "k_new"} + + +def test_keys_read_from_a_vault_are_not_pushed_back(): + vaults = FakeVaults({CACHED: "k_cached"}) + run_prepare(vaults, [CACHED, NEW], {NEW: "k_new"}, track_kid=NEW) + assert vaults.added == {NEW: "k_new"} + + +def test_batch_resolved_keys_are_cached_before_download(): + """The batch licence fills track.drm before prepare_drm, so dl caches it at the call site.""" + vaults = FakeVaults() + cmd = dl.__new__(dl) + cmd.log = SimpleNamespace(warning=print, info=print, debug=print) + cmd.service = "TEST" + cmd.vaults = vaults + cmd.vault_cache_tally = None + dl.LICENSE_KEY_CACHE.clear() + + drm = FakeWidevine([NEW]) + drm.content_keys = {NEW: "k_new"} + title = SimpleNamespace(tracks=[SimpleNamespace(drm=[drm]), SimpleNamespace(drm=None)]) + + cmd.cache_resolved_keys(title) + cmd.wait_vault_writes() + + assert vaults.added == {NEW: "k_new"} + + +def test_vault_hits_replicate_to_the_other_vaults(): + """The local path pushes a vault hit to every other vault; the server path must too.""" + vaults = FakeVaults({CACHED: "k_cached"}, vault_used="vault-a") + run_prepare(vaults, [CACHED], {}, track_kid=CACHED) + assert vaults.replicated == [(CACHED, "k_cached", "vault-a")] + assert vaults.added == {} + + +def test_tracks_sharing_one_drm_licence_and_push_once(): + """Two tracks on the same DRM run at the same time; the second must find the first one's keys.""" + vaults = FakeVaults() + cmd = make_cmd(vaults) + drm = FakeWidevine([NEW]) + calls = [] + + def licence(**kwargs): + calls.append(1) + time.sleep(0.2) + drm.content_keys.update({NEW: "k_new"}) + + threads = [threading.Thread(target=prepare, args=(cmd, drm, licence, NEW)) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + cmd.wait_vault_writes() + + assert len(calls) == 1 + assert vaults.pushes == [{NEW: "k_new"}] + + +def test_cdm_only_still_reuses_keys_from_the_run_cache(): + """--cdm-only skips the vaults, not LICENSE_KEY_CACHE, so a second track with the same KID does not licence twice.""" + vaults = FakeVaults() + cmd = make_cmd(vaults) + calls = [] + + for _ in range(2): + drm = FakeWidevine([NEW]) + + def licence(**kwargs): + calls.append(1) + drm.content_keys.update({NEW: "k_new"}) + + prepare(cmd, drm, licence, NEW, cdm_only=True) + cmd.wait_vault_writes() + + assert len(calls) == 1 + assert vaults.pushes == [{NEW: "k_new"}] diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index bed700a4..4c2582d8 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -2904,6 +2904,7 @@ class dl: if hasattr(service, "resolve_server_keys"): service.resolve_server_keys(title) + self.cache_resolved_keys(title) dl_start_time = time.time() @@ -3996,6 +3997,23 @@ 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 cache_resolved_keys(self, 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. + """ + keys = { + 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 + } + if keys: + self.LICENSE_KEY_CACHE.update(keys) + self.flush_vault_writes([partial(self.cache_keys_to_vaults, keys)]) + def cache_keys_to_vaults(self, content_keys: dict[UUID, str]) -> None: successful_caches = self.vaults.add_keys(content_keys) keys, caches = self.vault_cache_tally or (0, successful_caches) @@ -4054,29 +4072,44 @@ class dl: server_cdm = getattr(svc_for_cdm, "_server_cdm", getattr(self, "server_cdm", False)) if server_cdm: - if not cdm_only: + with self.drm_lock(drm): + 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: continue - content_key = self.LICENSE_KEY_CACHE.get(kid) or self.vaults.get_key(kid)[0] + 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: + pending_vault_writes.append( + partial(self.vaults.add_key, kid, content_key, excluding=vault_used) + ) if content_key: drm.content_keys[kid] = content_key self.LICENSE_KEY_CACHE[kid] = content_key - def missing_track_key() -> bool: - return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys) + known_keys = set(drm.content_keys) - if missing_track_key(): - try: - licence( - drm_system="playready" if drm.__class__.__name__ == "PlayReady" else "widevine", - challenge=b"", - ) - except Exception as e: - self.log.debug(f"Server CDM licence with an empty challenge failed: {e!r}") + def missing_track_key() -> bool: + return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys) + + if missing_track_key(): + try: + licence( + drm_system="playready" if drm.__class__.__name__ == "PlayReady" else "widevine", + challenge=b"", + ) + except Exception as e: + self.log.debug(f"Server CDM licence with an empty challenge failed: {e!r}") + + new_keys = {kid: key for kid, key in drm.content_keys.items() if kid not in known_keys} + 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) if not drm.content_keys: self.log.warning("Server CDM did not resolve any keys for this track")