From 4006593a8aea0c061330b8cb36f9c7da41e65597 Mon Sep 17 00:00:00 2001 From: Andy Date: Fri, 12 Sep 2025 06:38:14 +0000 Subject: [PATCH] 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 | 161 +++++++++++++++++++------------ unshackle/core/manifests/m3u8.py | 59 ++--------- unshackle/core/tracks/track.py | 77 +++++++++++++++ 3 files changed, 183 insertions(+), 114 deletions(-) diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 0bac6e7..f6dc0a3 100644 --- a/unshackle/commands/dl.py +++ b/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, - SubtitleCodecChoice, VideoCodecChoice) +from unshackle.core.utils.click_types import ( + 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: - from_vaults = drm.content_keys.copy() + if kid not in drm.content_keys and cdm_only: + need_license = True - try: - if self.service == "NF": - 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 + if need_license and not vaults_only: + from_vaults = drm.content_keys.copy() - for kid_, key in drm.content_keys.items(): - if key == "0" * 32: - key = f"[red]{key}[/]" - label = f"[text2]{kid_.hex}:{key}{is_track_kid}" - if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children): - cek_tree.add(label) + try: + if self.service == "NF": + 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 - drm.content_keys = { - kid_: key for kid_, key in drm.content_keys.items() if key and key.count("0") != len(key) - } + for kid_, key in drm.content_keys.items(): + 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. - # So we re-add the keys from vaults earlier overwriting blanks or removed KIDs data. - drm.content_keys.update(from_vaults) + drm.content_keys = { + kid_: key for kid_, key in drm.content_keys.items() if key and key.count("0") != len(key) + } - successful_caches = self.vaults.add_keys(drm.content_keys) - 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" - ) - break # licensing twice will be unnecessary + # The CDM keys may have returned blank content keys for KIDs we got from vaults. + # So we re-add the keys from vaults earlier overwriting blanks or removed KIDs data. + drm.content_keys.update(from_vaults) + + successful_caches = self.vaults.add_keys(drm.content_keys) + 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: - from_vaults = drm.content_keys.copy() + if kid not in drm.content_keys and cdm_only: + need_license = True - try: - drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate) - 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 + if need_license and not vaults_only: + from_vaults = drm.content_keys.copy() - for kid_, key in drm.content_keys.items(): - label = f"[text2]{kid_.hex}:{key}{is_track_kid}" - if not any(f"{kid_.hex}:{key}" in x.label for x in cek_tree.children): - cek_tree.add(label) + try: + drm.get_content_keys(cdm=self.cdm, licence=licence, certificate=certificate) + 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 - 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) - 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" - ) - break + drm.content_keys.update(from_vaults) + + successful_caches = self.vaults.add_keys(drm.content_keys) + 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" diff --git a/unshackle/core/manifests/m3u8.py b/unshackle/core/manifests/m3u8.py index b11d971..16fad1d 100644 --- a/unshackle/core/manifests/m3u8.py +++ b/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 [])) - need_pr = not any(isinstance(d, PlayReady) for t in tracks for d in (t.drm or [])) + bool(master.session_keys or HLS.parse_session_data_keys(master, session or Session())) - if (need_wv or need_pr) and tracks.videos: - if not session: - 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 + if True: + for t in tracks.videos + tracks.audio: + t.needs_drm_loading = True + t.session = session return tracks diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 2a02a02..9ff4939 100644 --- a/unshackle/core/tracks/track.py +++ b/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,