From 1009f67c2a9a774bd4b3d7b09de4ad1268fa2f02 Mon Sep 17 00:00:00 2001 From: pukkandan Date: Wed, 22 Sep 2021 05:27:07 +0530 Subject: [PATCH] [fragment,aria2c] Generalize and refactor some code --- yt_dlp/downloader/external.py | 27 ++++----------------- yt_dlp/downloader/fragment.py | 45 +++++++++++++++++++---------------- 2 files changed, 30 insertions(+), 42 deletions(-) diff --git a/yt_dlp/downloader/external.py b/yt_dlp/downloader/external.py index a0d346c12..025eb38cb 100644 --- a/yt_dlp/downloader/external.py +++ b/yt_dlp/downloader/external.py @@ -6,7 +6,7 @@ import subprocess 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 @@ from ..utils import ( ) -class ExternalFD(FileDownloader): +class ExternalFD(FragmentFD): SUPPORTED_PROTOCOLS = ('http', 'https', 'ftp', 'ftps') can_download_to_stdout = False @@ -142,6 +142,7 @@ class ExternalFD(FileDownloader): 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 @@ class ExternalFD(FileDownloader): 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 @@ class ExternalFD(FileDownloader): 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') ) diff --git a/yt_dlp/downloader/fragment.py b/yt_dlp/downloader/fragment.py index 10ab90ba6..ebdef27db 100644 --- a/yt_dlp/downloader/fragment.py +++ b/yt_dlp/downloader/fragment.py @@ -324,6 +324,29 @@ class FragmentFD(FileDownloader): '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 @@ class FragmentFD(FileDownloader): 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 @@ class FragmentFD(FileDownloader): 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: