Fix: Implement lazy DRM loading for multi-track key retrieval

- Add deferred DRM loading to M3U8 parser to mark tracks for later processing
- Optimize prepare_drm to load DRM just-in-time during download process

unshackle/commands/dl.py

@@ -58,8 +58,15 @@ from unshackle.core.tracks.attachment import Attachment
 from unshackle.core.tracks.hybrid import Hybrid
 from unshackle.core.utilities import get_system_fonts, is_close_match, time_elapsed_since
 from unshackle.core.utils import tags
-from unshackle.core.utils.click_types import (LANGUAGE_RANGE, QUALITY_LIST, SEASON_RANGE, ContextData, MultipleChoice,
+from unshackle.core.utils.click_types import (
-                                              SubtitleCodecChoice, VideoCodecChoice)
+    LANGUAGE_RANGE,
+    QUALITY_LIST,
+    SEASON_RANGE,
+    ContextData,
+    MultipleChoice,
+    SubtitleCodecChoice,
+    VideoCodecChoice,
+)
 from unshackle.core.utils.collections import merge_dict
 from unshackle.core.utils.subprocess import ffprobe
 from unshackle.core.vaults import Vaults
@@ -862,6 +869,10 @@ class dl:
 
             selected_tracks, tracks_progress_callables = title.tracks.tree(add_progress=True)
 
+            for track in title.tracks:
+                if hasattr(track, "needs_drm_loading") and track.needs_drm_loading:
+                    track.load_drm_if_needed(service)
+
             download_table = Table.grid()
             download_table.add_row(selected_tracks)
 
@@ -1149,7 +1160,11 @@ class dl:
                             progress.start_task(task_id)  # TODO: Needed?
                             audio_expected = not video_only and not no_audio
                             muxed_path, return_code, errors = task_tracks.mux(
-                                str(title), progress=partial(progress.update, task_id=task_id), delete=False, audio_expected=audio_expected, title_language=title.language
+                                str(title),
+                                progress=partial(progress.update, task_id=task_id),
+                                delete=False,
+                                audio_expected=audio_expected,
+                                title_language=title.language,
                             )
                             muxed_paths.append(muxed_path)
                             if return_code >= 2:
@@ -1249,7 +1264,12 @@ class dl:
                 if pre_existing_tree:
                     cek_tree = pre_existing_tree
 
-                for kid in drm.kids:
+                need_license = False
+                all_kids = list(drm.kids)
+                if track_kid and track_kid not in all_kids:
+                    all_kids.append(track_kid)
+
+                for kid in all_kids:
                     if kid in drm.content_keys:
                         continue
 
@@ -1269,46 +1289,51 @@ class dl:
                             if not pre_existing_tree:
                                 table.add_row(cek_tree)
                             raise Widevine.Exceptions.CEKNotFound(msg)
+                        else:
+                            need_license = True
 
-                    if kid not in drm.content_keys and not vaults_only:
+                    if kid not in drm.content_keys and cdm_only:
-                        from_vaults = drm.content_keys.copy()
+                        need_license = True
 
-                        try:
+                if need_license and not vaults_only:
-                            if self.service == "NF":
+                    from_vaults = drm.content_keys.copy()
-                                drm.get_NF_content_keys(cdm=self.cdm, licence=licence, certificate=certificate)
-                            else:
-                                drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate)
-                        except Exception as e:
-                            if isinstance(e, (Widevine.Exceptions.EmptyLicense, Widevine.Exceptions.CEKNotFound)):
-                                msg = str(e)
-                            else:
-                                msg = f"An exception occurred in the Service's license function: {e}"
-                            cek_tree.add(f"[logging.level.error]{msg}")
-                            if not pre_existing_tree:
-                                table.add_row(cek_tree)
-                            raise e
 
-                        for kid_, key in drm.content_keys.items():
+                    try:
-                            if key == "0" * 32:
+                        if self.service == "NF":
-                                key = f"[red]{key}[/]"
+                            drm.get_NF_content_keys(cdm=self.cdm, licence=licence, certificate=certificate)
-                            label = f"[text2]{kid_.hex}:{key}{is_track_kid}"
+                        else:
-                            if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children):
+                            drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate)
-                                cek_tree.add(label)
+                    except Exception as e:
+                        if isinstance(e, (Widevine.Exceptions.EmptyLicense, Widevine.Exceptions.CEKNotFound)):
+                            msg = str(e)
+                        else:
+                            msg = f"An exception occurred in the Service's license function: {e}"
+                        cek_tree.add(f"[logging.level.error]{msg}")
+                        if not pre_existing_tree:
+                            table.add_row(cek_tree)
+                        raise e
 
-                        drm.content_keys = {
+                    for kid_, key in drm.content_keys.items():
-                            kid_: key for kid_, key in drm.content_keys.items() if key and key.count("0") != len(key)
+                        if key == "0" * 32:
-                        }
+                            key = f"[red]{key}[/]"
+                        is_track_kid_marker = ["", "*"][kid_ == track_kid]
+                        label = f"[text2]{kid_.hex}:{key}{is_track_kid_marker}"
+                        if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children):
+                            cek_tree.add(label)
 
