mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 12:05:07 +00:00
fix(drm): give the zero-KID fallback only the track's own key
mp4decrypt and shaka use the last key for a repeated KID.
This commit is contained in:
imSp4rky
@@ -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'<KID ALGID="AESCTR" VALUE="{base64.b64encode(k.bytes_le).decode()}"></KID>' for k in kids)
|
||||
header = (
|
||||
'<WRMHEADER xmlns="http://schemas.microsoft.com/DRM/2007/03/PlayReadyHeader" version="4.3.0.0">'
|
||||
f"<DATA><PROTECTINFO><KIDS>{kid_xml}</KIDS></PROTECTINFO></DATA></WRMHEADER>"
|
||||
).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()) == []
|
||||
@@ -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}")
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
]
|
||||
|
||||
@@ -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,
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user