def test_delete_files(mocker):
mock_unlink = mocker.patch('os.unlink')
from lib.utils import Utils
for files_list in [['0'], ['0', '1'], ['0', '1', '2'],
['0', '1', '2', '3']]:
Utils.delete_files(files_list)
assert mock_unlink.call_count == len(files_list)
mock_unlink.assert_has_calls([call(x) for x in files_list])
mock_unlink.reset_mock()
def process_video(self, video_metadata, mkv_filepath, root_url, pause_ev, resume_ev, progress_callback_func,
video_quality='highest'):
"""
Download video and decrypt, join, encode to mkv
:return:
"""
if video_metadata.get('fcid'):
ttid = video_metadata['fcid']
flipped = True
else:
ttid = video_metadata['ttid']
flipped = False
number_of_tracks = int(video_metadata['tapNToggle'])
duration = int(video_metadata['actualDuration'])
encryption_keys = dict()
self.logger.info("[{}]: Starting download for {}".format(ttid, mkv_filepath))
# download media files for this video.
m3u8_content = self._download_m3u8(root_url, ttid, flipped, video_quality)
if m3u8_content:
summary, tracks_info = M3u8Parser(m3u8_content, num_tracks=number_of_tracks).parse()
download_dir = os.path.join(self.temp_downloads_dir, str(ttid))
os.makedirs(download_dir, exist_ok=True)
temp_files_to_delete = set()
ts_files = list()
items_processed = 0
for track_index, track_info in enumerate(tracks_info):
streams_to_join = list()
for item in track_info:
# download encrypted stream..
enc_stream_filepath = '{}/{}'.format(download_dir, item['file_number'])
temp_files_to_delete.add(enc_stream_filepath)
download_flag = False
while not download_flag:
if pause_ev.is_set():
self.logger.info("[{}]: Pausing download for {}".format(ttid, mkv_filepath))
resume_ev.wait()
self.logger.info("[{}]: Resuming download for {}".format(ttid, mkv_filepath))
resume_ev.clear()
try:
with open(enc_stream_filepath, 'wb') as fh:
content = requests.get(item['url']).content
fh.write(content)
download_flag = True
except TimeoutError:
self.logger.warning("[{}]: Timeout error. retrying download for {}...".format(
ttid, item['url']))
time.sleep(self.conf.get('retry_wait'))
# decrypt files if encrypted.
if item.get('encryption_method') == "NONE":
streams_to_join.append(enc_stream_filepath)
else:
if not encryption_keys.get(item['encryption_key_id']):
key = self.session.get(item['encryption_key_url']).content[2:]
key = key[::-1] # reverse the bytes.
encryption_keys[item['encryption_key_id']] = key
encryption_key = encryption_keys[item['encryption_key_id']]
decrypted_stream_filepath = Decrypter.decrypt(
encryption_key, enc_stream_filepath,
download_dir)
streams_to_join.append(decrypted_stream_filepath)
temp_files_to_delete.add(decrypted_stream_filepath)
# update progress bar
items_processed += 1
items_processed_percent = items_processed * 100 // summary.get('media_files')
progress_callback_func(items_processed_percent)
# All stream files for this track are decrypted, join them.
self.logger.info("[{}]: joining streams for track {} ..".format(ttid, track_index))
ts_file = Encoder.join(streams_to_join, download_dir, track_index)
ts_files.append(ts_file)
temp_files_to_delete.add(ts_file)
# Encode all ts files into a single output mkv.
os.makedirs(os.path.dirname(mkv_filepath), exist_ok=True)
success = Encoder.encode_mkv(ttid, ts_files, mkv_filepath, duration, self.conf.get('debug'))
if success:
self.logger.info("[{}]: Processed {}\n---".format(ttid, mkv_filepath))
# delete temp files.
if not self.conf.get('debug'):
Utils.delete_files(list(temp_files_to_delete))
os.rmdir(download_dir)
