zhaarey/unshackle
1
0
Fork 0
mirror of https://github.com/unshackle-dl/unshackle.git synced 2025-10-23 15:11:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
This commit is contained in:
Andy
2025-09-12 06:38:14 +00:00
parent 307be4549b
commit 4006593a8a
3 changed files with 183 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
unshackle/commands/dl.py
View File

@@ -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,8 +1289,13 @@ 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:
need_license = True
if need_license and not vaults_only:
from_vaults = drm.content_keys.copy()
try:
@@ -1291,7 +1316,8 @@ class dl:
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}"
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)
@@ -1308,7 +1334,6 @@ class dl:
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
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,8 +1398,13 @@ 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:
need_license = True
if need_license and not vaults_only:
from_vaults = drm.content_keys.copy()
try:
@@ -1385,7 +1420,8 @@ class dl:
raise e
for kid_, key in drm.content_keys.items():
label = f"[text2]{kid_.hex}:{key}{is_track_kid}"
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)
@@ -1396,7 +1432,6 @@ class dl:
f"Cached {len(drm.content_keys)} Key{'' if len(drm.content_keys) == 1 else 's'} to "
f"{successful_caches}/{len(self.vaults)} Vaults"
)
break
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"

57
unshackle/core/manifests/m3u8.py
View File

@@ -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):
if True:
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
t.needs_drm_loading = True
t.session = session
return tracks

77
unshackle/core/tracks/track.py
View File

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