unshackle/tests/remote/unit/test_server_cdm_vault_write.py

"""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"}]