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

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