fix(hls): prefer media playlist keys over session keys for accurate KID matching

Session keys from master playlists often contain PSSHs with multiple KIDs covering all tracks, causing licensing to return keys for wrong KIDs.

Changes:
- Unified DRM licensing logic for all downloaders
- Prefer media playlist EXT-X-KEY tags which contain track-specific KIDs
- Add filter_keys_for_cdm() to select keys matching configured CDM type
- Add get_track_kid_from_init() to extract KID from init segment with fallback to drm.kid from PSSH
- Track initial_drm_key to prevent double-licensing on first segment
- Simplify n_m3u8dl_re block to reuse common licensing flow
- Use strict PlayReady keyformat matching via PR_PSSH.SYSTEM_ID URN instead of loose substring match
- Fix PlayReady keyformat comparisons that incorrectly compared strings to PlayReadyCdm class
- Fix byterange header format in get_track_kid_from_init() to use HLS.calculate_byte_range()

Also fixes PlayReady keyformat matching in:
- unshackle/core/tracks/track.py
- unshackle/core/drm/playready.py

Fixes download failures where track_kid was null or mismatched, causing wrong content keys to be obtained during PlayReady/Widevine licensing.

unshackle/core/drm/playready.py

@@ -154,7 +154,9 @@ class PlayReady:
             pssh_boxes.extend(
                 Box.parse(base64.b64decode(x.uri.split(",")[-1]))
                 for x in (master.session_keys or master.keys)
-                if x and x.keyformat and "playready" in x.keyformat.lower()
+                if x and x.keyformat and x.keyformat.lower() in {
+                    f"urn:uuid:{PSSH.SYSTEM_ID}", "com.microsoft.playready"
+                }
             )
 
         init_data = track.get_init_segment(session=session)

unshackle/core/manifests/hls.py

@@ -12,6 +12,7 @@ from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Optional, Union
 from urllib.parse import urljoin
+from uuid import UUID
 from zlib import crc32
 
 import m3u8
@@ -260,7 +261,9 @@ class HLS:
                     sys.exit(1)
                 playlist_text = response.text
             else:
-                raise TypeError(f"Expected response to be a requests.Response or curl_cffi.Response, not {type(response)}")
+                raise TypeError(
+                    f"Expected response to be a requests.Response or curl_cffi.Response, not {type(response)}"
+                )
 
             master = m3u8.loads(playlist_text, uri=track.url)
 
@@ -268,23 +271,51 @@ class HLS:
             log.error("Track's HLS playlist has no segments, expecting an invariant M3U8 playlist.")
             sys.exit(1)
 
+        # Get session DRM as fallback but prefer media playlist keys for accurate KID matching
         if track.drm:
             session_drm = track.get_drm_for_cdm(cdm)
-            if isinstance(session_drm, (Widevine, PlayReady)):
-                # license and grab content keys
-                try:
-                    if not license_widevine:
-                        raise ValueError("license_widevine func must be supplied to use DRM")
-                    progress(downloaded="LICENSING")
-                    license_widevine(session_drm)
-                    progress(downloaded="[yellow]LICENSED")
-                except Exception:  # noqa
-                    DOWNLOAD_CANCELLED.set()  # skip pending track downloads
-                    progress(downloaded="[red]FAILED")
-                    raise
         else:
             session_drm = None
 
+        initial_drm_licensed = False
+        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)
+            media_playlist_key = HLS.get_supported_key(cdm_media_keys) if cdm_media_keys else None
+
+            if media_playlist_key:
+                media_drm = HLS.get_drm(media_playlist_key, session)
+                if isinstance(media_drm, (Widevine, PlayReady)):
+                    track_kid = HLS.get_track_kid_from_init(master, track, session) or media_drm.kid
+                    try:
+                        if not license_widevine:
+                            raise ValueError("license_widevine func must be supplied to use DRM")
+                        progress(downloaded="LICENSING")
+                        license_widevine(media_drm, track_kid=track_kid)
+                        progress(downloaded="[yellow]LICENSED")
+                        initial_drm_licensed = True
+                        initial_drm_key = media_playlist_key
+                        track.drm = [media_drm]
+                        session_drm = media_drm
+                    except Exception:  # noqa
+                        DOWNLOAD_CANCELLED.set()  # skip pending track downloads
+                        progress(downloaded="[red]FAILED")
+                        raise
+
+        # Fall back to session DRM if media playlist has no matching keys
+        if not initial_drm_licensed and session_drm and isinstance(session_drm, (Widevine, PlayReady)):
+            try:
+                if not license_widevine:
+                    raise ValueError("license_widevine func must be supplied to use DRM")
+                progress(downloaded="LICENSING")
+                license_widevine(session_drm)
+                progress(downloaded="[yellow]LICENSED")
+            except Exception:  # noqa
+                DOWNLOAD_CANCELLED.set()  # skip pending track downloads
+                progress(downloaded="[red]FAILED")
+                raise
+
         if DOWNLOAD_LICENCE_ONLY.is_set():
             progress(downloaded="[yellow]SKIPPED")
             return
