mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-10-29 10:12:15 +01:00
imSp4rky
7065cbf911
Only the first WRM header was read and licensed, so a manifest carrying several headers silently lost content keys. Collect KIDs from every header and license each distinct one until coverage is complete.
417 lines
15 KiB
Python
417 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
import pytest
|
|
from pyplayready.system.pssh import PSSH
|
|
from pyplayready.system.wrmheader import WRMHeader
|
|
|
|
from unshackle.core.drm import PlayReady
|
|
|
|
KID_A = UUID("00112233-4455-6677-8899-aabbccddeeff")
|
|
KID_B = UUID("ffeeddcc-bbaa-9988-7766-554433221100")
|
|
KEY_A = "00000000000000000000000000000001"
|
|
KEY_B = "00000000000000000000000000000002"
|
|
|
|
|
|
def _wrm_header(*kids: UUID, la_url: str = "https://example.invalid/") -> bytes:
|
|
"""A minimal v4.3 WRMHEADER naming the given KIDs, UTF-16LE as the PRO carries it."""
|
|
kid_xml = "".join(f'<KID ALGID="AESCTR" VALUE="{base64.b64encode(kid.bytes_le).decode()}"></KID>' for kid in kids)
|
|
xml = (
|
|
'<WRMHEADER xmlns="http://schemas.microsoft.com/DRM/2007/03/PlayReadyHeader" version="4.3.0.0">'
|
|
f"<DATA><PROTECTINFO><KIDS>{kid_xml}</KIDS></PROTECTINFO><LA_URL>{la_url}</LA_URL></DATA>"
|
|
"</WRMHEADER>"
|
|
)
|
|
return xml.encode("utf-16-le")
|
|
|
|
|
|
def _pro(*headers: bytes) -> bytes:
|
|
records = [{"type": 1, "length": len(h), "data": h} for h in headers]
|
|
return PSSH.PlayreadyHeader.build({"length": 6 + sum(4 + len(h) for h in headers), "records": records})
|
|
|
|
|
|
def _drm(*headers: bytes) -> PlayReady:
|
|
pro = _pro(*headers)
|
|
return PlayReady(pssh=PSSH(pro), pssh_b64=base64.b64encode(pro).decode())
|
|
|
|
|
|
class FakeKey:
|
|
def __init__(self, kid: UUID, key: str) -> None:
|
|
self.key_id = kid
|
|
self.key = bytes.fromhex(key)
|
|
|
|
|
|
class FakeCdm:
|
|
"""Answers each challenge with the keys mapped for the header's KIDs; records what it saw."""
|
|
|
|
def __init__(self, answers: dict[UUID, str], fail_on: set[UUID] | None = None) -> None:
|
|
self.answers = answers
|
|
self.fail_on = fail_on or set()
|
|
self.challenged: list[list[UUID]] = []
|
|
self.pssh_b64s: list[str] = []
|
|
self.required: list[list[UUID]] = []
|
|
self.attempted: list[list[UUID]] = []
|
|
self._by_session: dict[str, list[UUID]] = {}
|
|
self._next = 0
|
|
|
|
def open(self) -> str:
|
|
self._next += 1
|
|
return f"session-{self._next}"
|
|
|
|
def close(self, session_id: str) -> None:
|
|
self._by_session.pop(session_id, None)
|
|
|
|
def set_pssh_b64(self, pssh_b64: str, session_id: Any = None) -> None:
|
|
self.pssh_b64s.append(pssh_b64)
|
|
|
|
def set_required_kids(self, kids: list[UUID], session_id: Any = None) -> None:
|
|
self.required.append(list(kids))
|
|
|
|
def get_license_challenge(self, session_id: str, header: WRMHeader) -> str:
|
|
kids = [k.value for k in header.key_ids]
|
|
self.attempted.append(kids)
|
|
if self.fail_on & set(kids):
|
|
raise RuntimeError("licence refused")
|
|
self.challenged.append(kids)
|
|
self._by_session[session_id] = kids
|
|
return "challenge"
|
|
|
|
def parse_license(self, session_id: str, response: str) -> None:
|
|
pass
|
|
|
|
def get_keys(self, session_id: str) -> list[FakeKey]:
|
|
return [FakeKey(kid, self.answers[kid]) for kid in self._by_session[session_id] if kid in self.answers]
|
|
|
|
|
|
def _licence(**_: Any) -> str:
|
|
return "<License>x</License>"
|
|
|
|
|
|
def _cert(**_: Any) -> None:
|
|
return None
|
|
|
|
|
|
def test_kids_are_collected_from_every_header() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
|
|
assert len(drm.pssh.wrm_headers) == 2
|
|
assert drm.kids == [KID_A, KID_B]
|
|
|
|
|
|
def test_single_header_sends_one_challenge_with_the_original_pssh() -> None:
|
|
drm = _drm(_wrm_header(KID_A, KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert cdm.challenged == [[KID_A, KID_B]]
|
|
assert cdm.pssh_b64s == [drm.pssh_b64]
|
|
assert cdm.required == [[KID_A, KID_B]]
|
|
assert drm.content_keys == {KID_A: KEY_A, KID_B: KEY_B}
|
|
|
|
|
|
def test_two_distinct_headers_are_both_licensed_and_keys_merged() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert cdm.challenged == [[KID_A], [KID_B]]
|
|
assert drm.content_keys == {KID_A: KEY_A, KID_B: KEY_B}
|
|
# a remote CDM sends whatever set_pssh_b64 gave it, so each round carries only its own header
|
|
assert [len(PSSH(b).wrm_headers) for b in cdm.pssh_b64s] == [1, 1]
|
|
for pro_b64 in cdm.pssh_b64s:
|
|
pro = base64.b64decode(pro_b64)
|
|
assert PSSH.PlayreadyHeader.parse(pro).length == len(pro)
|
|
assert _pro(PSSH(pro).wrm_headers[0].dumps().encode("utf-16-le")) == pro
|
|
assert [PSSH(b).wrm_headers[0].dumps() for b in cdm.pssh_b64s] == [h.dumps() for h in drm.pssh.wrm_headers]
|
|
|
|
|
|
def test_stacked_licence_covering_every_kid_stops_after_one_challenge() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B})
|
|
cdm.get_keys = lambda session_id: [FakeKey(KID_A, KEY_A), FakeKey(KID_B, KEY_B)] # type: ignore[method-assign]
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert cdm.challenged == [[KID_A]]
|
|
assert drm.content_keys == {KID_A: KEY_A, KID_B: KEY_B}
|
|
|
|
|
|
def test_headers_naming_the_same_kids_cost_one_challenge() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_A))
|
|
cdm = FakeCdm({KID_A: KEY_A})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert cdm.challenged == [[KID_A]]
|
|
|
|
|
|
def test_a_header_already_covered_by_the_vault_is_skipped() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
drm.content_keys[KID_A] = KEY_A
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert cdm.challenged == [[KID_B]]
|
|
|
|
|
|
def test_an_uncovered_kid_warns_instead_of_raising(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
events: list[dict[str, Any]] = []
|
|
monkeypatch.setattr(
|
|
"unshackle.core.drm.playready.log_event",
|
|
lambda operation, **kw: events.append({"operation": operation, **kw}),
|
|
)
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert drm.content_keys == {KID_A: KEY_A}
|
|
missing = [e for e in events if e["operation"] == "drm_missing_keys"]
|
|
assert missing and missing[0]["level"] == "WARNING"
|
|
assert missing[0]["missing_kids"] == [KID_B.hex]
|
|
|
|
|
|
def test_a_refused_header_does_not_lose_the_keys_of_the_others() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B}, fail_on={KID_A})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert drm.content_keys == {KID_B: KEY_B}
|
|
|
|
|
|
def test_every_header_refused_raises() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({}, fail_on={KID_A, KID_B})
|
|
|
|
with pytest.raises(RuntimeError, match="licence refused"):
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
|
|
def test_get_drm_for_cdm_folds_sibling_playready_objects_into_the_first() -> None:
|
|
from pyplayready.cdm import Cdm as PlayReadyCdm
|
|
|
|
from unshackle.core.tracks import Video
|
|
|
|
first = _drm(_wrm_header(KID_A))
|
|
second = _drm(_wrm_header(KID_B))
|
|
track = Video(
|
|
id_="v", url="https://example.invalid/v.mpd", language="en", codec=Video.Codec.AVC, drm=[first, second]
|
|
)
|
|
fake_cdm = PlayReadyCdm.__new__(PlayReadyCdm)
|
|
|
|
chosen = track.get_drm_for_cdm(fake_cdm)
|
|
|
|
assert chosen is first
|
|
assert chosen.kids == [KID_A, KID_B]
|
|
assert [[k.value for k in h.key_ids] for h in chosen.distinct_headers()] == [[KID_A], [KID_B]]
|
|
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B})
|
|
chosen.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
assert chosen.content_keys == {KID_A: KEY_A, KID_B: KEY_B}
|
|
|
|
|
|
def test_repeated_get_drm_for_cdm_does_not_grow_the_header_list() -> None:
|
|
from pyplayready.cdm import Cdm as PlayReadyCdm
|
|
|
|
from unshackle.core.tracks import Video
|
|
|
|
# a KID-less header is exempt from the KID-set dedupe, so only identity-level idempotence protects it
|
|
first = _drm(_wrm_header(KID_A))
|
|
second = PlayReady(pssh=PSSH(_pro(_wrm_header(la_url="https://kidless.invalid/"))), kid=KID_B)
|
|
track = Video(
|
|
id_="v", url="https://example.invalid/v.mpd", language="en", codec=Video.Codec.AVC, drm=[first, second]
|
|
)
|
|
fake_cdm = PlayReadyCdm.__new__(PlayReadyCdm)
|
|
|
|
for _ in range(3):
|
|
chosen = track.get_drm_for_cdm(fake_cdm)
|
|
|
|
assert chosen is first
|
|
assert len(chosen.distinct_headers()) == 2
|
|
assert chosen.kids == [KID_A, KID_B]
|
|
|
|
cdm = FakeCdm({KID_A: KEY_A})
|
|
chosen.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
assert len(cdm.challenged) == 2
|
|
|
|
|
|
def test_absorb_keeps_every_sibling_under_concurrent_calls() -> None:
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
first = _drm(_wrm_header(KID_A))
|
|
# each call brings a sibling nobody else brings, so a lost update shows as a missing header
|
|
siblings = [PlayReady(pssh=PSSH(_pro(_wrm_header(la_url=f"https://x{i}.invalid/"))), kid=KID_B) for i in range(64)]
|
|
|
|
with ThreadPoolExecutor(8) as pool:
|
|
list(pool.map(lambda sibling: first.absorb(sibling), siblings))
|
|
|
|
assert len(first.distinct_headers()) == 65
|
|
assert first.kids == [KID_A, KID_B]
|
|
|
|
|
|
def test_a_track_carrying_playready_survives_deepcopy() -> None:
|
|
from copy import deepcopy
|
|
|
|
from unshackle.core.tracks import Video
|
|
|
|
first = _drm(_wrm_header(KID_A))
|
|
second = _drm(_wrm_header(KID_B))
|
|
first.absorb(second)
|
|
track = Video(
|
|
id_="v", url="https://example.invalid/v.mpd", language="en", codec=Video.Codec.AVC, drm=[first, second]
|
|
)
|
|
|
|
clone = deepcopy(track)
|
|
|
|
assert clone.drm[0] is not first
|
|
assert clone.drm[0].kids == [KID_A, KID_B]
|
|
assert len(clone.drm[0].distinct_headers()) == 2
|
|
|
|
|
|
def _pro_b64(*headers: bytes) -> str:
|
|
return base64.b64encode(_pro(*headers)).decode()
|
|
|
|
|
|
def _xml_header(version: str, body: str) -> bytes:
|
|
return (
|
|
f'<WRMHEADER xmlns="http://schemas.microsoft.com/DRM/2007/03/PlayReadyHeader" version="{version}">'
|
|
f"<DATA>{body}</DATA></WRMHEADER>"
|
|
).encode("utf-16-le")
|
|
|
|
|
|
def _b64le(kid: UUID) -> str:
|
|
return base64.b64encode(kid.bytes_le).decode()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("version", "body", "expected"),
|
|
[
|
|
(
|
|
"4.0.0.0",
|
|
f"<PROTECTINFO><KEYLEN>16</KEYLEN><ALGID>AESCTR</ALGID></PROTECTINFO><KID>{_b64le(KID_A)}</KID>",
|
|
[KID_A],
|
|
),
|
|
("4.1.0.0", f'<PROTECTINFO><KID ALGID="AESCTR" VALUE="{_b64le(KID_B)}"></KID></PROTECTINFO>', [KID_B]),
|
|
(
|
|
"4.2.0.0",
|
|
f'<PROTECTINFO><KIDS><KID ALGID="AESCTR" VALUE="{_b64le(KID_A)}"></KID>'
|
|
f'<KID ALGID="AESCTR" VALUE="{_b64le(KID_B)}"></KID></KIDS></PROTECTINFO>',
|
|
[KID_A, KID_B],
|
|
),
|
|
],
|
|
)
|
|
def test_kids_are_read_from_every_header_version(version: str, body: str, expected: list[UUID]) -> None:
|
|
header = _xml_header(version, body)
|
|
assert PlayReady(pssh=PSSH(_pro(header)), pssh_b64=_pro_b64(header)).kids == expected
|
|
|
|
|
|
def test_customattributes_kids_follow_the_header_kids() -> None:
|
|
kid_c = UUID("11111111-2222-3333-4444-555555555555")
|
|
header = _xml_header(
|
|
"4.1.0.0",
|
|
f'<PROTECTINFO><KID ALGID="AESCTR" VALUE="{_b64le(KID_A)}"></KID></PROTECTINFO>'
|
|
f'<CUSTOMATTRIBUTES><KIDS><KID VALUE="{_b64le(KID_B)}"/><KID VALUE="{_b64le(kid_c)}"/></KIDS></CUSTOMATTRIBUTES>',
|
|
)
|
|
|
|
drm = PlayReady(pssh=PSSH(_pro(header)), pssh_b64=_pro_b64(header))
|
|
|
|
# the header KID leads (it is what the licence server acts on); CUSTOMATTRIBUTES KIDs stay reachable
|
|
assert drm.kids == [KID_A, KID_B, kid_c]
|
|
assert drm.kid == KID_A
|
|
|
|
|
|
def test_device_revoked_on_the_first_header_is_not_swallowed() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B})
|
|
|
|
calls: list[int] = []
|
|
|
|
def licence(**_: Any) -> str:
|
|
calls.append(1)
|
|
if len(calls) == 1:
|
|
raise PlayReady.Exceptions.DeviceRevoked("revoked")
|
|
return "<License>x</License>"
|
|
|
|
with pytest.raises(PlayReady.Exceptions.DeviceRevoked):
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=licence)
|
|
assert cdm.challenged == [[KID_A]]
|
|
|
|
|
|
def test_a_revocation_hresult_in_a_refusal_becomes_device_revoked() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B})
|
|
|
|
calls: list[int] = []
|
|
|
|
def licence(**_: Any) -> str:
|
|
calls.append(1)
|
|
if len(calls) == 1:
|
|
raise RuntimeError("<StatusCode>0x8004C065</StatusCode>")
|
|
return "<License>x</License>"
|
|
|
|
with pytest.raises(PlayReady.Exceptions.DeviceRevoked):
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=licence)
|
|
assert cdm.challenged == [[KID_A]]
|
|
|
|
|
|
def _kidless_header() -> bytes:
|
|
return _wrm_header(la_url="https://kidless.invalid/")
|
|
|
|
|
|
def test_a_kidless_header_is_still_sent_when_the_vault_covered_every_kid() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _kidless_header())
|
|
drm.content_keys[KID_A] = KEY_A
|
|
cdm = FakeCdm({KID_A: KEY_A})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert cdm.challenged == [[]]
|
|
|
|
|
|
def test_a_kidless_header_is_still_sent_when_a_licence_round_covered_every_kid() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _kidless_header())
|
|
cdm = FakeCdm({KID_A: KEY_A})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert cdm.challenged == [[KID_A], []]
|
|
|
|
|
|
def test_a_header_that_yielded_nothing_is_not_sent_again_by_the_same_object() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A})
|
|
|
|
# every track of a title calls this on the shared object; the second call must not repeat B
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert cdm.challenged == [[KID_A], [KID_B]]
|
|
assert drm.content_keys == {KID_A: KEY_A}
|
|
|
|
|
|
def test_a_refused_header_is_not_sent_again_by_the_same_object() -> None:
|
|
drm = _drm(_wrm_header(KID_A), _wrm_header(KID_B))
|
|
cdm = FakeCdm({KID_A: KEY_A, KID_B: KEY_B}, fail_on={KID_B})
|
|
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
attempts_after_first = len(cdm.attempted)
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|
|
|
|
assert len(cdm.attempted) == attempts_after_first
|
|
assert drm.content_keys == {KID_A: KEY_A}
|
|
|
|
|
|
def test_a_single_refused_header_still_raises_on_every_call() -> None:
|
|
drm = _drm(_wrm_header(KID_A))
|
|
cdm = FakeCdm({}, fail_on={KID_A})
|
|
|
|
for _ in range(2):
|
|
with pytest.raises(RuntimeError, match="licence refused"):
|
|
drm.get_content_keys(cdm=cdm, certificate=_cert, licence=_licence)
|