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

@@ -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"

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})

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}")

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:

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)]

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):

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]

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):