@@ -341,12 +372,15 @@ class HLS:
 
         if downloader.__name__ == "n_m3u8dl_re":
             skip_merge = True
+            # session_drm already has correct content_keys from initial licensing above
+            n_m3u8dl_content_keys = session_drm.content_keys if session_drm else None
+
             downloader_args.update(
                 {
                     "output_dir": save_dir,
                     "filename": track.id,
                     "track": track,
-                    "content_keys": session_drm.content_keys if session_drm else None,
+                    "content_keys": n_m3u8dl_content_keys,
                 }
             )
 
@@ -390,7 +424,7 @@ class HLS:
             range_offset = 0
             map_data: Optional[tuple[m3u8.model.InitializationSection, bytes]] = None
             if session_drm:
-                encryption_data: Optional[tuple[Optional[m3u8.Key], DRM_T]] = (None, session_drm)
+                encryption_data: Optional[tuple[Optional[m3u8.Key], DRM_T]] = (initial_drm_key, session_drm)
             else:
                 encryption_data: Optional[tuple[Optional[m3u8.Key], DRM_T]] = None
 
@@ -571,6 +605,8 @@ class HLS:
                                     track_kid = track.get_key_id(map_data[1])
                                 else:
                                     track_kid = None
+                                if not track_kid:
+                                    track_kid = drm.kid
                                 progress(downloaded="LICENSING")
                                 license_widevine(drm, track_kid=track_kid)
                                 progress(downloaded="[yellow]LICENSED")
@@ -770,6 +806,60 @@ class HLS:
 
         return keys
 
