mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-13 12:35:05 +00:00
fix(hls): read the key id from the init segment when the PSSH has none
An EXT-X-KEY whose Widevine PSSH holds a content id and no key id raised KIDNotFound. Covers download, licence-only, --all-drm, session keys, import and the server CDM path. Fixes #166.
This commit is contained in:
imSp4rky
@@ -0,0 +1,132 @@
|
||||
"""An HLS Widevine key whose PSSH holds a content id and no Key ID takes the Key ID of its init segment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
import m3u8
|
||||
from pywidevine.license_protocol_pb2 import WidevinePsshData
|
||||
from pywidevine.pssh import PSSH
|
||||
|
||||
from unshackle.core.constants import DOWNLOAD_LICENCE_ONLY
|
||||
from unshackle.core.drm import PlayReady, Widevine
|
||||
from unshackle.core.import_service import hls_playlist_kids
|
||||
from unshackle.core.manifests.hls import HLS
|
||||
from unshackle.core.tracks import Video
|
||||
from unshackle.core.tracks.track import DownloadContext, Track
|
||||
|
||||
INIT_KID = UUID("11111111-1111-1111-1111-111111111111")
|
||||
PSSH_KID = UUID("22222222-2222-2222-2222-222222222222")
|
||||
TAG_KID = UUID("33333333-3333-3333-3333-333333333333")
|
||||
|
||||
ZERO_KID = UUID(int=0)
|
||||
|
||||
|
||||
def no_kid_pssh(content_id: bytes = b"content-id-only") -> str:
|
||||
return PSSH.new(
|
||||
system_id=PSSH.SystemId.Widevine,
|
||||
init_data=WidevinePsshData(content_id=content_id).SerializeToString(),
|
||||
).dumps()
|
||||
|
||||
|
||||
NO_KID_PSSH = no_kid_pssh()
|
||||
KEY_LINE = (
|
||||
'#EXT-X-KEY:METHOD=SAMPLE-AES-CTR,URI="data:text/plain;base64,{pssh}",'
|
||||
'KEYFORMAT="urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed",KEYFORMATVERSIONS="1"'
|
||||
)
|
||||
|
||||
|
||||
def keyed_segment(pssh: str, extra: str = "", init: str = "init.mp4") -> str:
|
||||
return f'#EXT-X-MAP:URI="{init}"\n{KEY_LINE.format(pssh=pssh)}{extra}\n#EXTINF:6.0,\n{init}.m4s\n'
|
||||
|
||||
|
||||
def playlist(*segments: str) -> str:
|
||||
body = "#EXT-X-DISCONTINUITY\n".join(segments)
|
||||
return f"#EXTM3U\n#EXT-X-VERSION:7\n#EXT-X-TARGETDURATION:6\n{body}#EXT-X-ENDLIST\n"
|
||||
|
||||
|
||||
def key_of(pssh: str, extra: str = "") -> m3u8.Key:
|
||||
return m3u8.loads(playlist(keyed_segment(pssh, extra))).keys[0]
|
||||
|
||||
|
||||
def licence_only(text: str, tmp_path: Path) -> list[tuple[list[UUID], Any]]:
|
||||
"""Run a licence-only download of the playlist and return the Key IDs of each licence call."""
|
||||
media = tmp_path / "media.m3u8"
|
||||
media.write_text(text, encoding="utf-8")
|
||||
track = Video(
|
||||
url="https://example.test/media.m3u8",
|
||||
language="en",
|
||||
descriptor=Track.Descriptor.HLS,
|
||||
codec=Video.Codec.AVC,
|
||||
)
|
||||
track.from_file = media
|
||||
licensed: list[tuple[list[UUID], Any]] = []
|
||||
ctx = DownloadContext(
|
||||
save_path=tmp_path / "video.mp4",
|
||||
save_dir=tmp_path / "video_segments",
|
||||
progress=lambda **_: None,
|
||||
license_widevine=lambda drm, track_kid=None: licensed.append((drm.kids, track_kid)),
|
||||
cdm=None,
|
||||
)
|
||||
DOWNLOAD_LICENCE_ONLY.set()
|
||||
try:
|
||||
HLS.download_track(track=track, ctx=ctx)
|
||||
finally:
|
||||
DOWNLOAD_LICENCE_ONLY.clear()
|
||||
return licensed
|
||||
|
||||
|
||||
def test_download_track_licenses_with_the_init_segment_kid(tmp_path: Path, monkeypatch: Any) -> None:
|
||||
# the probe itself needs the network and FFprobe; the order of the calls is under test
|
||||
monkeypatch.setattr(HLS, "get_track_kid_from_init", staticmethod(lambda *args, **kwargs: INIT_KID))
|
||||
assert licence_only(playlist(keyed_segment(NO_KID_PSSH)), tmp_path) == [([INIT_KID], INIT_KID)]
|
||||
|
||||
|
||||
def test_licence_only_reads_each_later_key_from_its_own_init_segment(tmp_path: Path, monkeypatch: Any) -> None:
|
||||
def probe(master: Any, track: Any, session: Any, init_section: Any = None) -> UUID:
|
||||
return PSSH_KID if init_section is not None and init_section.uri == "init2.mp4" else INIT_KID
|
||||
|
||||
monkeypatch.setattr(HLS, "get_track_kid_from_init", staticmethod(probe))
|
||||
text = playlist(keyed_segment(no_kid_pssh(b"first")), keyed_segment(no_kid_pssh(b"later"), init="init2.mp4"))
|
||||
assert licence_only(text, tmp_path) == [([INIT_KID], INIT_KID), ([PSSH_KID], PSSH_KID)]
|
||||
|
||||
|
||||
def test_get_drm_keeps_a_pssh_that_has_its_own_kid() -> None:
|
||||
pssh = PSSH.new(key_ids=[PSSH_KID], system_id=PSSH.SystemId.Widevine).dumps()
|
||||
drm = HLS.get_drm(key_of(pssh), kid=INIT_KID)
|
||||
assert drm.kids == [PSSH_KID]
|
||||
assert drm.pssh.dumps() == pssh
|
||||
|
||||
|
||||
def test_get_drm_prefers_the_tag_keyid_over_the_init_segment_kid() -> None:
|
||||
drm = HLS.get_drm(key_of(NO_KID_PSSH, f",KEYID=0x{TAG_KID.hex}"), kid=INIT_KID)
|
||||
assert drm.kids == [TAG_KID]
|
||||
|
||||
|
||||
def test_get_drm_replaces_a_placeholder_pssh_kid_and_ignores_a_placeholder_init_kid() -> None:
|
||||
placeholder = PSSH.new(key_ids=[ZERO_KID], system_id=PSSH.SystemId.Widevine).dumps()
|
||||
assert HLS.get_drm(key_of(placeholder), kid=INIT_KID).kids == [INIT_KID]
|
||||
assert HLS.get_drm(key_of(placeholder), kid=ZERO_KID).kids == [ZERO_KID]
|
||||
|
||||
|
||||
def test_session_key_without_a_kid_does_not_stop_the_parse() -> None:
|
||||
with_kid = PSSH.new(key_ids=[PSSH_KID], system_id=PSSH.SystemId.Widevine).dumps()
|
||||
drm = HLS.get_all_drm([key_of(NO_KID_PSSH), key_of(with_kid)])
|
||||
assert [d.kids for d in drm] == [[PSSH_KID]]
|
||||
|
||||
|
||||
def test_alternate_drm_builds_the_other_system_with_the_init_segment_kid() -> None:
|
||||
chosen = object.__new__(PlayReady)
|
||||
alternate = HLS.alternate_drm([key_of(NO_KID_PSSH)], chosen, kid=INIT_KID)
|
||||
assert [type(d) for d in alternate] == [Widevine] and alternate[0].kids == [INIT_KID]
|
||||
|
||||
|
||||
def test_import_reads_the_kid_of_a_playlist_whose_key_has_none() -> None:
|
||||
def get(url: str, headers: Any = None) -> Any:
|
||||
return SimpleNamespace(text=playlist(keyed_segment(NO_KID_PSSH)), ok=True, content=b"init")
|
||||
|
||||
track = SimpleNamespace(id="vid", url="https://example.test/media.m3u8", get_key_id=lambda init: INIT_KID)
|
||||
assert hls_playlist_kids(track, SimpleNamespace(get=get)) == {INIT_KID}
|
||||
@@ -21,6 +21,7 @@ from Cryptodome.Cipher import AES
|
||||
from Cryptodome.Util.Padding import pad
|
||||
from pywidevine.pssh import PSSH
|
||||
|
||||
from tests.core.test_hls_kid_from_init import KEY_LINE, no_kid_pssh
|
||||
from tests.core.test_ism_init import piff_fragment
|
||||
from unshackle.core.config import config
|
||||
from unshackle.core.constants import DOWNLOAD_CANCELLED
|
||||
@@ -264,3 +265,34 @@ def test_hls_aes128_keeps_post_merge(server, tmp_path, monkeypatch):
|
||||
|
||||
plain, rolling = download_both_ways(HLS, make, tmp_path, monkeypatch, False)
|
||||
assert rolling == plain == INIT + b"".join(PARTS)
|
||||
|
||||
|
||||
def test_hls_rekey_takes_the_kid_from_the_init_segment_of_the_new_key(server, tmp_path, monkeypatch):
|
||||
"""Each key carries a PSSH with a content id and no Key ID, so only the init segment names the KID."""
|
||||
later_kid = UUID(hex="22222222222222222222222222222222")
|
||||
later_init = bytes(range(100, 200))
|
||||
server.routes["/init2.mp4"] = later_init
|
||||
no_decrypt(monkeypatch, hls_module)
|
||||
monkeypatch.setattr(config, "merge_segments", False)
|
||||
|
||||
def key_line(content_id: bytes) -> str:
|
||||
return KEY_LINE.format(pssh=no_kid_pssh(content_id))
|
||||
|
||||
track = hls_track(server, False)
|
||||
lines = ["#EXTM3U", "#EXT-X-VERSION:7", "#EXT-X-TARGETDURATION:4", '#EXT-X-MAP:URI="init.mp4"', key_line(b"a")]
|
||||
for i in range(len(PARTS)):
|
||||
if i == 2:
|
||||
lines += ["#EXT-X-DISCONTINUITY", '#EXT-X-MAP:URI="init2.mp4"', key_line(b"b")]
|
||||
lines += ["#EXTINF:4.0,", f"part{i}.m4s"]
|
||||
server.routes["/media.m3u8"] = "\n".join(lines + ["#EXT-X-ENDLIST"]).encode()
|
||||
# the real probe needs FFprobe and a tenc box; the KID follows the init bytes the parser hands over
|
||||
monkeypatch.setattr(
|
||||
type(track), "get_key_id", lambda self, init_data=None, **kw: {later_init: later_kid}.get(init_data, KID)
|
||||
)
|
||||
|
||||
licensed: list[UUID] = []
|
||||
ctx = make_ctx(tmp_path, "rekey.mp4")
|
||||
ctx.license_widevine = lambda drm, track_kid=None: licensed.append(track_kid)
|
||||
HLS.download_track(track, ctx)
|
||||
|
||||
assert licensed == [KID, later_kid]
|
||||
|
||||
@@ -13,6 +13,7 @@ from uuid import UUID
|
||||
import pytest
|
||||
from pywidevine.pssh import PSSH as WvPSSH
|
||||
|
||||
from tests.core.test_hls_kid_from_init import NO_KID_PSSH
|
||||
from unshackle.core.api import handlers
|
||||
from unshackle.core.drm import Widevine
|
||||
from unshackle.core.tracks import Track
|
||||
@@ -22,13 +23,14 @@ pytestmark = pytest.mark.unit
|
||||
KID = UUID("161219ec3df64e0eadbb3827c0ccc46d")
|
||||
|
||||
|
||||
def _media_playlist() -> str:
|
||||
pssh = WvPSSH.new(key_ids=[KID], system_id=WvPSSH.SystemId.Widevine).dumps()
|
||||
def _media_playlist(pssh: str | None = None) -> str:
|
||||
pssh = pssh or WvPSSH.new(key_ids=[KID], system_id=WvPSSH.SystemId.Widevine).dumps()
|
||||
return "\n".join(
|
||||
[
|
||||
"#EXTM3U",
|
||||
"#EXT-X-VERSION:6",
|
||||
"#EXT-X-TARGETDURATION:6",
|
||||
'#EXT-X-MAP:URI="init.mp4"',
|
||||
f'#EXT-X-KEY:METHOD=SAMPLE-AES,URI="data:text/plain;base64,{pssh}",'
|
||||
'KEYFORMAT="urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed",KEYFORMATVERSIONS="1"',
|
||||
"#EXTINF:6.0,",
|
||||
@@ -49,3 +51,21 @@ def test_hls_descriptor_alone_reaches_the_media_playlist_keys():
|
||||
assert track.drm and isinstance(track.drm[0], Widevine), "media playlist DRM not read"
|
||||
assert track.drm[0].kids == [KID]
|
||||
assert base64.b64decode(handlers.extract_pssh_from_track(track, "widevine"))
|
||||
|
||||
|
||||
def test_key_without_a_kid_takes_it_from_the_init_segment():
|
||||
"""A PSSH that holds a content id and no Key ID is DRM, and the server must not call the track clear."""
|
||||
response = SimpleNamespace(text=_media_playlist(NO_KID_PSSH), ok=True, content=b"init")
|
||||
session = SimpleNamespace(get=lambda url, headers=None: response)
|
||||
track = SimpleNamespace(
|
||||
id="vid",
|
||||
url="https://cdn.example.invalid/media.m3u8",
|
||||
descriptor=Track.Descriptor.HLS,
|
||||
data={},
|
||||
drm=None,
|
||||
get_key_id=lambda init: KID,
|
||||
)
|
||||
|
||||
handlers.ensure_track_drm(track, session)
|
||||
|
||||
assert track.drm and track.drm[0].kids == [KID]
|
||||
|
||||
Reference in New Issue
Block a user