fix(api): let the server CDM see service-supplied keys and HLS track DRM

tests/remote/unit/test_ensure_track_drm_hls_descriptor.py

@@ -0,0 +1,51 @@
+"""The server must read HLS media-playlist DRM for a track the service built by hand.
+
+A service can make an HLS `Video` from a media playlist URL without going through
+`HLS.to_tracks`, so the track carries the HLS descriptor but no `data["hls"]`. The
+batch server CDM route then saw no DRM, told the client the track carried none, and left
+the client to license the same track again through the single-track route.
+"""
+
+import base64
+from types import SimpleNamespace
+from uuid import UUID
+
+import pytest
+from pywidevine.pssh import PSSH as WvPSSH
+
+from unshackle.core.api import handlers
+from unshackle.core.drm import Widevine
+from unshackle.core.tracks import Track
+
+pytestmark = pytest.mark.unit
+
+KID = UUID("161219ec3df64e0eadbb3827c0ccc46d")
+
+
+def _media_playlist() -> str:
+    pssh = WvPSSH.new(key_ids=[KID], system_id=WvPSSH.SystemId.Widevine).dumps()
+    return "\n".join(
+        [
+            "#EXTM3U",
+            "#EXT-X-VERSION:6",
+            "#EXT-X-TARGETDURATION:6",
+            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,",
+            "seg0.ts",
+            "#EXT-X-ENDLIST",
+        ]
+    )
+
+
+def test_hls_descriptor_alone_reaches_the_media_playlist_keys():
+    session = SimpleNamespace(get=lambda url: SimpleNamespace(text=_media_playlist()))
+    track = SimpleNamespace(
+        id="vid", url="https://cdn.example.invalid/media.m3u8", descriptor=Track.Descriptor.HLS, data={}, drm=None
+    )
+
+    handlers.ensure_track_drm(track, session)
+
+    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"))

tests/remote/unit/test_server_cdm_service_injected_keys.py

@@ -0,0 +1,117 @@
+"""The server CDM must return keys a service wrote onto the track's own DRM object.
+
+Some services skip the CDM exchange: the licence callback fetches the keys itself, writes
+them into `track.drm[0].content_keys`, and returns a value the CDM cannot parse. `dl.py`
+accepts the failed parse because the object it licensed already holds keys. The server
+licenses a rebuilt object, so it must read the track's own objects too.
+"""
+
+import base64
+from types import SimpleNamespace
+from uuid import UUID
+
+import pytest
+from pyplayready.system.pssh import PSSH as PRPSSH
+from pywidevine.pssh import PSSH as WvPSSH
+
+import unshackle.core.cdm as cdm_mod
+import unshackle.core.cdm.detect as detect_mod
+from unshackle.core.api import handlers
+from unshackle.core.drm import PlayReady, Widevine
+
+pytestmark = pytest.mark.unit
+
+KID_A = UUID("161219ec3df64e0eadbb3827c0ccc46d")
+KEY_A = "0" * 31 + "1"
+
+
+class _BypassCdm:
+    """A CDM that cannot parse the service's stand-in licence, like pywidevine or decrypt-labs."""
+
+    service_certificate_challenge = b"cert"
+
+    def open(self):
+        return b"sid"
+
+    def close(self, session_id):
+        pass
+
+    def set_pssh_b64(self, *a, **k):
+        pass
+
+    def get_license_challenge(self, session_id, pssh):
+        return b"challenge"
+
+    def parse_license(self, session_id, licence):
+        raise ValueError(f"not a licence: {licence!r}")
+
+    def get_keys(self, session_id, type_=None):
+        return []
+
+
+@pytest.fixture
+def server_env(monkeypatch):
+    monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", lambda *a, **k: _BypassCdm())
+    monkeypatch.setattr(detect_mod, "is_widevine_cdm", lambda cdm: True)
+    monkeypatch.setattr(detect_mod, "is_playready_cdm", lambda cdm: True)
+    monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None)
+    monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev")
+    monkeypatch.setattr(handlers, "check_vaults", lambda kids, name: None)
+    monkeypatch.setattr(handlers, "cache_to_vaults", lambda keys, name: None)
+    monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False)
+
+
+def _bypass_service(kind: str):
+    """A service that writes the key onto the track's own object and returns a stand-in licence."""
+
+    def licence(*, challenge, title, track, **_):
+        track.drm[0].content_keys[KID_A] = KEY_A
+        return b"bypass"
+
+    return SimpleNamespace(
+        get_widevine_service_certificate=lambda **_: None,
+        **{f"get_{kind}_license": licence},
+    )
+
+
+def test_widevine_bypass_keys_reach_the_client(server_env):
+    pssh = WvPSSH.new(key_ids=[KID_A], system_id=WvPSSH.SystemId.Widevine)
+    track = SimpleNamespace(id="vid", drm=[Widevine(pssh=pssh, kid=KID_A)])
+
+    keys = handlers.handle_single_server_cdm(
+        _bypass_service("widevine"), SimpleNamespace(), track, pssh.dumps(), "widevine", None
+    )
+
+    assert keys == {KID_A.hex: KEY_A}
+
+
+def test_playready_bypass_keys_reach_the_client(server_env):
+    xml = (
+        '<WRMHEADER xmlns="http://schemas.microsoft.com/DRM/2007/03/PlayReadyHeader" version="4.2.0.0">'
+        "<DATA><PROTECTINFO></PROTECTINFO><LA_URL>https://example.invalid/</LA_URL></DATA></WRMHEADER>"
+    )
+    record = xml.encode("utf-16-le")
+    pro = PRPSSH.PlayreadyHeader.build(
+        {"length": 10 + len(record), "records": [{"type": 1, "length": len(record), "data": record}]}
+    )
+    pro_b64 = base64.b64encode(pro).decode()
+    track = SimpleNamespace(id="vid", drm=[PlayReady(pssh=PRPSSH(pro), kid=KID_A, pssh_b64=pro_b64)])
+
+    keys = handlers.handle_single_server_cdm(
+        _bypass_service("playready"), SimpleNamespace(), track, pro_b64, "playready", None
+    )
+
+    assert keys == {KID_A.hex: KEY_A}
+
+
+def test_failed_licence_without_injected_keys_still_raises(server_env):
+    """The fold is not a licence to swallow real failures."""
+    pssh = WvPSSH.new(key_ids=[KID_A], system_id=WvPSSH.SystemId.Widevine)
+    track = SimpleNamespace(id="vid", drm=[Widevine(pssh=pssh, kid=KID_A)])
+    service = SimpleNamespace(
+        get_widevine_service_certificate=lambda **_: None,
+        get_widevine_license=lambda **_: b"garbage",
+    )
+
+    with pytest.raises(ValueError, match="not a licence"):
+        handlers.handle_single_server_cdm(service, SimpleNamespace(), track, pssh.dumps(), "widevine", None)