+    @staticmethod
+    def filter_keys_for_cdm(
+        keys: list[Union[m3u8.model.SessionKey, m3u8.model.Key]],
+        cdm: object,
+    ) -> list[Union[m3u8.model.SessionKey, m3u8.model.Key]]:
+        """
+        Filter EXT-X-KEY entries to only include those matching the CDM type.
+
+        This ensures we select the correct DRM system (Widevine vs PlayReady)
+        based on what CDM is configured, avoiding license request failures.
+        """
+        playready_urn = f"urn:uuid:{PR_PSSH.SYSTEM_ID}"
+        playready_keyformats = {playready_urn, "com.microsoft.playready"}
+        if isinstance(cdm, WidevineCdm):
+            return [k for k in keys if k.keyformat and k.keyformat.lower() == WidevineCdm.urn]
+        elif isinstance(cdm, PlayReadyCdm):
+            return [k for k in keys if k.keyformat and k.keyformat.lower() in playready_keyformats]
+        elif hasattr(cdm, "is_playready"):
+            if cdm.is_playready:
+                return [k for k in keys if k.keyformat and k.keyformat.lower() in playready_keyformats]
+            else:
+                return [k for k in keys if k.keyformat and k.keyformat.lower() == WidevineCdm.urn]
+        return keys
+
+    @staticmethod
+    def get_track_kid_from_init(
+        master: M3U8,
+        track: AnyTrack,
+        session: Union[Session, CurlSession],
+    ) -> Optional[UUID]:
+        """
+        Extract the track's Key ID from its init segment (EXT-X-MAP).
+
+        Returns None if no init segment exists or KID extraction fails.
+        The caller should fall back to drm.kid from the PSSH if this returns None.
+        """
+        map_section = next((seg.init_section for seg in master.segments if seg.init_section), None)
+        if not map_section:
+            return None
+
+        map_uri = urljoin(map_section.base_uri or master.base_uri or "", map_section.uri)
+        try:
+            if map_section.byterange:
+                byte_range = HLS.calculate_byte_range(map_section.byterange, 0)
+                headers = {"Range": f"bytes={byte_range}"}
+            else:
+                headers = {}
+            map_res = session.get(url=map_uri, headers=headers)
+            if map_res.ok:
+                return track.get_key_id(map_res.content)
+        except Exception:
+            pass
+        return None
+
     @staticmethod
     def get_supported_key(keys: list[Union[m3u8.model.SessionKey, m3u8.model.Key]]) -> Optional[m3u8.Key]:
         """
@@ -798,9 +888,9 @@ class HLS:
                 return key
             elif key.keyformat and key.keyformat.lower() == WidevineCdm.urn:
                 return key
-            elif key.keyformat and (
+            elif key.keyformat and key.keyformat.lower() in {
-                key.keyformat.lower() == PlayReadyCdm or "com.microsoft.playready" in key.keyformat.lower()
+                f"urn:uuid:{PR_PSSH.SYSTEM_ID}", "com.microsoft.playready"
-            ):
+            }:
                 return key
             else:
                 unsupported_systems.append(key.method + (f" ({key.keyformat})" if key.keyformat else ""))
@@ -837,9 +927,9 @@ class HLS:
                 pssh=WV_PSSH(key.uri.split(",")[-1]),
                 **key._extra_params,  # noqa
             )
-        elif key.keyformat and (
+        elif key.keyformat and key.keyformat.lower() in {
-            key.keyformat.lower() == PlayReadyCdm or "com.microsoft.playready" in key.keyformat.lower()
+            f"urn:uuid:{PR_PSSH.SYSTEM_ID}", "com.microsoft.playready"
-        ):
+        }:
             drm = PlayReady(
                 pssh=PR_PSSH(key.uri.split(",")[-1]),
                 pssh_b64=key.uri.split(",")[-1],

unshackle/core/tracks/track.py

@@ -215,7 +215,8 @@ class Track:
             # or when the subtitle has a direct file extension
             if self.downloader.__name__ == "n_m3u8dl_re" and (
                 self.descriptor == self.Descriptor.URL
-                or get_extension(self.url) in {
+                or get_extension(self.url)
+                in {
                     ".srt",
                     ".vtt",
                     ".ttml",
@@ -303,7 +304,9 @@ class Track:
                                 try:
                                     self.drm = [Widevine.from_track(self, session)]
                                 except Widevine.Exceptions.PSSHNotFound:
-                                    log.warning("No PlayReady or Widevine PSSH was found for this track, is it DRM free?")
+                                    log.warning(
+                                        "No PlayReady or Widevine PSSH was found for this track, is it DRM free?"
+                                    )
                         else:
                             try:
                                 self.drm = [Widevine.from_track(self, session)]
@@ -311,7 +314,9 @@ class Track:
                                 try:
                                     self.drm = [PlayReady.from_track(self, session)]
                                 except PlayReady.Exceptions.PSSHNotFound:
-                                    log.warning("No Widevine or PlayReady PSSH was found for this track, is it DRM free?")
+                                    log.warning(
+                                        "No Widevine or PlayReady PSSH was found for this track, is it DRM free?"
+                                    )
 
                     if self.drm:
                         track_kid = self.get_key_id(session=session)
@@ -548,7 +553,6 @@ class Track:
 
         try:
             import m3u8
-            from pyplayready.cdm import Cdm as PlayReadyCdm
             from pyplayready.system.pssh import PSSH as PR_PSSH
             from pywidevine.cdm import Cdm as WidevineCdm
             from pywidevine.pssh import PSSH as WV_PSSH
@@ -569,7 +573,7 @@ class Track:
                     pssh_b64 = key.uri.split(",")[-1]
                     drm = Widevine(pssh=WV_PSSH(pssh_b64))
                     drm_list.append(drm)
-                elif fmt == PlayReadyCdm or "com.microsoft.playready" in fmt:
+                elif fmt in {f"urn:uuid:{PR_PSSH.SYSTEM_ID}", "com.microsoft.playready"}:
                     pssh_b64 = key.uri.split(",")[-1]
                     drm = PlayReady(pssh=PR_PSSH(pssh_b64), pssh_b64=pssh_b64)
                     drm_list.append(drm)