feat(ClearKey): Pass session not proxy str in from_m3u_key method

This reduces the amount of connections being made by quite a bit for playlists that constantly change keys, or have new key data for every single segment (e.g., Pluto sometimes).

It also allows you to pass headers and cookies, while still also being able to supply a proxy.
This commit is contained in:
rlaphoenix 2024-03-08 15:44:41 +00:00
parent c974a41b6d
commit e0aa0e37d3
2 changed files with 30 additions and 12 deletions

View File

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

View File

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