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)