From 6cfbaa7db166105c5caa65b8e1149b54731231c6 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Mon, 29 May 2023 22:23:39 +0100
Subject: [PATCH] Pass cookies to the aria2c and requests downloaders

For aria2c I've simplified the operation by offloading most of the work for creating a cookie header by just re-doing what Python-requests does. This results in the exact same cookies Python-requests would have used in a requests.get() call or such. It supports multiple of the same-name cookies under different domains/paths based on the URI of the mock request.
---
 devine/commands/dl.py               |  1 +
 devine/core/downloaders/aria2c.py   | 19 ++++++++++++++++++-
 devine/core/downloaders/requests.py |  7 +++++--
 devine/core/manifests/dash.py       |  7 +++++++
 devine/core/manifests/hls.py        |  1 +
 5 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/devine/commands/dl.py b/devine/commands/dl.py
index 74ebc42..a8c69cd 100644
--- a/devine/commands/dl.py
+++ b/devine/commands/dl.py
@@ -907,6 +907,7 @@ class dl:
                             uri=track.url,
                             out=save_path,
                             headers=service.session.headers,
+                            cookies=service.session.cookies,
                             proxy=proxy if track.needs_proxy else None,
                             progress=progress
                         )
diff --git a/devine/core/downloaders/aria2c.py b/devine/core/downloaders/aria2c.py
index fa1ac52..55f883e 100644
--- a/devine/core/downloaders/aria2c.py
+++ b/devine/core/downloaders/aria2c.py
@@ -2,9 +2,12 @@ import asyncio
 import subprocess
 import textwrap
 from functools import partial
+from http.cookiejar import CookieJar
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional, Union, MutableMapping
 
+import requests
+from requests.cookies import RequestsCookieJar, get_cookie_header, cookiejar_from_dict
 from rich.text import Text
 
 from devine.core.config import config
@@ -16,6 +19,7 @@ async def aria2c(
     uri: Union[str, list[str]],
     out: Path,
     headers: Optional[dict] = None,
+    cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None,
     proxy: Optional[str] = None,
     silent: bool = False,
     segmented: bool = False,
@@ -73,7 +77,20 @@ async def aria2c(
         "-i", "-"
     ]
 
+    if cookies:
+        # use python-requests pre-existing code to convert a Jar/Dict to a header while
+        # also supporting multiple cookies of the same name with different domain/paths.
+        if isinstance(cookies, CookieJar):
+            cookiejar = cookies
+        else:
+            cookiejar = cookiejar_from_dict(cookies)
+        mock_request = requests.Request(url=uri)
+        cookie_header = get_cookie_header(cookiejar, mock_request)
+        arguments.extend(["--header", f"Cookie: {cookie_header}"])
+
     for header, value in (headers or {}).items():
+        if header.lower() == "cookie":
+            raise ValueError("You cannot set Cookies as a header manually, please use the `cookies` param.")
         if header.lower() == "accept-encoding":
             # we cannot set an allowed encoding, or it will return compressed
             # and the code is not set up to uncompress the data
diff --git a/devine/core/downloaders/requests.py b/devine/core/downloaders/requests.py
index 1ddf599..f593905 100644
--- a/devine/core/downloaders/requests.py
+++ b/devine/core/downloaders/requests.py
@@ -1,17 +1,18 @@
 import time
 from functools import partial
 from pathlib import Path
-from typing import Optional, Union, Any
+from typing import Optional, Union, Any, MutableMapping
 
 from requests import Session
+from requests.cookies import RequestsCookieJar
 from rich import filesize
-from rich.filesize import decimal
 
 
 def requests(
     uri: Union[str, list[str]],
     out: Path,
     headers: Optional[dict] = None,
+    cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None,
     proxy: Optional[str] = None,
     progress: Optional[partial] = None,
     *_: Any,
@@ -45,6 +46,8 @@ def requests(
             if k.lower() != "accept-encoding"
         }
         session.headers.update(headers)
+    if cookies:
+        session.cookies.update(cookies)
     if proxy:
         session.proxies.update({"all": proxy})
 
diff --git a/devine/core/manifests/dash.py b/devine/core/manifests/dash.py
index 7fa9ca9..86d8692 100644
--- a/devine/core/manifests/dash.py
+++ b/devine/core/manifests/dash.py
@@ -23,6 +23,7 @@ from lxml.etree import Element
 from pywidevine.cdm import Cdm as WidevineCdm
 from pywidevine.pssh import PSSH
 from requests import Session
+from requests.cookies import RequestsCookieJar
 from rich import filesize
 
 from devine.core.constants import AnyTrack
@@ -425,6 +426,7 @@ class DASH:
                         track=track,
                         proxy=proxy,
                         headers=session.headers,
+                        cookies=session.cookies,
                         bytes_range=bytes_range,
                         stop_event=stop_event
                     )
@@ -491,6 +493,7 @@ class DASH:
         track: AnyTrack,
         proxy: Optional[str] = None,
         headers: Optional[MutableMapping[str, str | bytes]] = None,
+        cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None,
         bytes_range: Optional[str] = None,
         stop_event: Optional[Event] = None
     ) -> int:
@@ -504,6 +507,9 @@ class DASH:
                 fix an invalid value in the TFHD box of Audio Tracks.
             proxy: Proxy URI to use when downloading the Segment file.
             headers: HTTP Headers to send when requesting the Segment file.
+            cookies: Cookies to send when requesting the Segment file. The actual cookies sent
+                will be resolved based on the URI among other parameters. Multiple cookies with
+                the same name but a different domain/path are resolved.
             bytes_range: Download only specific bytes of the Segment file using the Range header.
             stop_event: Prematurely stop the Download from beginning. Useful if ran from
                 a Thread Pool. It will raise a KeyboardInterrupt if set.
@@ -527,6 +533,7 @@ class DASH:
                     uri=url,
                     out=out_path,
                     headers=headers_,
+                    cookies=cookies,
                     proxy=proxy,
                     silent=attempts != 5,
                     segmented=True
diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py
index 7d8bc99..299dffd 100644
--- a/devine/core/manifests/hls.py
+++ b/devine/core/manifests/hls.py
@@ -428,6 +428,7 @@ class HLS:
                     uri=urljoin(segment.base_uri, segment.uri),
                     out=out_path,
                     headers=headers_,
+                    cookies=session.cookies,
                     proxy=proxy,
                     silent=attempts != 5,
                     segmented=True