From 1aa6315dc356a5014d6b896347a71e610ce3a958 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sun, 23 Aug 2026 16:32:23 -0600 Subject: [PATCH] feat(drm): let a track pin the DRM system it licenses with Adds Track.drm_preference so a service can pin a track to Widevine or PlayReady when a manifest advertises both but only one will license it. Honoured in DASH/HLS selection, the dl licence path and the remote/API track round-trip. Also fixes quality-keyed CDM maps sending audio to the default entry, and a licence callable binding the CDM before the swap. --- tests/remote/unit/test_handlers_serialize.py | 26 +++++++ .../unit/test_remote_service_helpers.py | 17 +++++ unshackle/commands/dl.py | 31 +++++++-- unshackle/core/api/handlers.py | 68 ++++++++++++------- unshackle/core/manifests/dash.py | 16 ++++- unshackle/core/manifests/hls.py | 21 +++++- unshackle/core/remote_service.py | 10 ++- unshackle/core/tracks/track.py | 43 +++++++++++- 8 files changed, 193 insertions(+), 39 deletions(-) diff --git a/tests/remote/unit/test_handlers_serialize.py b/tests/remote/unit/test_handlers_serialize.py index 3fa06ea3..1d967a1a 100644 --- a/tests/remote/unit/test_handlers_serialize.py +++ b/tests/remote/unit/test_handlers_serialize.py @@ -7,6 +7,7 @@ from langcodes import Language from unshackle.core.api.errors import APIError, APIErrorCode from unshackle.core.api.handlers import ( + drm_preference_name, sanitize_log, search_handler, serialize_audio_track, @@ -161,6 +162,31 @@ def test_serialize_audio_track_basic() -> None: assert d["descriptive"] is False +def test_serialize_video_track_drm_preference_round_trip() -> None: + d = serialize_video_track(video()) + assert d["drm_preference"] is None + v = video() + v.drm_preference = "pr" + assert serialize_video_track(v)["drm_preference"] == "pr" + + +def test_serialize_audio_track_drm_preference() -> None: + a = audio() + a.drm_preference = "widevine" + assert serialize_audio_track(a)["drm_preference"] == "widevine" + + +def test_drm_preference_name_maps_short_forms() -> None: + v = video() + assert drm_preference_name(v) is None + v.drm_preference = "pr" + assert drm_preference_name(v) == "playready" + v.drm_preference = "wv" + assert drm_preference_name(v) == "widevine" + v.drm_preference = "widevine" + assert drm_preference_name(v) == "widevine" + + def test_serialize_subtitle_track_basic() -> None: d = serialize_subtitle_track(subtitle(forced=True)) assert d["id"] == "sub-001" diff --git a/tests/remote/unit/test_remote_service_helpers.py b/tests/remote/unit/test_remote_service_helpers.py index 55cf5ab9..f496dc60 100644 --- a/tests/remote/unit/test_remote_service_helpers.py +++ b/tests/remote/unit/test_remote_service_helpers.py @@ -108,6 +108,23 @@ def test_build_tracks_aggregates() -> None: assert len(t.subtitles) == 1 +def test_build_tracks_restores_drm_preference() -> None: + data = { + "video": [{"id": "v", "codec": "AVC", "width": 1280, "height": 720, "drm_preference": "playready"}], + "audio": [{"id": "a", "codec": "AAC", "channels": 2, "drm_preference": "wv"}], + "subtitles": [], + } + t = build_tracks(data) + assert t.videos[0].drm_preference == "playready" + assert t.audio[0].drm_preference == "wv" + + +def test_build_tracks_ignores_unknown_drm_preference() -> None: + data = {"video": [{"id": "v", "codec": "AVC", "width": 1, "height": 1, "drm_preference": "fairplay"}]} + t = build_tracks(data) + assert t.videos[0].drm_preference is None + + def test_match_track_by_id() -> None: a = deserialize_video({"id": "v1", "codec": "AVC", "width": 1920, "height": 1080}) b = deserialize_video({"id": "v2", "codec": "AVC", "width": 1280, "height": 720}) diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 7d2884e2..91663c6b 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -2852,9 +2852,8 @@ class dl: track=track, ), licence=partial( - service.get_playready_license - if is_playready_cdm(self.cdm) - else service.get_widevine_license, + self.service_licence, + service, title=title, track=track, ), @@ -3919,6 +3918,18 @@ class dl: with cls.DRM_LOCKS_GUARD: return cls.DRM_LOCKS.setdefault(key or type(drm).__name__, Lock()) + def service_licence(self, service: Service, drm_system: Optional[str] = None, **kwargs: Any) -> Any: + """Call the service's licence function for the track's DRM system. + + prepare_drm passes drm_system so the choice follows the track, not the loaded CDM: + a track downloading at the same time can swap self.cdm. + """ + if drm_system is None: + drm_system = "playready" if is_playready_cdm(self.cdm) else "widevine" + if drm_system == "playready": + return service.get_playready_license(**kwargs) + return service.get_widevine_license(**kwargs) + def prepare_drm( self, drm: DRM_T, @@ -3999,6 +4010,12 @@ class dl: track_quality = None if isinstance(track, Video) and track.height: track_quality = track.height + elif title is not None: + track_quality = max((v.height for v in title.tracks.videos if v.height), default=None) + + track_cdm = self.cdm + + licence = partial(licence, drm_system="playready" if isinstance(drm, PlayReady) else "widevine") if not server_cdm: if isinstance(drm, Widevine): @@ -4010,6 +4027,7 @@ class dl: else: self.log.info("Switching to Widevine CDM for Widevine content") self.cdm = widevine_cdm + track_cdm = widevine_cdm else: raise ValueError(f"Title needs a Widevine CDM but {self.service} is configured with PlayReady.") @@ -4022,6 +4040,7 @@ class dl: else: self.log.info("Switching to PlayReady CDM for PlayReady content") self.cdm = playready_cdm + track_cdm = playready_cdm else: raise ValueError(f"Title needs a PlayReady CDM but {self.service} is configured with Widevine.") @@ -4148,9 +4167,9 @@ class dl: try: if self.service == "NF": - drm.get_NF_content_keys(cdm=self.cdm, licence=licence, certificate=certificate) + drm.get_NF_content_keys(cdm=track_cdm, licence=licence, certificate=certificate) else: - drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate) + drm.get_content_keys(cdm=track_cdm, licence=licence, certificate=certificate) except Exception as e: if drm.content_keys: self.log.debug(f"License call failed but keys already in content_keys: {e}") @@ -4338,7 +4357,7 @@ class dl: from_vaults = drm.content_keys.copy() try: - drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate) + drm.get_content_keys(cdm=track_cdm, licence=licence, certificate=certificate) except Exception as e: if drm.content_keys: self.log.debug(f"License call failed but keys already in content_keys: {e}") diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 49f477ab..d311007c 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -551,6 +551,14 @@ def descriptor_name(track: Any) -> Optional[str]: return enum_name(descriptor) if descriptor else None +def drm_preference_name(track: Any) -> Optional[str]: + """The track's drm_preference as a full DRM system name, "widevine" or "playready"; None when unset.""" + preference = getattr(track, "drm_preference", None) + if not preference: + return None + return {"wv": "widevine", "pr": "playready"}.get(preference, preference) + + def serialize_video_track(track: Video, include_url: bool = False) -> Dict[str, Any]: """Convert video track to JSON-serializable dict.""" codec_name = enum_name(track.codec) @@ -569,6 +577,7 @@ def serialize_video_track(track: Video, include_url: bool = False) -> Dict[str, "range_display": DYNAMIC_RANGE_MAP.get(range_name, range_name), "language": str(track.language) if track.language else None, "drm": serialize_drm(track.drm) if hasattr(track, "drm") and track.drm else None, + "drm_preference": getattr(track, "drm_preference", None), "descriptor": descriptor_name(track), } if include_url and hasattr(track, "url") and track.url: @@ -614,6 +623,7 @@ def serialize_audio_track(track: Audio, include_url: bool = False, is_original: "atmos": track.atmos if hasattr(track, "atmos") else False, "descriptive": track.descriptive if hasattr(track, "descriptive") else False, "drm": serialize_drm(track.drm) if hasattr(track, "drm") and track.drm else None, + "drm_preference": getattr(track, "drm_preference", None), "descriptor": descriptor_name(track), } if include_url and hasattr(track, "url") and track.url: @@ -2783,6 +2793,7 @@ async def session_license_handler( config_cdm_type = detect_cdm_type_for_service(service_tag, app_config) all_keys: Dict[str, Dict[str, str]] = {} + drm_types: Dict[str, str] = {} seen_pssh: set[str] = set() actual_drm_type: Optional[str] = None @@ -2797,41 +2808,45 @@ async def session_license_handler( title = find_title_for_track(tid, session) + if config_cdm_type == "playready": + candidates = ["playready", "widevine"] + elif config_cdm_type == "widevine": + candidates = ["widevine", "playready"] + else: + candidates = [ + name + for name, has_device in (("widevine", has_wv_device), ("playready", has_pr_device)) + if has_device + ] + + preferred = drm_preference_name(track) + if preferred: + if preferred in candidates: + candidates.remove(preferred) + candidates.insert(0, preferred) + else: + log.warning( + f"Track {sanitize_log(tid[:12])} wants {preferred} DRM " + "but the server has no device for it, using the configured DRM instead" + ) + track_drm_type = None pssh_str = None - if config_cdm_type == "playready": - pssh_str = extract_pssh_from_track(track, "playready") + for candidate in candidates: + pssh_str = extract_pssh_from_track(track, candidate) if pssh_str: - track_drm_type = "playready" - if not pssh_str: - pssh_str = extract_pssh_from_track(track, "widevine") - if pssh_str: - track_drm_type = "widevine" - elif config_cdm_type == "widevine": - pssh_str = extract_pssh_from_track(track, "widevine") - if pssh_str: - track_drm_type = "widevine" - if not pssh_str: - pssh_str = extract_pssh_from_track(track, "playready") - if pssh_str: - track_drm_type = "playready" - else: - if has_wv_device: - pssh_str = extract_pssh_from_track(track, "widevine") - if pssh_str: - track_drm_type = "widevine" - if not pssh_str and has_pr_device: - pssh_str = extract_pssh_from_track(track, "playready") - if pssh_str: - track_drm_type = "playready" + track_drm_type = candidate + break if not pssh_str or not track_drm_type: continue if pssh_str in seen_pssh: - for prev_keys in all_keys.values(): + for prev_tid, prev_keys in all_keys.items(): if prev_keys: all_keys[tid] = prev_keys + if prev_tid in drm_types: + drm_types[tid] = drm_types[prev_tid] break continue seen_pssh.add(pssh_str) @@ -2841,6 +2856,7 @@ async def session_license_handler( if keys: all_keys[tid] = keys if track_drm_type: + drm_types[tid] = 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") @@ -2850,6 +2866,8 @@ async def session_license_handler( response: Dict[str, Any] = {"keys": all_keys} if actual_drm_type: response["drm_type"] = actual_drm_type + if drm_types: + response["drm_types"] = drm_types return web.json_response(response) if not track_id: diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 5220429f..418d4d8a 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -22,14 +22,13 @@ from pywidevine.cdm import Cdm as WidevineCdm from pywidevine.pssh import PSSH from requests import Session -from unshackle.core.cdm.detect import is_playready_cdm from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine from unshackle.core.events import events from unshackle.core.session import RnetSession from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume -from unshackle.core.tracks.track import assert_fragments_decrypted +from unshackle.core.tracks.track import DRM_PREFERENCE_TYPES, assert_fragments_decrypted from unshackle.core.utilities import is_close_match, log_event, try_ensure_utf8 from unshackle.core.utils.redact import safe_display_url from unshackle.core.utils.xml import load_xml @@ -334,6 +333,17 @@ class DASH: else: track.drm = existing_drm + if track.drm and track.drm_preference: + wanted = DRM_PREFERENCE_TYPES[track.drm_preference] + preferred = [drm_obj for drm_obj in track.drm if isinstance(drm_obj, wanted)] + if preferred: + track.drm = preferred + else: + log.warning( + f"Track wants {track.drm_preference} DRM but this manifest does not offer it, " + "using the manifest's DRM instead" + ) + if pre_existing_keys and track.drm: for drm_obj in track.drm: if hasattr(drm_obj, "content_keys"): @@ -406,7 +416,7 @@ class DASH: track.data["dash"]["segment_durations"] = segment_durations if not track.drm and init_data and isinstance(track, (Video, Audio)): - prefers_playready = is_playready_cdm(cdm) + prefers_playready = track.prefers_playready(cdm) if prefers_playready: try: track.drm = [PlayReady.from_init_data(init_data)] diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index e2e24701..dcd3b8db 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -34,7 +34,7 @@ from unshackle.core.drm import DRM_T, ClearKey, MonaLisa, PlayReady, Widevine from unshackle.core.events import events from unshackle.core.session import RnetResponse, RnetSession from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume -from unshackle.core.tracks.track import assert_fragments_decrypted +from unshackle.core.tracks.track import DRM_PREFERENCE_TYPES, assert_fragments_decrypted from unshackle.core.utilities import get_extension, is_close_match, log_event, try_ensure_utf8 from unshackle.core.utils.redact import safe_display_url from unshackle.core.utils.subprocess import log_tool_run @@ -567,7 +567,7 @@ class HLS: initial_drm_key = None # Track the EXT-X-KEY used for initial licensing media_keys = [k for k in (master.keys or []) if k is not None] if media_keys: - cdm_media_keys = HLS.filter_keys_for_cdm(media_keys, cdm) + cdm_media_keys = HLS.filter_keys_for_cdm(media_keys, cdm, track.drm_preference) media_playlist_key = HLS.get_supported_key(cdm_media_keys) if cdm_media_keys else None if media_playlist_key: @@ -919,7 +919,7 @@ class HLS: segment_keys = getattr(segment, "keys", None) if segment_keys and segment not in unwanted_segments: if cdm: - cdm_segment_keys = HLS.filter_keys_for_cdm(segment_keys, cdm) + cdm_segment_keys = HLS.filter_keys_for_cdm(segment_keys, cdm, track.drm_preference) key = ( HLS.get_supported_key(cdm_segment_keys) if cdm_segment_keys @@ -1190,15 +1190,30 @@ class HLS: def filter_keys_for_cdm( keys: list[Union[m3u8.model.SessionKey, m3u8.model.Key]], cdm: object, + drm_preference: Optional[str] = None, ) -> list[Union[m3u8.model.SessionKey, m3u8.model.Key]]: """ Filter EXT-X-KEY entries to only include those matching the CDM type. This makes sure that we select the correct DRM system (Widevine or PlayReady) for the CDM in the configuration, so the challenge does not fail. + + A track's drm_preference wins over the CDM. If the playlist has no EXT-X-KEY entry + for it, the CDM decides as usual. """ playready_urn = f"urn:uuid:{PR_PSSH.SYSTEM_ID}" playready_keyformats = {playready_urn, "com.microsoft.playready"} + if drm_preference: + if DRM_PREFERENCE_TYPES[drm_preference] is PlayReady: + preferred = [k for k in keys if k.keyformat and k.keyformat.lower() in playready_keyformats] + else: + preferred = [k for k in keys if k.keyformat and k.keyformat.lower() == WidevineCdm.urn] + if preferred: + return preferred + logging.getLogger("HLS").warning( + f"Track wants {drm_preference} DRM but this playlist has no EXT-X-KEY entry for it, " + "using the CDM's DRM instead" + ) if is_widevine_cdm(cdm): return [k for k in keys if k.keyformat and k.keyformat.lower() == WidevineCdm.urn] elif is_playready_cdm(cdm): diff --git a/unshackle/core/remote_service.py b/unshackle/core/remote_service.py index 138b786f..78b3ca85 100644 --- a/unshackle/core/remote_service.py +++ b/unshackle/core/remote_service.py @@ -212,6 +212,12 @@ def build_tracks(data: Dict[str, Any]) -> Tracks: drm_objs = reconstruct_drm(track_data.get("drm")) if drm_objs: track_obj.drm = drm_objs + preference = track_data.get("drm_preference") + if preference: + try: + track_obj.drm_preference = preference + except ValueError: + log.warning(f"Ignoring unknown drm_preference {preference!r} from the server for track {track_obj.id}") tracks.attachments = [ Attachment(url=a["url"], name=a.get("name"), mime_type=a.get("mime_type"), description=a.get("description")) for a in data.get("attachments", []) @@ -710,6 +716,7 @@ class RemoteService: ) keys_by_track = resp.get("keys", {}) server_drm_type = resp.get("drm_type", drm_type) + drm_types_by_track = resp.get("drm_types", {}) self._server_cdm_type = server_drm_type self.log.debug(f"Server responded with drm_type={server_drm_type}, keys for {len(keys_by_track)} track(s)") @@ -719,7 +726,8 @@ class RemoteService: continue kid_list = list(track_keys.keys()) - drm_obj = self.create_drm_stub(server_drm_type, kid_list) + track_drm_type = drm_types_by_track.get(str(track.id), server_drm_type) + drm_obj = self.create_drm_stub(track_drm_type, kid_list) for kid_hex, key_hex in track_keys.items(): drm_obj.content_keys[UUID(hex=kid_hex)] = key_hex track.drm = [drm_obj] diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index fbb61c2b..ec765649 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -30,6 +30,13 @@ from unshackle.core.tracks import resume from unshackle.core.utilities import get_boxes, log_event, try_ensure_utf8 from unshackle.core.utils.subprocess import ffprobe +DRM_PREFERENCE_TYPES: dict[str, Union[type[Widevine], type[PlayReady]]] = { + "wv": Widevine, + "widevine": Widevine, + "pr": PlayReady, + "playready": PlayReady, +} + def direct_session(session: Union[Session, "RnetSession"]) -> Session: """requests.Session with copied headers/cookies and no proxy.""" @@ -320,6 +327,7 @@ class Track: self.needs_repack = needs_repack self.name = name self.drm = drm + self._drm_preference: Optional[str] = None self.edition: list[str] = [edition] if isinstance(edition, str) else (edition or []) self.session = session self.downloader = downloader @@ -401,6 +409,33 @@ class Track: value = defaultdict(dict, **value) self._data = value + @property + def drm_preference(self) -> Optional[str]: + """ + DRM system this track must license with, one of the names in DRM_PREFERENCE_TYPES. + + None (the default) lets the loaded CDM choose. Set it when the manifest advertises more + than one DRM system but only one of them licenses this track. + """ + return self._drm_preference + + @drm_preference.setter + def drm_preference(self, value: Optional[str]) -> None: + if value is None: + self._drm_preference = None + return + if not isinstance(value, str) or value.lower() not in DRM_PREFERENCE_TYPES: + raise ValueError( + f"Expected drm_preference to be one of {sorted(DRM_PREFERENCE_TYPES)} or None, not {value!r}" + ) + self._drm_preference = value.lower() + + def prefers_playready(self, cdm: Optional[object]) -> bool: + """Whether to try PlayReady before Widevine. The track's preference wins over the loaded CDM.""" + if self._drm_preference: + return DRM_PREFERENCE_TYPES[self._drm_preference] is PlayReady + return is_playready_cdm(cdm) + def download( self, session: Union[Session, "RnetSession"], @@ -493,7 +528,7 @@ class Track: elif self.descriptor == self.Descriptor.URL: try: if not self.drm and track_type in ("Video", "Audio"): - if is_playready_cdm(cdm): + if self.prefers_playready(cdm): try: self.drm = [PlayReady.from_track(self, session)] except PlayReady.Exceptions.PSSHNotFound: @@ -705,6 +740,12 @@ class Track: if not self.drm: return None + if self._drm_preference: + wanted = DRM_PREFERENCE_TYPES[self._drm_preference] + for drm in self.drm: + if isinstance(drm, wanted): + return drm + if is_widevine_cdm(cdm): for drm in self.drm: if isinstance(drm, Widevine):