-                        # The CDM keys may have returned blank content keys for KIDs we got from vaults.
+                    drm.content_keys = {
-                        # So we re-add the keys from vaults earlier overwriting blanks or removed KIDs data.
+                        kid_: key for kid_, key in drm.content_keys.items() if key and key.count("0") != len(key)
-                        drm.content_keys.update(from_vaults)
+                    }
 
-                        successful_caches = self.vaults.add_keys(drm.content_keys)
+                    # The CDM keys may have returned blank content keys for KIDs we got from vaults.
-                        self.log.info(
+                    # So we re-add the keys from vaults earlier overwriting blanks or removed KIDs data.
-                            f"Cached {len(drm.content_keys)} Key{'' if len(drm.content_keys) == 1 else 's'} to "
+                    drm.content_keys.update(from_vaults)
-                            f"{successful_caches}/{len(self.vaults)} Vaults"
+
-                        )
+                    successful_caches = self.vaults.add_keys(drm.content_keys)
-                        break  # licensing twice will be unnecessary
+                    self.log.info(
+                        f"Cached {len(drm.content_keys)} Key{'' if len(drm.content_keys) == 1 else 's'} to "
+                        f"{successful_caches}/{len(self.vaults)} Vaults"
+                    )
 
                 if track_kid and track_kid not in drm.content_keys:
                     msg = f"No Content Key for KID {track_kid.hex} was returned in the License"
@@ -1348,7 +1373,12 @@ class dl:
                 if pre_existing_tree:
                     cek_tree = pre_existing_tree
 
-                for kid in drm.kids:
+                need_license = False
+                all_kids = list(drm.kids)
+                if track_kid and track_kid not in all_kids:
+                    all_kids.append(track_kid)
+
+                for kid in all_kids:
                     if kid in drm.content_keys:
                         continue
 
@@ -1368,35 +1398,40 @@ class dl:
                             if not pre_existing_tree:
                                 table.add_row(cek_tree)
                             raise PlayReady.Exceptions.CEKNotFound(msg)
+                        else:
+                            need_license = True
 
-                    if kid not in drm.content_keys and not vaults_only:
+                    if kid not in drm.content_keys and cdm_only:
-                        from_vaults = drm.content_keys.copy()
+                        need_license = True
 
-                        try:
+                if need_license and not vaults_only:
-                            drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate)
+                    from_vaults = drm.content_keys.copy()
-                        except Exception as e:
-                            if isinstance(e, (PlayReady.Exceptions.EmptyLicense, PlayReady.Exceptions.CEKNotFound)):
-                                msg = str(e)
-                            else:
-                                msg = f"An exception occurred in the Service's license function: {e}"
-                            cek_tree.add(f"[logging.level.error]{msg}")
-                            if not pre_existing_tree:
-                                table.add_row(cek_tree)
-                            raise e
 
-                        for kid_, key in drm.content_keys.items():
+                    try:
-                            label = f"[text2]{kid_.hex}:{key}{is_track_kid}"
+                        drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate)
-                            if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children):
+                    except Exception as e:
-                                cek_tree.add(label)
+                        if isinstance(e, (PlayReady.Exceptions.EmptyLicense, PlayReady.Exceptions.CEKNotFound)):
+                            msg = str(e)
+                        else:
+                            msg = f"An exception occurred in the Service's license function: {e}"
+                        cek_tree.add(f"[logging.level.error]{msg}")
+                        if not pre_existing_tree:
+                            table.add_row(cek_tree)
+                        raise e
 
-                        drm.content_keys.update(from_vaults)
+                    for kid_, key in drm.content_keys.items():
+                        is_track_kid_marker = ["", "*"][kid_ == track_kid]
+                        label = f"[text2]{kid_.hex}:{key}{is_track_kid_marker}"
+                        if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children):
+                            cek_tree.add(label)
 
-                        successful_caches = self.vaults.add_keys(drm.content_keys)
+                    drm.content_keys.update(from_vaults)
-                        self.log.info(
+
-                            f"Cached {len(drm.content_keys)} Key{'' if len(drm.content_keys) == 1 else 's'} to "
+                    successful_caches = self.vaults.add_keys(drm.content_keys)
-                            f"{successful_caches}/{len(self.vaults)} Vaults"
+                    self.log.info(
-                        )
+                        f"Cached {len(drm.content_keys)} Key{'' if len(drm.content_keys) == 1 else 's'} to "
-                        break
+                        f"{successful_caches}/{len(self.vaults)} Vaults"
+                    )
 
                 if track_kid and track_kid not in drm.content_keys:
                     msg = f"No Content Key for KID {track_kid.hex} was returned in the License"

unshackle/core/manifests/m3u8.py

@@ -2,17 +2,11 @@
 
 from __future__ import annotations
 
-from typing import Optional, Union
+from typing import Optional
 
-import httpx
 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
 from requests import Session
 
-from unshackle.core.drm import PlayReady, Widevine
 from unshackle.core.manifests.hls import HLS
 from unshackle.core.tracks import Tracks
 
