mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 12:05:07 +00:00
116 lines
4.0 KiB
Python
116 lines
4.0 KiB
Python
"""Key selection helpers of the HLS parser: session-data keys, CDM filtering and key choice."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
from typing import Optional
|
|
|
|
import m3u8
|
|
import pytest
|
|
|
|
from unshackle.core.cdm.detect import cdm_type_stub
|
|
from unshackle.core.manifests.hls import HLS
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
WV_URN = "urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed"
|
|
PR_URN = "urn:uuid:9a04f079-9840-4286-ab92-e65be0885f95"
|
|
PR_NAME = "com.microsoft.playready"
|
|
KEY_URI = "data:text/plain;base64,AAAA"
|
|
|
|
|
|
def _master(value: str, data_id: str = "com.apple.hls.keys") -> m3u8.M3U8:
|
|
text = (
|
|
"#EXTM3U\n"
|
|
f'#EXT-X-SESSION-DATA:DATA-ID="{data_id}",VALUE="{value}"\n'
|
|
"#EXT-X-STREAM-INF:BANDWIDTH=1000,RESOLUTION=640x360\n"
|
|
"video.m3u8\n"
|
|
)
|
|
return m3u8.loads(text)
|
|
|
|
|
|
def _b64(items: list[dict]) -> str:
|
|
return base64.b64encode(json.dumps(items).encode()).decode()
|
|
|
|
|
|
def _item(method: str, keyformat: Optional[str] = None) -> dict:
|
|
item = {"method": method, "uri": KEY_URI}
|
|
if keyformat:
|
|
item["keyformat"] = keyformat
|
|
return item
|
|
|
|
|
|
def _key(method: str, keyformat: Optional[str] = None) -> m3u8.model.Key:
|
|
return m3u8.model.Key(method=method, base_uri="", uri=KEY_URI, keyformat=keyformat)
|
|
|
|
|
|
def test_session_data_playready_urn_key_is_kept() -> None:
|
|
keys = HLS.parse_session_data_keys(_master(_b64([_item("SAMPLE-AES-CTR", PR_URN)])))
|
|
assert [k.keyformat for k in keys] == [PR_URN]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("method", "keyformat"),
|
|
[
|
|
("SAMPLE-AES-CTR", WV_URN),
|
|
("SAMPLE-AES-CTR", PR_NAME),
|
|
("SAMPLE-AES-CTR", PR_URN.upper()),
|
|
("AES-128", None),
|
|
("ISO-23001-7", None),
|
|
],
|
|
)
|
|
def test_session_data_supported_keys_are_kept(method: str, keyformat: Optional[str]) -> None:
|
|
keys = HLS.parse_session_data_keys(_master(_b64([_item(method, keyformat)])))
|
|
assert [k.keyformat for k in keys] == [keyformat]
|
|
|
|
|
|
def test_session_data_unknown_keyformat_is_dropped() -> None:
|
|
assert HLS.parse_session_data_keys(_master(_b64([_item("SAMPLE-AES", "com.example.drm")]))) == []
|
|
|
|
|
|
def test_session_data_plain_json_value_is_parsed() -> None:
|
|
manifest = _master("placeholder")
|
|
# a quoted attribute cannot hold raw JSON quotes, so set the value after parsing
|
|
manifest.session_data[0].value = json.dumps([_item("SAMPLE-AES-CTR", WV_URN)])
|
|
keys = HLS.parse_session_data_keys(manifest)
|
|
assert len(keys) == 1
|
|
|
|
|
|
def test_session_data_value_not_json_yields_nothing() -> None:
|
|
assert HLS.parse_session_data_keys(_master("not json at all")) == []
|
|
|
|
|
|
def test_session_data_other_data_id_is_ignored() -> None:
|
|
manifest = _master(_b64([_item("SAMPLE-AES-CTR", WV_URN)]), data_id="com.example.other")
|
|
assert HLS.parse_session_data_keys(manifest) == []
|
|
|
|
|
|
@pytest.mark.parametrize("playready_keyformat", [PR_URN, PR_NAME])
|
|
def test_filter_keys_follows_cdm(playready_keyformat: str) -> None:
|
|
wv = _key("SAMPLE-AES-CTR", WV_URN)
|
|
pr = _key("SAMPLE-AES-CTR", playready_keyformat)
|
|
assert HLS.filter_keys_for_cdm([wv, pr], cdm_type_stub("widevine")) == [wv]
|
|
assert HLS.filter_keys_for_cdm([wv, pr], cdm_type_stub("playready")) == [pr]
|
|
|
|
|
|
def test_filter_keys_preference_wins_over_cdm() -> None:
|
|
wv = _key("SAMPLE-AES-CTR", WV_URN)
|
|
pr = _key("SAMPLE-AES-CTR", PR_URN)
|
|
assert HLS.filter_keys_for_cdm([wv, pr], cdm_type_stub("widevine"), drm_preference="playready") == [pr]
|
|
assert HLS.filter_keys_for_cdm([wv, pr], cdm_type_stub("playready"), drm_preference="wv") == [wv]
|
|
|
|
|
|
def test_filter_keys_unsatisfied_preference_falls_back_to_cdm() -> None:
|
|
wv = _key("SAMPLE-AES-CTR", WV_URN)
|
|
assert HLS.filter_keys_for_cdm([wv], cdm_type_stub("widevine"), drm_preference="playready") == [wv]
|
|
|
|
|
|
def test_get_supported_key_none_method_returns_none() -> None:
|
|
assert HLS.get_supported_key([_key("NONE")]) is None
|
|
|
|
|
|
def test_get_supported_key_unsupported_only_raises() -> None:
|
|
with pytest.raises(NotImplementedError, match="SAMPLE-AES"):
|
|
HLS.get_supported_key([_key("SAMPLE-AES", "com.example.drm")])
|