diff --git a/tests/core/test_zero_kid_fallback.py b/tests/core/test_zero_kid_fallback.py new file mode 100644 index 00000000..dd41f3b3 --- /dev/null +++ b/tests/core/test_zero_kid_fallback.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import base64 +from typing import Union +from uuid import UUID + +import pytest +from pyplayready.system.pssh import PSSH as PlayReadyPSSH +from pywidevine.pssh import PSSH as WidevinePSSH + +from unshackle.core.drm import PlayReady, Widevine, drm_from_dict, key_args +from unshackle.core.drm.key_args import ZERO_KID + +VIDEO_KID = UUID("00112233-4455-6677-8899-aabbccddeeff") +AUDIO_KID = UUID("ffeeddcc-bbaa-9988-7766-554433221100") +VIDEO_KEY = "0000000000000000000000000000000a" +AUDIO_KEY = "0000000000000000000000000000000b" + +DRM = Union[Widevine, PlayReady] + + +def _playready_pssh(*kids: UUID) -> tuple[PlayReadyPSSH, str]: + kid_xml = "".join(f'' for k in kids) + header = ( + '' + f"{kid_xml}" + ).encode("utf-16-le") + pro = PlayReadyPSSH.PlayreadyHeader.build( + {"length": 6 + 4 + len(header), "records": [{"type": 1, "length": len(header), "data": header}]} + ) + return PlayReadyPSSH(pro), base64.b64encode(pro).decode() + + +def _widevine(*kids: UUID, kid: UUID | None = None) -> Widevine: + return Widevine(pssh=WidevinePSSH.new(system_id=WidevinePSSH.SystemId.Widevine, key_ids=list(kids)), kid=kid) + + +def _playready(*kids: UUID, kid: UUID | None = None) -> PlayReady: + pssh, pssh_b64 = _playready_pssh(*kids) + return PlayReady(pssh=pssh, kid=kid, pssh_b64=pssh_b64) + + +def _with_keys(drm: DRM, keys: dict[UUID, str]) -> DRM: + drm.content_keys.update(keys) + return drm + + +BOTH_KEYS = {VIDEO_KID: VIDEO_KEY, AUDIO_KID: AUDIO_KEY} + + +def _zero_entries(args: list[str]) -> list[str]: + return [a.split(":", 1)[1] for a in args if a.startswith(f"{ZERO_KID}:")] + + +def _shaka_zero_entries(keys: str) -> list[str]: + return [e.rsplit("key=", 1)[1] for e in keys.split(",") if f"key_id={ZERO_KID}" in e] + + +def _shaka_keys(drm: DRM) -> str: + return key_args.shaka_keys(drm.content_keys, drm.own_kids()) + + +OWN_KID_KNOWN = [ + pytest.param(lambda: _widevine(VIDEO_KID), id="widevine-pssh-kid"), + pytest.param(lambda: _widevine(VIDEO_KID, AUDIO_KID, kid=VIDEO_KID), id="widevine-constructor-kid"), + pytest.param(lambda: _playready(VIDEO_KID), id="playready-pssh-kid"), + pytest.param(lambda: _playready(VIDEO_KID, AUDIO_KID, kid=VIDEO_KID), id="playready-constructor-kid"), +] + + +@pytest.mark.parametrize("make", OWN_KID_KNOWN) +def test_several_keys_give_the_zero_kid_only_the_own_key(make) -> None: + drm = _with_keys(make(), BOTH_KEYS) + + args = drm.mp4decrypt_key_args() + + assert f"{VIDEO_KID.hex}:{VIDEO_KEY}" in args and f"{AUDIO_KID.hex}:{AUDIO_KEY}" in args + assert _zero_entries(args) == [VIDEO_KEY] + assert _shaka_zero_entries(_shaka_keys(drm)) == [VIDEO_KEY] + + +@pytest.mark.parametrize("make", [_widevine, _playready]) +def test_a_single_key_keeps_its_zero_kid_entry(make) -> None: + drm = _with_keys(make(VIDEO_KID), {VIDEO_KID: VIDEO_KEY}) + + assert drm.mp4decrypt_key_args() == ["--key", f"{VIDEO_KID.hex}:{VIDEO_KEY}", "--key", f"{ZERO_KID}:{VIDEO_KEY}"] + assert _shaka_keys(drm) == ( + f"label=0:key_id={VIDEO_KID.hex}:key={VIDEO_KEY},label=1:key_id={ZERO_KID}:key={VIDEO_KEY}" + ) + + +@pytest.mark.parametrize("make", [_widevine, _playready]) +def test_several_keys_and_no_single_own_kid_give_no_zero_kid_entry(make) -> None: + drm = _with_keys(make(VIDEO_KID, AUDIO_KID), BOTH_KEYS) + + assert _zero_entries(drm.mp4decrypt_key_args()) == [] + assert _shaka_zero_entries(_shaka_keys(drm)) == [] + + +def test_shaka_decrypt_passes_the_built_keys(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + """The shaka-packager call uses the same builder the tests above cover.""" + from unshackle.core.drm import widevine as widevine_module + + seen: list[list[str]] = [] + + class StopPopen(Exception): + pass + + def fake_popen(cmd, **kwargs): + seen.append([str(c) for c in cmd]) + raise StopPopen + + monkeypatch.setattr(widevine_module.subprocess, "Popen", fake_popen) + monkeypatch.setattr(widevine_module.binaries, "ShakaPackager", "shaka-packager") + monkeypatch.setattr(widevine_module.config.directories, "temp", tmp_path) + drm = _with_keys(_widevine(VIDEO_KID), BOTH_KEYS) + media = tmp_path / "track.mp4" + media.write_bytes(b"x") + + with pytest.raises(StopPopen): + drm.decrypt_with_shaka_packager(media) + + keys = seen[0][seen[0].index("--keys") + 1] + assert _shaka_zero_entries(keys) == [VIDEO_KEY] + + +def test_a_rebuilt_widevine_does_not_take_the_first_pssh_kid_as_its_own() -> None: + drm = drm_from_dict(_widevine(VIDEO_KID, AUDIO_KID).to_dict()) + drm.content_keys.update(BOTH_KEYS) + + assert _zero_entries(drm.mp4decrypt_key_args()) == [] diff --git a/unshackle/core/drm/__init__.py b/unshackle/core/drm/__init__.py index d9893e2f..bae6603a 100644 --- a/unshackle/core/drm/__init__.py +++ b/unshackle/core/drm/__init__.py @@ -34,7 +34,11 @@ def drm_from_dict(data: dict[str, Any]) -> Union[Widevine, PlayReady, ClearKeyCE elif system == "Widevine": from pywidevine.pssh import PSSH as WidevinePSSH - drm = Widevine(pssh=WidevinePSSH(pssh_b64), kid=kids[0] if kids else None) + wv_pssh = WidevinePSSH(pssh_b64) + # kids repeats the PSSH KIDs, so kids[0] is not the track's own KID; only a PSSH + # without KIDs needs it + + drm = Widevine(pssh=wv_pssh, kid=kids[0] if kids and not wv_pssh.key_ids else None) else: raise ValueError(f"Unsupported DRM system for reconstruction: {system!r}") diff --git a/unshackle/core/drm/key_args.py b/unshackle/core/drm/key_args.py new file mode 100644 index 00000000..7b310e15 --- /dev/null +++ b/unshackle/core/drm/key_args.py @@ -0,0 +1,62 @@ +"""Decrypter key arguments shared by the Widevine and PlayReady DRM systems.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +ZERO_KID = "00" * 16 + + +def _hex(value: Any) -> str: + """Return a KID or content key as lowercase hex without dashes.""" + if isinstance(value, bytes): + return value.hex() + if hasattr(value, "hex") and not isinstance(value, str): + return str(value.hex) + return str(value).replace("-", "").lower() + + +def zero_kid_key(content_keys: dict[Any, Any], own_kids: Iterable[Any]) -> Optional[str]: + """Return the content key to give the all-zero KID, or None when no content key is safe. + + Some tracks carry an all-zero default KID in the init segment and signal the real KID + out of band. The decrypter then needs the track's content key under the zero KID as well. + + A decrypter applies the last content key given for a repeated KID. Thus only one content + key can go under the zero KID, and it must be the content key of the track: + + - The DRM holds exactly one content key: that content key. + - The DRM holds several content keys: the content key of the one own KID that has one. + - Otherwise: None. A missing content key makes the decrypt fail loudly, but an incorrect + content key gives output that does not decode, and the verify step then flags good + content keys as bad. + """ + keys = {_hex(kid): _hex(key) for kid, key in content_keys.items()} + if ZERO_KID in keys or not keys: + return None + if len(keys) == 1: + return next(iter(keys.values())) + matches = {_hex(kid) for kid in own_kids} & keys.keys() + if len(matches) == 1: + return keys[matches.pop()] + return None + + +def mp4decrypt_key_args(content_keys: dict[Any, Any], own_kids: Iterable[Any]) -> list[str]: + """Build the mp4decrypt --key arguments: every content key, then the zero-KID fallback.""" + key_args: list[str] = [] + for kid, key in content_keys.items(): + key_args.extend(["--key", f"{_hex(kid)}:{_hex(key)}"]) + fallback = zero_kid_key(content_keys, own_kids) + if fallback: + key_args.extend(["--key", f"{ZERO_KID}:{fallback}"]) + return key_args + + +def shaka_keys(content_keys: dict[Any, Any], own_kids: Iterable[Any]) -> str: + """Build the shaka-packager --keys value: every content key, then the zero-KID fallback.""" + entries = [f"label={i}:key_id={_hex(kid)}:key={_hex(key)}" for i, (kid, key) in enumerate(content_keys.items())] + fallback = zero_kid_key(content_keys, own_kids) + if fallback: + entries.append(f"label={len(entries)}:key_id={ZERO_KID}:key={fallback}") + return ",".join(entries) diff --git a/unshackle/core/drm/playready.py b/unshackle/core/drm/playready.py index cddb04e7..f75d9d61 100644 --- a/unshackle/core/drm/playready.py +++ b/unshackle/core/drm/playready.py @@ -24,6 +24,7 @@ from unshackle.core import binaries from unshackle.core.config import config from unshackle.core.console import console from unshackle.core.constants import AnyTrack +from unshackle.core.drm import key_args as drm_key_args from unshackle.core.utilities import get_boxes, log_event from unshackle.core.utils.subprocess import ffprobe @@ -55,6 +56,7 @@ class PlayReady: if extra_kid not in kids: kids.append(extra_kid) + self._kid: Optional[UUID] = None if kid: if isinstance(kid, str): kid = UUID(hex=kid) @@ -62,6 +64,7 @@ class PlayReady: kid = UUID(bytes=kid) if not isinstance(kid, UUID): raise ValueError(f"Expected kid to be a {UUID}, str, or bytes, not {kid!r}") + self._kid = kid if kid not in kids: kids.append(kid) @@ -563,23 +566,19 @@ class PlayReady: output_size=path.stat().st_size if path.exists() else 0, ) - def mp4decrypt_key_args(self) -> list[str]: - """Build the mp4decrypt --key arguments for every content key.""" - key_args = [] - for kid, key in self.content_keys.items(): - kid_hex = kid.hex if hasattr(kid, "hex") else str(kid).replace("-", "") - key_hex = key if isinstance(key, str) else key.hex() - key_args.extend(["--key", f"{kid_hex}:{key_hex}"]) + def own_kids(self) -> list[UUID]: + """Return the KIDs that identify this track's own content key. - # Fallback for tracks whose tenc default_KID is all-zero and whose real - # KID is signalled out-of-band: emit a zero-KID entry per content key. - zero_kid = "00" * 16 - existing_kids = {kid.hex if hasattr(kid, "hex") else str(kid).replace("-", "") for kid in self.content_keys} - if zero_kid not in existing_kids: - for key in self.content_keys.values(): - key_hex = key if isinstance(key, str) else key.hex() - key_args.extend(["--key", f"{zero_kid}:{key_hex}"]) - return key_args + The KID given to the constructor (from the manifest or the init segment) names the + track exactly. A PSSH can name the KIDs of other renditions too. + """ + if self._kid and self._kid in self.content_keys: + return [self._kid] + return self.kids + + def mp4decrypt_key_args(self) -> list[str]: + """Build the mp4decrypt --key arguments for every content key and the zero-KID fallback.""" + return drm_key_args.mp4decrypt_key_args(self.content_keys, self.own_kids()) def decrypt_with_mp4decrypt(self, path: Path) -> None: """Decrypt using mp4decrypt""" @@ -633,18 +632,7 @@ class PlayReady: f"input={path},stream=0,output={output_path},output_format=MP4", "--enable_raw_key_decryption", "--keys", - ",".join( - [ - *[ - f"label={i}:key_id={kid.hex}:key={key.lower()}" - for i, (kid, key) in enumerate(self.content_keys.items()) - ], - *[ - f"label={i}:key_id={'00' * 16}:key={key.lower()}" - for i, (kid, key) in enumerate(self.content_keys.items(), len(self.content_keys)) - ], - ] - ), + drm_key_args.shaka_keys(self.content_keys, self.own_kids()), "--temp_dir", config.directories.temp, ] diff --git a/unshackle/core/drm/widevine.py b/unshackle/core/drm/widevine.py index a6cbf407..377e2b9b 100644 --- a/unshackle/core/drm/widevine.py +++ b/unshackle/core/drm/widevine.py @@ -21,6 +21,7 @@ from unshackle.core import binaries from unshackle.core.config import config from unshackle.core.console import console from unshackle.core.constants import AnyTrack +from unshackle.core.drm import key_args as drm_key_args from unshackle.core.utilities import get_boxes, log_event from unshackle.core.utils.subprocess import ffprobe @@ -315,23 +316,19 @@ class Widevine: output_size=path.stat().st_size if path.exists() else 0, ) - def mp4decrypt_key_args(self) -> list[str]: - """Build the mp4decrypt --key arguments for every content key.""" - key_args = [] - for kid, key in self.content_keys.items(): - kid_hex = kid.hex if hasattr(kid, "hex") else str(kid).replace("-", "") - key_hex = key if isinstance(key, str) else key.hex() - key_args.extend(["--key", f"{kid_hex}:{key_hex}"]) + def own_kids(self) -> list[UUID]: + """Return the KIDs that identify this track's own content key. - # Fallback for tracks whose tenc default_KID is all-zero and whose real - # KID is signalled out-of-band: emit a zero-KID entry per content key. - zero_kid = "00" * 16 - existing_kids = {kid.hex if hasattr(kid, "hex") else str(kid).replace("-", "") for kid in self.content_keys} - if zero_kid not in existing_kids: - for key in self.content_keys.values(): - key_hex = key if isinstance(key, str) else key.hex() - key_args.extend(["--key", f"{zero_kid}:{key_hex}"]) - return key_args + The KID given to the constructor (from the manifest or the init segment) names the + track exactly. A PSSH can name the KIDs of other renditions too. + """ + if self._kid and self._kid in self.content_keys: + return [self._kid] + return self.kids + + def mp4decrypt_key_args(self) -> list[str]: + """Build the mp4decrypt --key arguments for every content key and the zero-KID fallback.""" + return drm_key_args.mp4decrypt_key_args(self.content_keys, self.own_kids()) def decrypt_with_mp4decrypt(self, path: Path) -> None: """Decrypt using mp4decrypt""" @@ -385,19 +382,7 @@ class Widevine: f"input={path},stream=0,output={output_path},output_format=MP4", "--enable_raw_key_decryption", "--keys", - ",".join( - [ - *[ - "label={}:key_id={}:key={}".format(i, kid.hex, key.lower()) - for i, (kid, key) in enumerate(self.content_keys.items()) - ], - *[ - # some services use a blank KID on the file, but real KID for license server - "label={}:key_id={}:key={}".format(i, "00" * 16, key.lower()) - for i, (kid, key) in enumerate(self.content_keys.items(), len(self.content_keys)) - ], - ] - ), + drm_key_args.shaka_keys(self.content_keys, self.own_kids()), "--temp_dir", config.directories.temp, ]