@@ -21,54 +15,17 @@ def parse(
     master: m3u8.M3U8,
     language: str,
     *,
-    session: Optional[Union[Session, httpx.Client]] = None,
+    session: Optional[Session] = None,
 ) -> Tracks:
-    """Parse a variant playlist to ``Tracks`` with DRM information."""
+    """Parse a variant playlist to ``Tracks`` with basic information, defer DRM loading."""
     tracks = HLS(master, session=session).to_tracks(language)
 
-    need_wv = not any(isinstance(d, Widevine) for t in tracks for d in (t.drm or []))
+    bool(master.session_keys or HLS.parse_session_data_keys(master, session or Session()))
-    need_pr = not any(isinstance(d, PlayReady) for t in tracks for d in (t.drm or []))
 
-    if (need_wv or need_pr) and tracks.videos:
+    if True:
-        if not session:
+        for t in tracks.videos + tracks.audio:
-            session = Session()
+            t.needs_drm_loading = True
-
+            t.session = session
-        session_keys = list(master.session_keys or [])
-        session_keys.extend(HLS.parse_session_data_keys(master, session))
-
-        for drm_obj in HLS.get_all_drm(session_keys):
-            if need_wv and isinstance(drm_obj, Widevine):
-                for t in tracks.videos + tracks.audio:
-                    t.drm = [d for d in (t.drm or []) if not isinstance(d, Widevine)] + [drm_obj]
-                need_wv = False
-            elif need_pr and isinstance(drm_obj, PlayReady):
-                for t in tracks.videos + tracks.audio:
-                    t.drm = [d for d in (t.drm or []) if not isinstance(d, PlayReady)] + [drm_obj]
-                need_pr = False
-            if not need_wv and not need_pr:
-                break
-
-    if (need_wv or need_pr) and tracks.videos:
-        first_video = tracks.videos[0]
-        playlist = m3u8.load(first_video.url)
-        for key in playlist.keys or []:
-            if not key or not key.keyformat:
-                continue
-            fmt = key.keyformat.lower()
-            if need_wv and fmt == WidevineCdm.urn:
-                pssh_b64 = key.uri.split(",")[-1]
-                drm = Widevine(pssh=WV_PSSH(pssh_b64))
-                for t in tracks.videos + tracks.audio:
-                    t.drm = [d for d in (t.drm or []) if not isinstance(d, Widevine)] + [drm]
-                need_wv = False
-            elif need_pr and (fmt == PlayReadyCdm or "com.microsoft.playready" in fmt):
-                pssh_b64 = key.uri.split(",")[-1]
-                drm = PlayReady(pssh=PR_PSSH(pssh_b64), pssh_b64=pssh_b64)
-                for t in tracks.videos + tracks.audio:
-                    t.drm = [d for d in (t.drm or []) if not isinstance(d, PlayReady)] + [drm]
-                need_pr = False
-            if not need_wv and not need_pr:
-                break
 
     return tracks
 

unshackle/core/tracks/track.py

@@ -473,6 +473,83 @@ class Track:
                 if tenc.key_ID.int != 0:
                     return tenc.key_ID
 
+    def load_drm_if_needed(self, service=None) -> bool:
+        """
+        Load DRM information for this track if it was deferred during parsing.
+
+        Args:
+            service: Service instance that can fetch track-specific DRM info
+
+        Returns:
+            True if DRM was loaded or already present, False if failed
+        """
+        if not getattr(self, "needs_drm_loading", False):
+            return bool(self.drm)
+
+        if self.drm:
+            self.needs_drm_loading = False
+            return True
+
+        if not service or not hasattr(service, "get_track_drm"):
+            return self.load_drm_from_playlist()
+
+        try:
+            track_drm = service.get_track_drm(self)
+            if track_drm:
+                self.drm = track_drm if isinstance(track_drm, list) else [track_drm]
+                self.needs_drm_loading = False
+                return True
+        except Exception as e:
+            raise ValueError(f"Failed to load DRM from service for track {self.id}: {e}")
+
+        return self.load_drm_from_playlist()
+
+    def load_drm_from_playlist(self) -> bool:
+        """
+        Fallback method to load DRM by fetching this track's individual playlist.
+        """
+        if self.drm:
+            self.needs_drm_loading = False
+            return True
+
+        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
+
+            session = getattr(self, "session", None) or Session()
+
+            response = session.get(self.url)
+            playlist = m3u8.loads(response.text, self.url)
+
+            drm_list = []
+
+            for key in playlist.keys or []:
+                if not key or not key.keyformat:
+                    continue
+
+                fmt = key.keyformat.lower()
+                if fmt == WidevineCdm.urn:
+                    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:
+                    pssh_b64 = key.uri.split(",")[-1]
+                    drm = PlayReady(pssh=PR_PSSH(pssh_b64), pssh_b64=pssh_b64)
+                    drm_list.append(drm)
+
+            if drm_list:
+                self.drm = drm_list
+                self.needs_drm_loading = False
+                return True
+
+        except Exception as e:
+            raise ValueError(f"Failed to load DRM from playlist for track {self.id}: {e}")
+
+        return False
+
     def get_init_segment(
         self,
         maximum_size: int = 20000,