Get Cookie Header for each URL in aria2c

This commit is contained in:
rlaphoenix 2024-01-09 12:41:15 +00:00
parent 552a0f13a5
commit 9f04676b5c
1 changed files with 24 additions and 28 deletions

View File

@ -33,22 +33,11 @@ async def aria2c(
If multiple URLs are provided they will be downloaded in the provided order If multiple URLs are provided they will be downloaded in the provided order
to the output directory. They will not be merged together. to the output directory. They will not be merged together.
""" """
if isinstance(uri, list) and len(uri) == 1: if not isinstance(uri, list):
uri = uri[0] uri = [uri]
if isinstance(uri, list):
segmented = True if cookies and not isinstance(cookies, CookieJar):
uri = "\n".join([ cookies = cookiejar_from_dict(cookies)
f"{url}\n"
f"\tdir={out}\n"
f"\tout={i:08}.mp4"
for i, url in enumerate(uri)
])
if out.is_file():
raise ValueError("Provided multiple segments to download, expecting directory path")
elif "\t" not in uri:
uri = f"{uri}\n" \
f"\tdir={out.parent}\n" \
f"\tout={out.name}"
executable = get_binary_path("aria2c", "aria2") executable = get_binary_path("aria2c", "aria2")
if not executable: if not executable:
@ -60,6 +49,24 @@ async def aria2c(
async with start_pproxy(proxy) as pproxy_: async with start_pproxy(proxy) as pproxy_:
return await aria2c(uri, out, headers, cookies, pproxy_, silent, segmented, progress, *args) return await aria2c(uri, out, headers, cookies, pproxy_, silent, segmented, progress, *args)
multiple_urls = len(uri) > 1
url_files = []
for i, url in enumerate(uri):
url_text = url
if multiple_urls:
url_text += f"\n\tdir={out.parent}"
url_text += f"\n\tout={out.name}"
else:
url_text += f"\n\tdir={out}"
url_text += f"\n\tout={i:08}.mp4"
if cookies:
mock_request = requests.Request(url=url)
cookie_header = get_cookie_header(cookies, mock_request)
if cookie_header:
url_text += f"\n\theader=Cookie: {cookie_header}"
url_files.append(url_text)
url_file = "\n".join(url_files)
max_concurrent_downloads = int(config.aria2c.get("max_concurrent_downloads", 5)) max_concurrent_downloads = int(config.aria2c.get("max_concurrent_downloads", 5))
max_connection_per_server = int(config.aria2c.get("max_connection_per_server", 1)) max_connection_per_server = int(config.aria2c.get("max_connection_per_server", 1))
split = int(config.aria2c.get("split", 5)) split = int(config.aria2c.get("split", 5))
@ -92,17 +99,6 @@ async def aria2c(
*args *args
] ]
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": if header.lower() == "cookie":
raise ValueError("You cannot set Cookies as a header manually, please use the `cookies` param.") raise ValueError("You cannot set Cookies as a header manually, please use the `cookies` param.")
@ -126,7 +122,7 @@ async def aria2c(
stdout=[subprocess.PIPE, subprocess.DEVNULL][silent] stdout=[subprocess.PIPE, subprocess.DEVNULL][silent]
) )
p.stdin.write(uri.encode()) p.stdin.write(url_file.encode())
await p.stdin.drain() await p.stdin.drain()
p.stdin.close() p.stdin.close()