feat(drm): license a track only when its own KID has no key

A vault key for the track's KID now stands in for the licence on the client and the server (PlayReady included). decrypt_verified licenses any other header KID the file uses.

tests/core/test_licence_turns.py

@@ -72,7 +72,7 @@ def widevine(*kids: UUID) -> Widevine:
     return Widevine(pssh=PSSH.new(system_id=PSSH.SystemId.Widevine, key_ids=list(kids)))
 
 
-def prepare(cmd: dl, drm: Widevine, track_kid: UUID, table: Table, **kwargs: Any) -> None:
+def prepare(cmd: dl, drm: Widevine, track_kid: Optional[UUID], table: Table, **kwargs: Any) -> None:
     cmd.prepare_drm(
         drm,
         track=SimpleNamespace(drm=[drm]),
@@ -220,7 +220,25 @@ def test_a_refused_licence_without_the_track_key_reports_the_licence_error(
     assert "Licence request failed" in str(tree.children[-1].label)
 
 
-def test_a_refused_licence_keeps_the_known_track_key(
+def test_a_held_track_key_needs_no_licence(cmd: dl, licences: list[str], monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(dl, "LICENSE_KEY_CACHE", {KID_A: KEYS[KID_A]})
+    drm = widevine(KID_A, KID_B)
+
+    prepare(cmd, drm, KID_A, new_table())
+
+    assert licences == []
+    assert drm.content_keys == {KID_A: KEYS[KID_A]}
+
+
+def test_force_licenses_a_held_track_key_again(cmd: dl, licences: list[str], monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(dl, "LICENSE_KEY_CACHE", {KID_A: KEYS[KID_A]})
+
+    prepare(cmd, widevine(KID_A, KID_B), KID_A, new_table(), force=True)
+
+    assert len(licences) == 1
+
+
+def test_a_refused_licence_keeps_the_known_keys_when_the_track_kid_is_unknown(
     cmd: dl, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
 ) -> None:
     monkeypatch.setattr(Widevine, "get_content_keys", refuse_licence)
@@ -228,7 +246,7 @@ def test_a_refused_licence_keeps_the_known_track_key(
     monkeypatch.setattr(dl, "DRM_LOCKS", {})
     drm = widevine(KID_A, KID_B)
 
-    prepare(cmd, drm, KID_A, new_table())
+    prepare(cmd, drm, None, new_table())
 
     assert drm.content_keys == {KID_A: KEYS[KID_A]}
     assert [r.getMessage() for r in caplog.records if r.levelname == "WARNING"] == [

tests/core/test_mp4decrypt_track_key.py

@@ -131,7 +131,18 @@ def test_segment_decrypt_needs_a_key_for_a_kidless_track(
     drm.content_keys[key_kid] = KEY
 
     assert can_use(drm, "mp4decrypt", movie(two_sinf(A, piff_first=True))) is usable
-    assert can_use(drm, "mp4decrypt", movie(encv(A))) is True
+    assert can_use(drm, "mp4decrypt", movie(encv(A))) is usable
+
+
+def test_segment_decrypt_waits_for_a_key_to_every_kid_the_drm_names(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(config, "decrypt_segments", True, raising=False)
+    monkeypatch.setattr(binaries, "Mp4decrypt", tmp_path / "mp4decrypt")
+    drm = _widevine(A, B)
+    drm.content_keys[A] = KEY
+
+    assert can_use(drm, "mp4decrypt", movie(encv(A))) is False
 
 
 class StubDRM:

tests/core/test_vault_bad_keys.py

@@ -384,6 +384,32 @@ def test_stale_backup_from_a_killed_run_is_not_restored(tmp_path: Path, monkeypa
     assert path.read_bytes() == b"cipher"
 
 
+@pytest.mark.parametrize(
+    ("windows", "licensed"),
+    [({KID: [(0.0, 2.0)], SIBLING: [(2.0, 4.0)]}, [SIBLING]), ({KID: [(0.0, 2.0)]}, []), (None, [SIBLING])],
+    ids=["file-uses-it", "file-does-not", "unmapped"],
+)
+def test_an_unkeyed_kid_the_file_uses_is_licensed_before_the_decrypt(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, windows: Optional[dict], licensed: list[UUID]
+) -> None:
+    cmd = make_cmd(monkeypatch, Vaults("SVC"))
+    drm = FakeDRM({KID: GOOD})
+    drm.kids = [KID, SIBLING]
+    monkeypatch.setattr(
+        dl_module.verify,
+        "kid_windows",
+        lambda path: None if windows is None else dl_module.verify.KidMap(windows, True),
+    )
+    path = tmp_path / "video.mp4"
+    path.write_bytes(b"cipher")
+    licences: list[UUID] = []
+
+    cmd.decrypt_verified(drm, path, lambda drm, track_kid: licences.append(track_kid), KID)
+
+    assert licences == licensed
+    assert drm.decrypted_with == [{KID: GOOD}]
+
+
 def test_drm_without_content_keys_still_decrypts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     """HLS AES-128 ClearKey has no KID:KEY map; the hook must pass it straight through."""
     cmd = make_cmd(monkeypatch, Vaults("SVC"))

tests/remote/unit/test_extract_pssh_strict.py

@@ -56,7 +56,15 @@ async def _run_batch(monkeypatch, tracks, cdm_type, calls):
         return session
 
     def fake_single(
-        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, vault_only=False, shared=None
+        service,
+        title,
+        track,
+        pssh_str,
+        drm_type,
+        request,
+        sources=None,
+        refusal=None,
+        **_,
     ):
         calls.append((pssh_str, drm_type))
         return {"00" * 16: "key"}

tests/remote/unit/test_remote_service_server_cdm_license.py

@@ -132,3 +132,12 @@ def test_a_refused_server_cdm_licence_logs_one_warning(caplog):
 
     assert svc.proxy_license(b"chal", track, "widevine") == b"chal"
     assert [r.levelname for r in caplog.records if r.levelname in ("WARNING", "ERROR")] == ["WARNING"]
+
+
+def test_the_kid_dl_needs_goes_to_the_server():
+    svc = _service("widevine", {"keys": {KID.hex: "k"}, "drm_type": "widevine"})
+    track = SimpleNamespace(id="vid", drm=[_Widevine()])
+
+    svc.get_widevine_license(challenge=b"", title=None, track=track, kid=KID)
+
+    assert svc.client.posts[0]["kid"] == KID.hex

tests/remote/unit/test_server_cdm_key_verification.py

@@ -115,12 +115,25 @@ async def test_batch_reports_vault_keys_from_the_init_segment_pssh(monkeypatch):
     session = SessionEntry(session_id="sess", service_tag="SVC", service_instance=SimpleNamespace())
     session.tracks = {"vid": track}
 
+    probed: list = []
+
     async def validated(sid, req):
         return session
 
     def fake_single(
-        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, vault_only=False, shared=None
+        service,
+        title,
+        track,
+        pssh_str,
+        drm_type,
+        request,
+        sources=None,
+        refusal=None,
+        vault_only=False,
+        shared=None,
+        track_kid=None,
     ):
+        probed.append(track_kid)
         if pssh_str == "init":
             sources[CDM_KID.hex] = "sqlite"
             return {CDM_KID.hex: "v"}
@@ -147,6 +160,66 @@ async def test_batch_reports_vault_keys_from_the_init_segment_pssh(monkeypatch):
     assert payload["keys"] == {"vid": {CDM_KID.hex: "v"}}
     assert payload["vault_keys"] == [CDM_KID.hex]
     assert session.served_keys == {CDM_KID.hex: ("v", "sqlite")}
+    assert probed == [CDM_KID, CDM_KID]
+
+
+async def test_batch_asks_again_when_a_cached_vault_answer_lacks_the_track_kid(monkeypatch):
+    """Renditions that share a header: a vault answer for one track's KID does not serve a track with another KID."""
+    kids = {"a": KID, "b": CDM_KID, "c": CDM_KID}
+    session = SessionEntry(session_id="sess", service_tag="SVC", service_instance=SimpleNamespace())
+    session.tracks = {
+        tid: SimpleNamespace(id=tid, drm=[_WV("shared", [KID, CDM_KID])], get_key_id=lambda init, k=kid: k)
+        for tid, kid in kids.items()
+    }
+    asked: list = []
+
+    async def validated(sid, req):
+        return session
+
+    def fake_single(
+        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, track_kid=None, **_
+    ):
+        asked.append(track_kid)
+        if track_kid == KID:
+            sources[KID.hex] = "sqlite"
+            return {KID.hex: "v"}
+        return {KID.hex: "v", CDM_KID.hex: "c"}
+
+    monkeypatch.setattr(handlers, "get_validated_session", validated)
+    monkeypatch.setattr(handlers, "server_cdm_allowed", lambda request, tag: True)
+    monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None)
+    monkeypatch.setattr(handlers, "find_title_for_track", lambda tid, session: SimpleNamespace())
+    monkeypatch.setattr(handlers, "detect_cdm_type_for_service", lambda tag, cfg: "widevine")
+    monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev")
+    monkeypatch.setattr(handlers, "fetch_init_segment", lambda track, sess: b"init")
+    monkeypatch.setattr(handlers, "handle_single_server_cdm", fake_single)
+    monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False)
+
+    resp = await handlers.session_license_handler(
+        {"mode": "server_cdm", "track_ids": list(kids), "drm_type": "widevine"}, "sess", None
+    )
+
+    payload = json.loads(resp.body)
+    assert asked == [KID, CDM_KID]
+    assert payload["keys"]["c"] == {KID.hex: "v", CDM_KID.hex: "c"}
+
+
+@pytest.mark.parametrize("drm", [["drm"], []], ids=["with-drm", "clear"])
+def test_the_server_reads_the_head_of_a_one_file_track_with_drm(drm):
+    from unshackle.core.tracks import Video
+
+    fetched: list = []
+    track = SimpleNamespace(
+        id="v",
+        data={},
+        descriptor=Video.Descriptor.URL,
+        drm=drm,
+        get_init_segment=lambda session: fetched.append(session) or b"head",
+    )
+
+    head = handlers.fetch_init_segment(track, "svc-session")
+
+    assert (head, fetched) == ((b"head", ["svc-session"]) if drm else (None, []))
 
 
 async def test_server_tells_a_remote_vault_only_when_it_served_the_pair(monkeypatch, session):

tests/remote/unit/test_server_cdm_manifest_kid.py

@@ -27,6 +27,8 @@ from unshackle.core.drm import PlayReady
 
 pytestmark = pytest.mark.unit
 
+REAL_CHECK_VAULTS = handlers.check_vaults
+
 KID_A = UUID("161219ec3df64e0eadbb3827c0ccc46d")
 KID_B = UUID("43de497c6d1945f49acc2951696ec628")
 
@@ -161,3 +163,59 @@ async def test_uncovered_kid_reaches_the_client(server_env, fake_playready, monk
     )
 
     assert any(KID_B.hex in r["message"] for r in buf.since(0)), "client gets no reason for the missing key"
+
+
+class _VaultsHolding:
+    """Server vaults that hold a content key for ``held`` only."""
+
+    def __init__(self, *held: UUID):
+        self.held = held
+        self.vaults = [SimpleNamespace(name="v1")]
+
+    def get_key(self, kid):
+        return ("k" + kid.hex[:4], self.vaults[0]) if kid in self.held else (None, None)
+
+
+@pytest.mark.parametrize(
+    ("track_kid", "found"), [(None, None), (KID_A, {KID_A.hex}), (KID_B, None)], ids=["every-kid", "held", "not-held"]
+)
+def test_check_vaults_answers_for_the_track_kid(monkeypatch, track_kid, found):
+    monkeypatch.setattr(handlers, "load_server_vaults", lambda name: _VaultsHolding(KID_A))
+
+    hit = handlers.check_vaults([KID_A, KID_B], "SVC", track_kid=track_kid)
+
+    assert (set(hit[0]) if hit else None) == found
+
+
+@pytest.mark.parametrize("track_kid", [KID_A, None], ids=["track-kid-held", "track-kid-unknown"])
+def test_a_playready_track_takes_the_vault_key_only_for_its_known_kid(server_env, monkeypatch, track_kid):
+    monkeypatch.setattr(handlers, "check_vaults", REAL_CHECK_VAULTS)
+    monkeypatch.setattr(handlers, "load_server_vaults", lambda name: _VaultsHolding(KID_A))
+    licensed: list = []
+    monkeypatch.setattr(
+        PlayReady,
+        "get_content_keys",
+        lambda self, cdm, certificate, licence: licensed.append(self) or self.content_keys.update({KID_A: "cdm"}),
+    )
+    pro_b64 = _kidless_pro_b64()
+    track = SimpleNamespace(
+        id="vid", drm=[PlayReady(pssh=RealPSSH(base64.b64decode(pro_b64)), kid=KID_A, pssh_b64=pro_b64)]
+    )
+    sources: dict = {}
+
+    keys = handlers.handle_single_server_cdm(
+        SimpleNamespace(), SimpleNamespace(), track, pro_b64, "playready", None, sources, track_kid=track_kid
+    )
+
+    assert bool(licensed) is (track_kid is None)
+    assert sources == ({KID_A.hex: "v1"} if track_kid else {})
+    assert set(keys) == {KID_A.hex}
+
+
+def test_the_client_kid_counts_only_when_the_track_drm_names_it():
+    track = SimpleNamespace(drm=[SimpleNamespace(kids=[KID_A])])
+
+    assert handlers.client_track_kid(KID_A.hex, track) == KID_A
+    assert handlers.client_track_kid(KID_B.hex, track) is None
+    assert handlers.client_track_kid("not-hex", track) is None
+    assert handlers.client_track_kid(None, track) is None

tests/remote/unit/test_server_cdm_max_height.py

@@ -397,7 +397,15 @@ def licence_env(monkeypatch):
     vault = {}
 
     def fake_single(
-        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, vault_only=False, shared=None
+        service,
+        title,
+        track,
+        pssh_str,
+        drm_type,
+        request,
+        sources=None,
+        refusal=None,
+        **_,
     ):
         calls.append((track.id, pssh_str))
         if pssh_str in vault:

tests/remote/unit/test_server_cdm_multikey.py

@@ -134,7 +134,15 @@ async def test_batch_shares_full_bundle_per_pssh(monkeypatch):
     bundle_b = {"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa": "k_b"}
 
     def fake_single(
-        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, vault_only=False, shared=None
+        service,
+        title,
+        track,
+        pssh_str,
+        drm_type,
+        request,
+        sources=None,
+        refusal=None,
+        **_,
     ):
         calls.append(pssh_str)
         return dict(bundle_a) if pssh_str == pssh_a else dict(bundle_b)