From bae71c0194f5d71bcb765e72a3f0f9b9d2f5ed11 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Fri, 4 Sep 2026 22:19:21 -0600 Subject: [PATCH] fix(serve): resolve DRM keys for manifests with no ContentProtection PSSH Read the init segment server-side so server_cdm licenses every track's real KID, and re-license when the manifest PSSH misses it. --- tests/remote/unit/test_extract_pssh_strict.py | 89 ++++++++ tests/remote/unit/test_server_cdm_multikey.py | 8 +- unshackle/commands/dl.py | 17 ++ unshackle/core/api/handlers.py | 191 +++++++++++++----- unshackle/core/import_service.py | 15 +- unshackle/core/remote_service.py | 4 + 6 files changed, 263 insertions(+), 61 deletions(-) create mode 100644 tests/remote/unit/test_extract_pssh_strict.py diff --git a/tests/remote/unit/test_extract_pssh_strict.py b/tests/remote/unit/test_extract_pssh_strict.py new file mode 100644 index 00000000..1b50436f --- /dev/null +++ b/tests/remote/unit/test_extract_pssh_strict.py @@ -0,0 +1,89 @@ +"""extract_pssh_from_track must answer for the requested DRM system only. + +A loose match made the batch candidate loop a no-op: the first candidate always won, so the +loop licensed a PlayReady-only track as Widevine by accident. A Widevine device can still +license PlayReady-only content, but only as a deliberate step in the batch loop. +""" + +import json +from types import SimpleNamespace + +import pytest + +from unshackle.core.api import handlers + +pytestmark = pytest.mark.unit + + +class _WV: + def __init__(self, b64): + self._pssh = SimpleNamespace(dumps=lambda: b64) + + +_WV.__name__ = "Widevine" + + +class _PR: + def __init__(self, b64): + self.data = {"pssh_b64": b64} + + +_PR.__name__ = "PlayReady" + + +def test_extractor_is_type_strict(): + track = SimpleNamespace(drm=[_PR("pr")]) + assert handlers.extract_pssh_from_track(track, "playready") == "pr" + assert handlers.extract_pssh_from_track(track, "widevine") is None + + track = SimpleNamespace(drm=[_WV("wv"), _PR("pr")]) + assert handlers.extract_pssh_from_track(track, "widevine") == "wv" + assert handlers.extract_pssh_from_track(track, "playready") == "pr" + assert handlers.extract_pssh_from_track(SimpleNamespace(drm=None), "widevine") is None + + +async def _run_batch(monkeypatch, tracks, cdm_type, calls): + session = SimpleNamespace( + service_tag="EXAMPLE", + service_instance=SimpleNamespace(), + tracks=SimpleNamespace(get=tracks.get), + ) + + async def fake_get_session(session_id, request): + return session + + def fake_single(service, title, track, pssh_str, drm_type, request): + calls.append((pssh_str, drm_type)) + return {"00" * 16: "key"} + + monkeypatch.setattr(handlers, "get_validated_session", fake_get_session) + monkeypatch.setattr(handlers, "server_cdm_allowed", lambda request, tag: True) + monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None) + monkeypatch.setattr(handlers, "find_title_for_track", lambda tid, session: SimpleNamespace()) + monkeypatch.setattr(handlers, "detect_cdm_type_for_service", lambda tag, cfg: cdm_type) + monkeypatch.setattr(handlers, "handle_single_server_cdm", fake_single) + monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False) + resp = await handlers.session_license_handler( + {"mode": "server_cdm", "track_ids": list(tracks), "drm_type": cdm_type}, "sess", None + ) + return json.loads(resp.body) + + +async def test_playready_only_track_uses_playready_when_preferred(monkeypatch): + calls: list = [] + payload = await _run_batch(monkeypatch, {"vid": SimpleNamespace(id="vid", drm=[_PR("pr")])}, "playready", calls) + assert calls == [("pr", "playready")] + assert payload["drm_types"] == {"vid": "playready"} + + +async def test_playready_only_track_converted_for_widevine_device(monkeypatch): + calls: list = [] + payload = await _run_batch(monkeypatch, {"vid": SimpleNamespace(id="vid", drm=[_PR("pr")])}, "widevine", calls) + assert calls == [("pr", "widevine")] # deliberate: Widevine.__init__ converts the PlayReady PSSH + assert payload["drm_types"] == {"vid": "widevine"} + + +async def test_widevine_pssh_wins_for_widevine_device(monkeypatch): + calls: list = [] + await _run_batch(monkeypatch, {"vid": SimpleNamespace(id="vid", drm=[_PR("pr"), _WV("wv")])}, "widevine", calls) + assert calls == [("wv", "widevine")] diff --git a/tests/remote/unit/test_server_cdm_multikey.py b/tests/remote/unit/test_server_cdm_multikey.py index f4d6d8d2..6c072bad 100644 --- a/tests/remote/unit/test_server_cdm_multikey.py +++ b/tests/remote/unit/test_server_cdm_multikey.py @@ -16,6 +16,7 @@ import pytest import pywidevine.pssh as wv_pssh_mod import unshackle.core.cdm as cdm_mod +import unshackle.core.cdm.detect as detect_mod import unshackle.core.drm as drm_mod from unshackle.core.api import handlers @@ -50,7 +51,8 @@ def playready_env(monkeypatch): monkeypatch.setattr(pr_pssh_mod, "PSSH", _FakePRPSSH) monkeypatch.setattr(drm_mod, "PlayReady", _FakePlayReady) monkeypatch.setattr(cdm_mod, "load_cdm", lambda *a, **k: object()) - monkeypatch.setattr(handlers, "ensure_track_drm", lambda track: None) + 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.config, "serve", {"users": {}}, raising=False) @@ -94,7 +96,7 @@ def test_widevine_still_uses_vault_shortcut(monkeypatch): monkeypatch.setattr(wv_pssh_mod, "PSSH", _FakeWvPSSH) monkeypatch.setattr(drm_mod, "Widevine", _FakeWidevine) - monkeypatch.setattr(handlers, "ensure_track_drm", lambda track: None) + 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: {kid.hex: "vaultkey"}) monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False) @@ -132,7 +134,7 @@ async def test_batch_shares_full_bundle_per_pssh(monkeypatch): monkeypatch.setattr(handlers, "get_validated_session", fake_get_session) monkeypatch.setattr(handlers, "server_cdm_allowed", lambda request, tag: True) - monkeypatch.setattr(handlers, "ensure_track_drm", lambda track: None) + monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None) monkeypatch.setattr(handlers, "find_title_for_track", lambda tid, session: SimpleNamespace()) monkeypatch.setattr(handlers, "detect_cdm_type_for_service", lambda tag, cfg: "playready") monkeypatch.setattr(handlers, "drm_preference_name", lambda track: None) diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 2473571d..360539fd 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -4035,9 +4035,26 @@ class dl: drm.content_keys[kid] = content_key self.LICENSE_KEY_CACHE[kid] = content_key + def missing_track_key() -> bool: + return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys) + + if missing_track_key(): + try: + licence( + drm_system="playready" if drm.__class__.__name__ == "PlayReady" else "widevine", + challenge=b"", + ) + except Exception as e: + self.log.debug(f"Server CDM licence with an empty challenge failed: {e!r}") + if not drm.content_keys: self.log.warning("Server CDM did not resolve any keys for this track") return + if track_kid and track_kid not in drm.content_keys: + msg = f"No Content Key for KID {track_kid.hex} was returned by the server CDM" + if isinstance(drm, PlayReady): + raise PlayReady.Exceptions.CEKNotFound(msg) + raise Widevine.Exceptions.CEKNotFound(msg) svc = getattr(self, "_remote_service", None) server_drm_type = getattr(svc, "_server_cdm_type", None) if svc else None drm_name = {"widevine": "Widevine", "playready": "PlayReady"}.get( diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 31830dc9..304b0518 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -3424,30 +3424,74 @@ def find_title_for_track(track_id: str, session: Any) -> Any: def extract_pssh_from_track(track: Any, drm_type: str) -> Optional[str]: - """Extract PSSH base64 string from a track's DRM objects.""" - if not track.drm: - return None - pssh_b64 = None - for drm_obj in track.drm: + """The base64 PSSH of the track's DRM object for `drm_type`; None when the track has none of that system.""" + for drm_obj in track.drm or []: drm_class = drm_obj.__class__.__name__ - if drm_class == "Widevine" and hasattr(drm_obj, "_pssh") and drm_obj._pssh: - if hasattr(drm_obj._pssh, "dumps"): - pssh_b64 = drm_obj._pssh.dumps() - if drm_type == "widevine": - break - elif drm_class == "PlayReady": - if hasattr(drm_obj, "data") and drm_obj.data.get("pssh_b64"): - pssh_b64 = drm_obj.data["pssh_b64"] - if drm_type == "playready": - break - return pssh_b64 + if drm_type == "widevine" and drm_class == "Widevine": + pssh = getattr(drm_obj, "_pssh", None) + if pssh and hasattr(pssh, "dumps"): + return pssh.dumps() + elif drm_type == "playready" and drm_class == "PlayReady": + pssh_b64 = getattr(drm_obj, "data", {}).get("pssh_b64") + if pssh_b64: + return pssh_b64 + return None -def ensure_track_drm(track: Any) -> None: +def fetch_init_segment(track: Any, session: Any = None) -> Optional[bytes]: + """The init segment of a DASH track; None for a track that is not DASH, or when the fetch fails.""" + from unshackle.core.manifests import DASH as DASHManifest + + dash = (getattr(track, "data", None) or {}).get("dash") or {} + if not all(dash.get(k) is not None for k in ("period", "adaptation_set", "representation", "manifest")): + return None + + session = getattr(track, "session", None) or session + if session is None: + return None + + try: + return DASHManifest.get_period_segments( + period=dash["period"], + adaptation_set=dash["adaptation_set"], + representation=dash["representation"], + manifest=dash["manifest"], + track=track, + track_url=track.url, + session=session, + )[0] + except Exception as e: + log.warning(f"Init segment fetch failed for {track.id}: {e!r}") + return None + + +def drm_from_init_segment(track: Any, session: Any = None, init_data: Optional[bytes] = None) -> list: + """Read the DRM headers out of a DASH track's init segment. + + Some manifests carry no ContentProtection PSSH; the real header reaches the player only + in the init segment, so the server must read it there too. + """ + from unshackle.core.drm import PlayReady, Widevine + + if init_data is None: + init_data = fetch_init_segment(track, session) + if not init_data: + return [] + + drms = [] + for drm_cls in (Widevine, PlayReady): + try: + drms.append(drm_cls.from_init_data(init_data)) + except Exception as e: + log.debug(f"No usable {drm_cls.__name__} PSSH in the init segment of {track.id}: {e!r}") + return drms + + +def ensure_track_drm(track: Any, session: Any = None, init_data: Optional[bytes] = None) -> None: """Extract DRM from manifest data if track has none. - Supports DASH (ContentProtection elements), HLS (EXT-X-KEY from - playlist fetch), and ISM (ProtectionHeader elements). + Supports DASH (ContentProtection elements, then the init segment), HLS + (EXT-X-KEY from playlist fetch), and ISM (ProtectionHeader elements). """ if track.drm: return @@ -3462,6 +3506,10 @@ def ensure_track_drm(track: Any) -> None: if track.drm: return + track.drm = drm_from_init_segment(track, session, init_data) or None + if track.drm: + return + if track.data.get("hls") and track.url: try: import m3u8 @@ -3469,7 +3517,8 @@ def ensure_track_drm(track: Any) -> None: from unshackle.core.drm import PlayReady, Widevine from unshackle.core.manifests import HLS - playlist = m3u8.load(track.url) + http = getattr(track, "session", None) or session + playlist = m3u8.loads(http.get(track.url).text, uri=track.url) if http else m3u8.load(track.url) keys = [k for k in (playlist.keys or []) + (playlist.session_keys or []) if k is not None] for key in keys: try: @@ -3599,9 +3648,10 @@ def handle_single_server_cdm( import base64 from unshackle.core.cdm import load_cdm + from unshackle.core.cdm.detect import is_playready_cdm, is_widevine_cdm from unshackle.core.config import config as app_config - ensure_track_drm(track) + ensure_track_drm(track, getattr(service, "session", None)) if not pssh_b64: pssh_b64 = extract_pssh_from_track(track, drm_type) @@ -3623,6 +3673,8 @@ def handle_single_server_cdm( device_name = resolve_device_name(user_config, drm_type, service.__class__.__name__) 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( cdm=cdm, certificate=lambda challenge, **_: None, @@ -3646,6 +3698,8 @@ def handle_single_server_cdm( return vault_keys 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( cdm=cdm, certificate=lambda challenge, **_: service.get_widevine_service_certificate( @@ -3789,12 +3843,46 @@ async def session_license_handler( drm_type_by_pssh: Dict[str, str] = {} actual_drm_type: Optional[str] = None + def license_track( + track: Any, title: Any, candidates: list + ) -> tuple[Dict[str, str], Optional[str], Optional[str]]: + """`(keys, drm_type, pssh)` for the track's current DRM, licensing once per unique PSSH.""" + for candidate in candidates: + pssh_str = extract_pssh_from_track(track, candidate) + if not pssh_str and candidate == "widevine": + pssh_str = extract_pssh_from_track(track, "playready") + if pssh_str: + log.info( + f"Track {sanitize_log(str(track.id)[:12])} has only a PlayReady PSSH, licensing it with the Widevine device" + ) + if pssh_str: + break + else: + return {}, None, None + + if pssh_str not in keys_by_pssh: + keys_by_pssh[pssh_str] = {} + try: + keys = handle_single_server_cdm(service, title, track, pssh_str, candidate, request) + if keys: + keys_by_pssh[pssh_str] = keys + drm_type_by_pssh[pssh_str] = candidate + except SystemExit: + log.warning( + f"Service exited while resolving keys for track {sanitize_log(str(track.id)[:12])}, skipping" + ) + except (Exception, SystemExit) as e: + log.warning(f"Failed to resolve keys for track {sanitize_log(str(track.id)[:12])}: {e}") + return keys_by_pssh[pssh_str], drm_type_by_pssh.get(pssh_str), pssh_str + for tid in track_ids: track = session.tracks.get(tid) if not track: continue - ensure_track_drm(track) + svc_session = getattr(service, "session", None) + init_data = fetch_init_segment(track, svc_session) + ensure_track_drm(track, svc_session, init_data) if not track.drm: continue @@ -3822,37 +3910,34 @@ async def session_license_handler( "but the server has no device for it, using the configured DRM instead" ) - track_drm_type = None - pssh_str = None - for candidate in candidates: - pssh_str = extract_pssh_from_track(track, candidate) - if pssh_str: - track_drm_type = candidate - break + keys, track_drm_type, pssh_str = license_track(track, title, candidates) - if not pssh_str or not track_drm_type: - continue + track_kid = None + if init_data: + try: + track_kid = track.get_key_id(init_data) + except Exception as e: + log.debug(f"KID probe failed for {sanitize_log(tid[:12])}: {e!r}") + if track_kid and track_kid.hex not in keys: + manifest_drm = track.drm + track.drm = drm_from_init_segment(track, init_data=init_data) or manifest_drm + init_keys, init_drm_type, init_pssh = license_track(track, title, candidates) + if init_pssh != pssh_str: + log.info( + f"The manifest PSSH gave no content key for KID {track_kid.hex} of track " + f"{sanitize_log(tid[:12])}, tried the init segment PSSH" + ) + if track_kid.hex in init_keys: + keys, track_drm_type = init_keys, init_drm_type + else: + track.drm = manifest_drm + log.warning(f"No content key for KID {track_kid.hex} of track {sanitize_log(tid[:12])}") - if pssh_str in keys_by_pssh: - all_keys[tid] = keys_by_pssh[pssh_str] - if pssh_str in drm_type_by_pssh: - drm_types[tid] = drm_type_by_pssh[pssh_str] - continue - - keys_by_pssh[pssh_str] = {} - try: - keys = handle_single_server_cdm(service, title, track, pssh_str, track_drm_type, request) - if keys: - keys_by_pssh[pssh_str] = keys - all_keys[tid] = keys - if track_drm_type: - drm_types[tid] = track_drm_type - drm_type_by_pssh[pssh_str] = track_drm_type - actual_drm_type = track_drm_type - except SystemExit: - log.warning(f"Service exited while resolving keys for track {sanitize_log(tid[:12])}, skipping") - except (Exception, SystemExit) as e: - log.warning(f"Failed to resolve keys for track {sanitize_log(tid[:12])}: {e}") + if keys: + all_keys[tid] = keys + if keys and track_drm_type: + drm_types[tid] = track_drm_type + actual_drm_type = track_drm_type response: Dict[str, Any] = {"keys": all_keys} if actual_drm_type: @@ -3878,8 +3963,8 @@ async def session_license_handler( pssh_b64 = data.get("pssh") if pssh_b64: - if not track.drm: - track.drm = [] + cls_name = {"playready": "PlayReady", "widevine": "Widevine"}.get(drm_type) + track.drm = [d for d in (track.drm or []) if d.__class__.__name__ != cls_name] if drm_type == "playready": track.pr_pssh = pssh_b64 from pyplayready.system.pssh import PSSH as PlayReadyPSSH diff --git a/unshackle/core/import_service.py b/unshackle/core/import_service.py index 8838e303..e3cc7f26 100644 --- a/unshackle/core/import_service.py +++ b/unshackle/core/import_service.py @@ -355,7 +355,7 @@ class ImportService: kid_hexes = [kid.hex for kid in pool] for track in title.tracks: - if not isinstance(track, (Video, Audio)) or not self.track_is_encrypted(track): + if not isinstance(track, (Video, Audio)) or not self.track_is_encrypted(track, self.session): continue drm_obj = track.drm[0] if track.drm else RemoteService.create_drm_stub(system, kid_hexes) for kid, key in pool.items(): @@ -364,13 +364,14 @@ class ImportService: self._server_cdm_type = drm_obj.__class__.__name__.lower() @staticmethod - def track_is_encrypted(track: Any) -> bool: - """True if the track carries DRM or its manifest declares protection. + def track_is_encrypted(track: Any, session: Optional[requests.Session] = None) -> bool: + """True if the track carries DRM, its manifest declares protection, or its init segment does. This method examines ISM as well as DASH, because ISM.download_track reads only ``track.drm``. A rung that misses content key injection here downloads encrypted and - muxes without error. DASH re-derives - its DRM from the manifest elements, so it survives the same omission. + muxes without error. A DASH manifest with no ContentProtection can still describe an + encrypted track, with the PSSH only in the init segment, so that case probes the init + segment with `session`. """ if track.drm: return True @@ -380,6 +381,10 @@ class ImportService: for element in (dash.get("representation"), dash.get("adaptation_set")): if element is not None and element.findall("ContentProtection"): return True + from unshackle.core.api.handlers import drm_from_init_segment + + if drm_from_init_segment(track, session): + return True ism = data.get("ism") if ism: manifest = ism.get("manifest") diff --git a/unshackle/core/remote_service.py b/unshackle/core/remote_service.py index 3ca7d630..48b54089 100644 --- a/unshackle/core/remote_service.py +++ b/unshackle/core/remote_service.py @@ -1112,6 +1112,8 @@ class RemoteService: for track in title.tracks: track_keys = keys_by_track.get(str(track.id), {}) if not track_keys: + if str(track.id) in track_ids: + self.log.warning(f"Server CDM returned no content keys for track {track.id}") continue kid_list = list(track_keys.keys()) @@ -1256,6 +1258,8 @@ class RemoteService: return challenge except Exception as e: self.log.warning("server_cdm license failed: %s", e) + else: + self.log.warning(f"Track {track.id} has no {drm_type} PSSH to send to the server CDM") return challenge payload = {