144 lines
6.4 KiB
Python
144 lines
6.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
import concurrent.futures
|
|
import logging
|
|
import os
|
|
import sys
|
|
|
|
from Crypto.Cipher import AES
|
|
|
|
from m3u8_to_mp4 import v2_abstract_task_processor
|
|
from m3u8_to_mp4.helpers import path_helper
|
|
from m3u8_to_mp4.helpers import printer_helper
|
|
from m3u8_to_mp4.networks.synchronous.sync_http_requester import request_for
|
|
|
|
|
|
def download_segment(segment_url, customized_http_header):
|
|
response_code, response_content = request_for(segment_url,
|
|
customized_http_header=customized_http_header)
|
|
|
|
return response_code, response_content
|
|
|
|
|
|
class MultiThreadsFileCrawler(v2_abstract_task_processor.AbstractFileCrawler):
|
|
|
|
def _fetch_segments_to_local_tmpdir(self, key_segments_pairs):
|
|
if len(key_segments_pairs) < 1:
|
|
return
|
|
|
|
progress_bar = printer_helper.ProcessBar(self.num_fetched_ts_segments,
|
|
self.num_fetched_ts_segments + len(
|
|
key_segments_pairs),
|
|
'segment set',
|
|
'downloading...',
|
|
'downloaded segments successfully!')
|
|
|
|
key_url_encrypted_data_triple = list()
|
|
with concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=self.num_concurrent) as executor:
|
|
while len(key_segments_pairs) > 0:
|
|
future_2_key_and_url = {
|
|
executor.submit(download_segment, segment_url,
|
|
self.customized_http_header): (
|
|
key, segment_url) for key, segment_url in key_segments_pairs}
|
|
|
|
response_code, response_data = None, None
|
|
for future in concurrent.futures.as_completed(
|
|
future_2_key_and_url):
|
|
key, segment_url = future_2_key_and_url[future]
|
|
try:
|
|
response_code, response_data = future.result()
|
|
except Exception as exc:
|
|
logging.exception(
|
|
'{} generated an exception: {}'.format(segment_url,
|
|
exc))
|
|
|
|
if response_code == 200:
|
|
key_url_encrypted_data_triple.append(
|
|
(key, segment_url, response_data))
|
|
|
|
key_segments_pairs.remove((key, segment_url))
|
|
progress_bar.update()
|
|
|
|
if len(key_segments_pairs) > 0:
|
|
sys.stdout.write('\n')
|
|
logging.info(
|
|
'{} segments are failed to download, retry...'.format(
|
|
len(key_segments_pairs)))
|
|
|
|
logging.info('decrypt and dump segments...')
|
|
for key, segment_url, encrypted_data in key_url_encrypted_data_triple:
|
|
file_name = path_helper.resolve_file_name_by_uri(segment_url)
|
|
file_path = os.path.join(self.tmpdir, file_name)
|
|
|
|
if key is not None:
|
|
crypt_ls = {"AES-128": AES}
|
|
crypt_obj = crypt_ls[key.method]
|
|
cryptor = crypt_obj.new(key.value.encode(),
|
|
crypt_obj.MODE_CBC)
|
|
encrypted_data = cryptor.decrypt(encrypted_data)
|
|
|
|
with open(file_path, 'wb') as fin:
|
|
fin.write(encrypted_data)
|
|
|
|
|
|
class MultiThreadsUriCrawler(v2_abstract_task_processor.AbstractUriCrawler):
|
|
|
|
def _fetch_segments_to_local_tmpdir(self, key_segments_pairs):
|
|
if len(key_segments_pairs) < 1:
|
|
return
|
|
|
|
progress_bar = printer_helper.ProcessBar(self.num_fetched_ts_segments,
|
|
self.num_fetched_ts_segments + len(
|
|
key_segments_pairs),
|
|
'segment set',
|
|
'downloading...',
|
|
'downloaded segments successfully!')
|
|
|
|
key_url_encrypted_data_triple = list()
|
|
with concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=self.num_concurrent) as executor:
|
|
while len(key_segments_pairs) > 0:
|
|
future_2_key_and_url = {
|
|
executor.submit(download_segment, segment_url,
|
|
self.customized_http_header): (
|
|
key, segment_url) for key, segment_url in key_segments_pairs}
|
|
|
|
response_code, response_data = None, None
|
|
for future in concurrent.futures.as_completed(
|
|
future_2_key_and_url):
|
|
key, segment_url = future_2_key_and_url[future]
|
|
try:
|
|
response_code, response_data = future.result()
|
|
except Exception as exc:
|
|
logging.exception(
|
|
'{} generated an exception: {}'.format(segment_url,
|
|
exc))
|
|
|
|
if response_code == 200:
|
|
key_url_encrypted_data_triple.append(
|
|
(key, segment_url, response_data))
|
|
|
|
key_segments_pairs.remove((key, segment_url))
|
|
progress_bar.update()
|
|
|
|
if len(key_segments_pairs) > 0:
|
|
sys.stdout.write('\n')
|
|
logging.info(
|
|
'{} segments are failed to download, retry...'.format(
|
|
len(key_segments_pairs)))
|
|
|
|
logging.info('decrypt and dump segments...')
|
|
for key, segment_url, encrypted_data in key_url_encrypted_data_triple:
|
|
file_name = path_helper.resolve_file_name_by_uri(segment_url)
|
|
file_path = os.path.join(self.tmpdir, file_name)
|
|
|
|
if key is not None:
|
|
crypt_ls = {"AES-128": AES}
|
|
crypt_obj = crypt_ls[key.method]
|
|
cryptor = crypt_obj.new(key.value.encode(),
|
|
crypt_obj.MODE_CBC)
|
|
encrypted_data = cryptor.decrypt(encrypted_data)
|
|
|
|
with open(file_path, 'wb') as fin:
|
|
fin.write(encrypted_data)
|