From e0aa0e37d330834d412c4b47d85ba8e487ef5d9c Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Fri, 8 Mar 2024 15:44:41 +0000
Subject: [PATCH] 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.
---
 devine/core/drm/clearkey.py  | 24 +++++++++++++++++++-----
 devine/core/manifests/hls.py | 18 +++++++++++-------
 2 files changed, 30 insertions(+), 12 deletions(-)

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(