mirror of
https://github.com/yt-dlp/yt-dlp.git
synced 2024-08-25 00:31:10 +02:00
[fragment,aria2c] Generalize and refactor some code
This commit is contained in:
parent
bd6f722de8
commit
1009f67c2a
@ -6,7 +6,7 @@
|
||||
import sys
|
||||
import time
|
||||
|
||||
from .common import FileDownloader
|
||||
from .fragment import FragmentFD
|
||||
from ..aes import aes_cbc_decrypt_bytes
|
||||
from ..compat import (
|
||||
compat_setenv,
|
||||
@ -30,7 +30,7 @@
|
||||
)
|
||||
|
||||
|
||||
class ExternalFD(FileDownloader):
|
||||
class ExternalFD(FragmentFD):
|
||||
SUPPORTED_PROTOCOLS = ('http', 'https', 'ftp', 'ftps')
|
||||
can_download_to_stdout = False
|
||||
|
||||
@ -142,6 +142,7 @@ def _call_downloader(self, tmpfilename, info_dict):
|
||||
self.report_error('Giving up after %s fragment retries' % fragment_retries)
|
||||
return -1
|
||||
|
||||
decrypt_fragment = self.decrypter(info_dict)
|
||||
dest, _ = sanitize_open(tmpfilename, 'wb')
|
||||
for frag_index, fragment in enumerate(info_dict['fragments']):
|
||||
fragment_filename = '%s-Frag%d' % (tmpfilename, frag_index)
|
||||
@ -153,21 +154,7 @@ def _call_downloader(self, tmpfilename, info_dict):
|
||||
continue
|
||||
self.report_error('Unable to open fragment %d' % frag_index)
|
||||
return -1
|
||||
decrypt_info = fragment.get('decrypt_info')
|
||||
if decrypt_info:
|
||||
if decrypt_info['METHOD'] == 'AES-128':
|
||||
iv = decrypt_info.get('IV') or compat_struct_pack('>8xq', fragment['media_sequence'])
|
||||
decrypt_info['KEY'] = decrypt_info.get('KEY') or self.ydl.urlopen(
|
||||
self._prepare_url(info_dict, info_dict.get('_decryption_key_url') or decrypt_info['URI'])).read()
|
||||
encrypted_data = src.read()
|
||||
decrypted_data = aes_cbc_decrypt_bytes(encrypted_data, decrypt_info['KEY'], iv)
|
||||
dest.write(decrypted_data)
|
||||
else:
|
||||
fragment_data = src.read()
|
||||
dest.write(fragment_data)
|
||||
else:
|
||||
fragment_data = src.read()
|
||||
dest.write(fragment_data)
|
||||
dest.write(decrypt_fragment(fragment, src.read()))
|
||||
src.close()
|
||||
if not self.params.get('keep_fragments', False):
|
||||
os.remove(encodeFilename(fragment_filename))
|
||||
@ -181,10 +168,6 @@ def _call_downloader(self, tmpfilename, info_dict):
|
||||
self.to_stderr(stderr.decode('utf-8', 'replace'))
|
||||
return p.returncode
|
||||
|
||||
def _prepare_url(self, info_dict, url):
|
||||
headers = info_dict.get('http_headers')
|
||||
return sanitized_Request(url, None, headers) if headers else url
|
||||
|
||||
|
||||
class CurlFD(ExternalFD):
|
||||
AVAILABLE_OPT = '-V'
|
||||
@ -518,7 +501,7 @@ class AVconvFD(FFmpegFD):
|
||||
_BY_NAME = dict(
|
||||
(klass.get_basename(), klass)
|
||||
for name, klass in globals().items()
|
||||
if name.endswith('FD') and name != 'ExternalFD'
|
||||
if name.endswith('FD') and name not in ('ExternalFD', 'FragmentFD')
|
||||
)
|
||||
|
||||
|
||||
|
@ -324,6 +324,29 @@ def _prepare_external_frag_download(self, ctx):
|
||||
'fragment_index': 0,
|
||||
})
|
||||
|
||||
def decrypter(self, info_dict):
|
||||
_key_cache = {}
|
||||
|
||||
def _get_key(url):
|
||||
if url not in _key_cache:
|
||||
_key_cache[url] = self.ydl.urlopen(self._prepare_url(info_dict, url)).read()
|
||||
return _key_cache[url]
|
||||
|
||||
def decrypt_fragment(fragment, frag_content):
|
||||
decrypt_info = fragment.get('decrypt_info')
|
||||
if not decrypt_info or decrypt_info['METHOD'] != 'AES-128':
|
||||
return frag_content
|
||||
iv = decrypt_info.get('IV') or compat_struct_pack('>8xq', fragment['media_sequence'])
|
||||
decrypt_info['KEY'] = decrypt_info.get('KEY') or _get_key(info_dict.get('_decryption_key_url') or decrypt_info['URI'])
|
||||
# Don't decrypt the content in tests since the data is explicitly truncated and it's not to a valid block
|
||||
# size (see https://github.com/ytdl-org/youtube-dl/pull/27660). Tests only care that the correct data downloaded,
|
||||
# not what it decrypts to.
|
||||
if self.params.get('test', False):
|
||||
return frag_content
|
||||
return aes_cbc_decrypt_bytes(frag_content, decrypt_info['KEY'], iv)
|
||||
|
||||
return decrypt_fragment
|
||||
|
||||
def download_and_append_fragments(self, ctx, fragments, info_dict, *, pack_func=None, finish_func=None):
|
||||
fragment_retries = self.params.get('fragment_retries', 0)
|
||||
is_fatal = (lambda idx: idx == 0) if self.params.get('skip_unavailable_fragments', True) else (lambda _: True)
|
||||
@ -369,26 +392,6 @@ def download_fragment(fragment, ctx):
|
||||
return False, frag_index
|
||||
return frag_content, frag_index
|
||||
|
||||
_key_cache = {}
|
||||
|
||||
def _get_key(url):
|
||||
if url not in _key_cache:
|
||||
_key_cache[url] = self.ydl.urlopen(self._prepare_url(info_dict, url)).read()
|
||||
return _key_cache[url]
|
||||
|
||||
def decrypt_fragment(fragment, frag_content):
|
||||
decrypt_info = fragment.get('decrypt_info')
|
||||
if not decrypt_info or decrypt_info['METHOD'] != 'AES-128':
|
||||
return frag_content
|
||||
iv = decrypt_info.get('IV') or compat_struct_pack('>8xq', fragment['media_sequence'])
|
||||
decrypt_info['KEY'] = decrypt_info.get('KEY') or _get_key(info_dict.get('_decryption_key_url') or decrypt_info['URI'])
|
||||
# Don't decrypt the content in tests since the data is explicitly truncated and it's not to a valid block
|
||||
# size (see https://github.com/ytdl-org/youtube-dl/pull/27660). Tests only care that the correct data downloaded,
|
||||
# not what it decrypts to.
|
||||
if self.params.get('test', False):
|
||||
return frag_content
|
||||
return aes_cbc_decrypt_bytes(frag_content, decrypt_info['KEY'], iv)
|
||||
|
||||
def append_fragment(frag_content, frag_index, ctx):
|
||||
if not frag_content:
|
||||
if not is_fatal(frag_index - 1):
|
||||
@ -402,6 +405,8 @@ def append_fragment(frag_content, frag_index, ctx):
|
||||
self._append_fragment(ctx, pack_func(frag_content, frag_index))
|
||||
return True
|
||||
|
||||
decrypt_fragment = self.decrypter(info_dict)
|
||||
|
||||
max_workers = self.params.get('concurrent_fragment_downloads', 1)
|
||||
if can_threaded_download and max_workers > 1:
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user