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.
This commit is contained in:
rlaphoenix 2023-05-29 22:23:39 +01:00
parent 1ff4858ca7
commit 6cfbaa7db1
5 changed files with 32 additions and 3 deletions

View File

@ -907,6 +907,7 @@ class dl:
uri=track.url, uri=track.url,
out=save_path, out=save_path,
headers=service.session.headers, headers=service.session.headers,
cookies=service.session.cookies,
proxy=proxy if track.needs_proxy else None, proxy=proxy if track.needs_proxy else None,
progress=progress progress=progress
) )

View File

@ -2,9 +2,12 @@ import asyncio
import subprocess import subprocess
import textwrap import textwrap
from functools import partial from functools import partial
from http.cookiejar import CookieJar
from pathlib import Path 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 rich.text import Text
from devine.core.config import config from devine.core.config import config
@ -16,6 +19,7 @@ async def aria2c(
uri: Union[str, list[str]], uri: Union[str, list[str]],
out: Path, out: Path,
headers: Optional[dict] = None, headers: Optional[dict] = None,
cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None,
proxy: Optional[str] = None, proxy: Optional[str] = None,
silent: bool = False, silent: bool = False,
segmented: bool = False, segmented: bool = False,
@ -73,7 +77,20 @@ async def aria2c(
"-i", "-" "-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(): 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": if header.lower() == "accept-encoding":
# we cannot set an allowed encoding, or it will return compressed # we cannot set an allowed encoding, or it will return compressed
# and the code is not set up to uncompress the data # and the code is not set up to uncompress the data

View File

@ -1,17 +1,18 @@
import time import time
from functools import partial from functools import partial
from pathlib import Path from pathlib import Path
from typing import Optional, Union, Any from typing import Optional, Union, Any, MutableMapping
from requests import Session from requests import Session
from requests.cookies import RequestsCookieJar
from rich import filesize from rich import filesize
from rich.filesize import decimal
def requests( def requests(
uri: Union[str, list[str]], uri: Union[str, list[str]],
out: Path, out: Path,
headers: Optional[dict] = None, headers: Optional[dict] = None,
cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None,
proxy: Optional[str] = None, proxy: Optional[str] = None,
progress: Optional[partial] = None, progress: Optional[partial] = None,
*_: Any, *_: Any,
@ -45,6 +46,8 @@ def requests(
if k.lower() != "accept-encoding" if k.lower() != "accept-encoding"
} }
session.headers.update(headers) session.headers.update(headers)
if cookies:
session.cookies.update(cookies)
if proxy: if proxy:
session.proxies.update({"all": proxy}) session.proxies.update({"all": proxy})

View File

@ -23,6 +23,7 @@ from lxml.etree import Element
from pywidevine.cdm import Cdm as WidevineCdm from pywidevine.cdm import Cdm as WidevineCdm
from pywidevine.pssh import PSSH from pywidevine.pssh import PSSH
from requests import Session from requests import Session
from requests.cookies import RequestsCookieJar
from rich import filesize from rich import filesize
from devine.core.constants import AnyTrack from devine.core.constants import AnyTrack
@ -425,6 +426,7 @@ class DASH:
track=track, track=track,
proxy=proxy, proxy=proxy,
headers=session.headers, headers=session.headers,
cookies=session.cookies,
bytes_range=bytes_range, bytes_range=bytes_range,
stop_event=stop_event stop_event=stop_event
) )
@ -491,6 +493,7 @@ class DASH:
track: AnyTrack, track: AnyTrack,
proxy: Optional[str] = None, proxy: Optional[str] = None,
headers: Optional[MutableMapping[str, str | bytes]] = None, headers: Optional[MutableMapping[str, str | bytes]] = None,
cookies: Optional[Union[MutableMapping[str, str], RequestsCookieJar]] = None,
bytes_range: Optional[str] = None, bytes_range: Optional[str] = None,
stop_event: Optional[Event] = None stop_event: Optional[Event] = None
) -> int: ) -> int:
@ -504,6 +507,9 @@ class DASH:
fix an invalid value in the TFHD box of Audio Tracks. fix an invalid value in the TFHD box of Audio Tracks.
proxy: Proxy URI to use when downloading the Segment file. proxy: Proxy URI to use when downloading the Segment file.
headers: HTTP Headers to send when requesting 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. 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 stop_event: Prematurely stop the Download from beginning. Useful if ran from
a Thread Pool. It will raise a KeyboardInterrupt if set. a Thread Pool. It will raise a KeyboardInterrupt if set.
@ -527,6 +533,7 @@ class DASH:
uri=url, uri=url,
out=out_path, out=out_path,
headers=headers_, headers=headers_,
cookies=cookies,
proxy=proxy, proxy=proxy,
silent=attempts != 5, silent=attempts != 5,
segmented=True segmented=True

View File

@ -428,6 +428,7 @@ class HLS:
uri=urljoin(segment.base_uri, segment.uri), uri=urljoin(segment.base_uri, segment.uri),
out=out_path, out=out_path,
headers=headers_, headers=headers_,
cookies=session.cookies,
proxy=proxy, proxy=proxy,
silent=attempts != 5, silent=attempts != 5,
segmented=True segmented=True