unshackle/core/api/handlers.py

@@ -3528,7 +3528,7 @@ def ensure_track_drm(track: Any, session: Any = None, init_data: Optional[bytes]
         if track.drm:
             return
 
-    if track.data.get("hls") and track.url:
+    if (track.data.get("hls") or descriptor_name(track) == "HLS") and track.url:
         try:
             import m3u8
 
@@ -3654,6 +3654,36 @@ def cache_to_vaults(keys: Dict[str, str], service_name: str) -> None:
         log.warning(f"Failed to cache keys to vaults: {e}")
 
 
+def track_injected_keys(track: Any, drm_class: str) -> Dict[str, str]:
+    """`KID:KEY` pairs a service wrote onto the track's own DRM objects of `drm_class`.
+
+    Some services skip the CDM exchange: the licence callback fetches the keys itself,
+    writes them into ``track.drm[0].content_keys`` and returns a value the CDM cannot
+    parse. ``dl.py`` accepts that failed parse because the object it licensed holds the
+    keys. The server licenses a rebuilt object, so it reads the track's objects instead.
+    """
+    keys: Dict[str, str] = {}
+    for drm in getattr(track, "drm", None) or []:
+        if drm.__class__.__name__ != drm_class:
+            continue
+        for kid, key in (getattr(drm, "content_keys", None) or {}).items():
+            keys[kid.hex if hasattr(kid, "hex") else str(kid).replace("-", "")] = key
+    return keys
+
+
+def license_with_track_fallback(drm: Any, track: Any, **kwargs: Any) -> Dict[str, str]:
+    """Run `drm.get_content_keys` and merge in the keys the service put on the track itself."""
+    try:
+        drm.get_content_keys(**kwargs)
+    except Exception:
+        if not track_injected_keys(track, drm.__class__.__name__):
+            raise
+    return {
+        **track_injected_keys(track, drm.__class__.__name__),
+        **{kid.hex: key for kid, key in drm.content_keys.items()},
+    }
+
+
 def handle_single_server_cdm(
     service: Any,
     title: Any,
@@ -3697,14 +3727,15 @@ def handle_single_server_cdm(
         cdm = load_cdm(device_name, service_name=service.__class__.__name__)
         if not is_playready_cdm(cdm):
             raise APIError(APIErrorCode.INVALID_INPUT, f"CDM device '{device_name}' is not a PlayReady device")
-        pr_drm.get_content_keys(
+        keys = license_with_track_fallback(
+            pr_drm,
+            track,
             cdm=cdm,
             certificate=lambda challenge, **_: None,
             licence=lambda **kw: service.get_playready_license(
                 **declared_kwargs(service.get_playready_license, {**kw, "title": title, "track": track})
             ),
         )
-        keys = {kid.hex: key for kid, key in pr_drm.content_keys.items()}
     elif drm_type == "widevine":
         from pywidevine.pssh import PSSH as WvPSSH
 
@@ -3724,7 +3755,9 @@ def handle_single_server_cdm(
         cdm = load_cdm(device_name, service_name=service.__class__.__name__)
         if not is_widevine_cdm(cdm):
             raise APIError(APIErrorCode.INVALID_INPUT, f"CDM device '{device_name}' is not a Widevine device")
-        wv_drm.get_content_keys(
+        keys = license_with_track_fallback(
+            wv_drm,
+            track,
             cdm=cdm,
             certificate=lambda challenge, **_: service.get_widevine_service_certificate(
                 challenge=challenge, title=title, track=track
@@ -3733,7 +3766,6 @@ def handle_single_server_cdm(
                 **declared_kwargs(service.get_widevine_license, {**kw, "title": title, "track": track})
             ),
         )
-        keys = {kid.hex: key for kid, key in wv_drm.content_keys.items()}
     else:
         raise APIError(
             APIErrorCode.INVALID_PARAMETERS,