diff --git a/devine/core/drm/clearkey.py b/devine/core/drm/clearkey.py index a4cd269..b5bf743 100644 --- a/devine/core/drm/clearkey.py +++ b/devine/core/drm/clearkey.py @@ -6,10 +6,10 @@ from pathlib import Path from typing import Optional, Union from urllib.parse import urljoin -import requests from Cryptodome.Cipher import AES from Cryptodome.Util.Padding import pad, unpad from m3u8.model import Key +from requests import Session class ClearKey: @@ -58,14 +58,29 @@ class ClearKey: shutil.move(decrypted_path, path) @classmethod - def from_m3u_key(cls, m3u_key: Key, proxy: Optional[str] = None) -> ClearKey: + def from_m3u_key(cls, m3u_key: Key, session: Optional[Session] = None) -> ClearKey: + """ + Load a ClearKey from an M3U(8) Playlist's EXT-X-KEY. + + Parameters: + m3u_key: A Key object parsed from a m3u(8) playlist using + the `m3u8` library. + session: Optional session used to request external URIs with. + Useful to set headers, proxies, cookies, and so forth. + """ if not isinstance(m3u_key, Key): raise ValueError(f"Provided M3U Key is in an unexpected type {m3u_key!r}") + if not isinstance(session, (Session, type(None))): + raise TypeError(f"Expected session to be a {Session}, not a {type(session)}") + if not m3u_key.method.startswith("AES"): raise ValueError(f"Provided M3U Key is not an AES Clear Key, {m3u_key.method}") if not m3u_key.uri: raise ValueError("No URI in M3U Key, unable to get Key.") + if not session: + session = Session() + if m3u_key.uri.startswith("data:"): media_types, data = m3u_key.uri[5:].split(",") media_types = media_types.split(";") @@ -74,12 +89,11 @@ class ClearKey: key = data else: url = urljoin(m3u_key.base_uri, m3u_key.uri) - res = requests.get( + res = session.get( url=url, headers={ "User-Agent": "smartexoplayer/1.1.0 (Linux;Android 8.0.0) ExoPlayerLib/2.13.3" - }, - proxies={"all": proxy} if proxy else None + } ) res.raise_for_status() if not res.content: diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py index 7073e76..67305ac 100644 --- a/devine/core/manifests/hls.py +++ b/devine/core/manifests/hls.py @@ -469,7 +469,7 @@ class HLS: if key is None: encryption_data = None elif not encryption_data or encryption_data[0] != key: - drm = HLS.get_drm(key, proxy) + drm = HLS.get_drm(key, session) if isinstance(drm, Widevine): try: if map_data: @@ -586,8 +586,7 @@ class HLS: # break elif key.method == "AES-128": return key - # # TODO: Use a session instead of creating a new connection within - # encryption_data = (key, ClearKey.from_m3u_key(key, proxy)) + # encryption_data = (key, ClearKey.from_m3u_key(key, session)) # break elif key.method == "ISO-23001-7": return key @@ -613,21 +612,26 @@ class HLS: @staticmethod def get_drm( key: Union[m3u8.model.SessionKey, m3u8.model.Key], - proxy: Optional[str] = None + session: Optional[requests.Session] = None ) -> DRM_T: """ Convert HLS EXT-X-KEY data to an initialized DRM object. Parameters: key: m3u8 key system (EXT-X-KEY) object. - proxy: Optional proxy string used for requesting AES-128 URIs. + session: Optional session used to request AES-128 URIs. + Useful to set headers, proxies, cookies, and so forth. Raises a NotImplementedError if the key system is not supported. """ + if not isinstance(session, (Session, type(None))): + raise TypeError(f"Expected session to be a {Session}, not {type(session)}") + if not session: + session = Session() + # TODO: Add support for 'SAMPLE-AES', 'AES-CTR', 'AES-CBC', 'ClearKey' if key.method == "AES-128": - # TODO: Use a session instead of creating a new connection within - drm = ClearKey.from_m3u_key(key, proxy) + drm = ClearKey.from_m3u_key(key, session) elif key.method == "ISO-23001-7": drm = Widevine( pssh=PSSH.new(