diff --git a/main.py b/main.py index 3f10793..3783f79 100644 --- a/main.py +++ b/main.py @@ -18,7 +18,7 @@ from utils.input import ( ) from utils.oqee import OqeeClient from utils.downloader import get_keys -from utils.utilities import verify_cmd, merge_segments, decrypt, verify_mp4ff +from utils.utilities import verify_cmd, decrypt, verify_mp4ff from utils.times import ( convert_date_to_sec, convert_sec_to_ticks, @@ -190,36 +190,49 @@ if __name__ == "__main__": mpd_content = get_manifest(dash_id) manifest_info = parse_mpd_manifest(mpd_content) - avc_sel = selections_avc["video"] - avc_init_segment = avc_sel["segments"]["initialization"] - avc_track_id = avc_init_segment.split("/")[-1].split("_init")[0] + avc_tick = None + if "video" in selections_avc: + avc_sel = selections_avc["video"] + avc_init_segment = avc_sel["segments"]["initialization"] + avc_track_id = avc_init_segment.split("/")[-1].split("_init")[0] - logger.info("Bruteforcing AVC video track %s...", avc_track_id) - avc_valid_ticks = asyncio.run( - bruteforce(avc_track_id, start_tick_user, batch_size) - ) - if len(avc_valid_ticks) == 0: - logger.error("No valid ticks found for AVC video.") - sys.exit(1) - avc_tick = avc_valid_ticks[0] - logger.info("AVC video tick found: %s", avc_tick) - - hevc_sel = selections_hevc["video"] - hevc_init_segment = hevc_sel["segments"]["initialization"] - hevc_track_id = hevc_init_segment.split("/")[-1].split("_init")[0] - - logger.info("Bruteforcing HEVC video track %s...", hevc_track_id) - hevc_valid_ticks = asyncio.run( - bruteforce(hevc_track_id, start_tick_user, batch_size) - ) - if len(hevc_valid_ticks) == 0: - logger.warning( - "No valid ticks found for HEVC video. HEVC tracks will be removed from manifest." + logger.info("Bruteforcing AVC video track %s...", avc_track_id) + avc_valid_ticks = asyncio.run( + bruteforce(avc_track_id, start_tick_user, batch_size) ) - hevc_tick = None + if len(avc_valid_ticks) == 0: + logger.warning( + "No valid ticks found for AVC video. AVC tracks will be removed from manifest." + ) + else: + avc_tick = avc_valid_ticks[0] + logger.info("AVC video tick found: %s", avc_tick) else: - hevc_tick = hevc_valid_ticks[0] - logger.info("HEVC video tick found: %s", hevc_tick) + logger.warning( + "No AVC video track available. AVC tracks will be removed from manifest." + ) + + hevc_tick = None + if "video" in selections_hevc: + hevc_sel = selections_hevc["video"] + hevc_init_segment = hevc_sel["segments"]["initialization"] + hevc_track_id = hevc_init_segment.split("/")[-1].split("_init")[0] + + logger.info("Bruteforcing HEVC video track %s...", hevc_track_id) + hevc_valid_ticks = asyncio.run( + bruteforce(hevc_track_id, start_tick_user, batch_size) + ) + if len(hevc_valid_ticks) == 0: + logger.warning( + "No valid ticks found for HEVC video. HEVC tracks will be removed from manifest." + ) + else: + hevc_tick = hevc_valid_ticks[0] + logger.info("HEVC video tick found: %s", hevc_tick) + else: + logger.warning( + "No HEVC video track available. HEVC tracks will be removed from manifest." + ) # Bruteforce for audio (same for all audio tracks) audio_sel = selections_avc["audio"] @@ -248,13 +261,20 @@ if __name__ == "__main__": rep_codec = rep.get("codecs", "") or rep.get("codec", "") if rep.get("segments") and rep["segments"].get("timeline"): if rep_codec.startswith("avc"): - for seg in rep["segments"]["timeline"]: - seg["t"] = avc_tick - logger.debug( - "Updated AVC tick for video rep %s (codec: %s)", - rep.get("id"), - rep_codec, - ) + if avc_tick is not None: + for seg in rep["segments"]["timeline"]: + seg["t"] = avc_tick + logger.debug( + "Updated AVC tick for video rep %s (codec: %s)", + rep.get("id"), + rep_codec, + ) + else: + reps_to_remove.append(rep) + logger.debug( + "Marking AVC rep %s for removal (no valid tick)", + rep.get("id"), + ) elif rep_codec.startswith( "hvc" ) or rep_codec.startswith("hev"): @@ -273,7 +293,7 @@ if __name__ == "__main__": rep.get("id"), ) - # Remove HEVC representations + # Remove representations with no valid tick for rep in reps_to_remove: adaptation_info["representations"].remove(rep) @@ -494,17 +514,6 @@ if __name__ == "__main__": ) ) - # Merge video and audio - video_file = f"{output_dir}/temp_video.mp4" - audio_file = f"{output_dir}/temp_audio.mp4" - - data["file"] = video_file if content_type == "video" else audio_file - merge_segments( - output_dir, - track_id, - video_file if content_type == "video" else audio_file, - ) - kid = get_kid(output_dir, track_id) data["kid"] = kid key = None @@ -539,7 +548,6 @@ if __name__ == "__main__": for content_type, data in [("video", video_data), ("audio", audio_data)]: track_id = data["track_id"] - file = data["file"] kid = data["kid"] key = None @@ -550,7 +558,7 @@ if __name__ == "__main__": init_path = get_init(output_dir, track_id) dec_file = f"{output_dir}/dec_{content_type}.mp4" - decrypt(file, init_path, dec_file, key) + decrypt(f"{output_dir}/segments_{track_id}", init_path, dec_file, key) track_id_video = video_data["track_id"] track_id_audio = audio_data["track_id"] @@ -607,8 +615,7 @@ if __name__ == "__main__": os.remove(f"{output_dir}/dec_audio.mp4") os.remove(f"{output_dir}/video.mp4") os.remove(f"{output_dir}/audio.mp4") - os.remove(f"{output_dir}/temp_video.mp4") - os.remove(f"{output_dir}/temp_audio.mp4") + shutil.rmtree(f"{output_dir}/segments_{video_data['track_id']}") shutil.rmtree(f"{output_dir}/segments_{audio_data['track_id']}") diff --git a/pyproject.toml b/pyproject.toml index f04f5b7..264b420 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "oqee-rewind" -version = "0.1.0" +version = "1.0" description = "Oqee TV Live Downloader" readme = "README.md" requires-python = ">=3.9" diff --git a/utils/input.py b/utils/input.py index e4399cb..947ff76 100644 --- a/utils/input.py +++ b/utils/input.py @@ -439,8 +439,16 @@ def select_track(content_dict, quality_spec, content_type): candidates.extend(tracks) if not candidates: - logger.warning("No %s track found for '%s'.", content_type, quality_spec) - return None + if filter_part: + logger.warning( + "No %s track found for '%s'. Falling back to unfiltered selection.", + content_type, quality_spec + ) + for key, tracks in content_dict.items(): + candidates.extend(tracks) + if not candidates: + logger.warning("No %s track found for '%s'.", content_type, quality_spec) + return None if pref == "best": selected = max(candidates, key=lambda x: x["bandwidth"]) diff --git a/utils/stream.py b/utils/stream.py index fe56d14..06a1950 100644 --- a/utils/stream.py +++ b/utils/stream.py @@ -456,8 +456,19 @@ async def save_segments( # Build list of all segments to download segments_to_download = [(start_tick + i * duration, i) for i in range(rep_nb)] + + # In case of resuming, check which segments are already downloaded + already_downloaded = [ + int(f.split(".")[0]) + for f in os.listdir(f"{output_folder}/segments_{track_id}") + if f.endswith(".m4s") and f.split(".")[0].isdigit() + ] + segments_to_download = [ + (tick, rep) for tick, rep in segments_to_download if tick not in already_downloaded + ] + retry_list = [] - successful = 0 + successful = len(already_downloaded) async with aiohttp.ClientSession() as session: # Process segments in batches diff --git a/utils/times.py b/utils/times.py index 4d410c6..cdf1f15 100644 --- a/utils/times.py +++ b/utils/times.py @@ -3,6 +3,7 @@ import asyncio import datetime import time +from zoneinfo import ZoneInfo import aiohttp from tqdm import tqdm @@ -10,6 +11,8 @@ from tqdm import tqdm from utils.stream import fetch_segment from utils.logging_config import logger +FRANCE_TZ = ZoneInfo("Europe/Paris") + def convert_ticks_to_sec(ticks, timescale): """Convert ticks to seconds.""" @@ -21,24 +24,24 @@ def convert_sec_to_ticks(seconds, timescale): return seconds * timescale -def convert_sec_to_date(seconds, offset_hours=1): - """Convert seconds to datetime with offset.""" - dt = datetime.datetime.utcfromtimestamp(seconds) + datetime.timedelta( - hours=offset_hours +def convert_sec_to_date(seconds): + """Convert UTC seconds to France local datetime.""" + dt = datetime.datetime.fromtimestamp(seconds, tz=datetime.UTC).astimezone( + FRANCE_TZ ) - return dt + return dt.replace(tzinfo=None) -def convert_date_to_sec(dt, offset_hours=1): - """Convert datetime to seconds with offset.""" - epoch = datetime.datetime(1970, 1, 1) - utc_dt = dt - datetime.timedelta(hours=offset_hours) - return (utc_dt - epoch).total_seconds() +def convert_date_to_sec(dt): + """Convert France local datetime to UTC seconds.""" + aware = dt.replace(tzinfo=FRANCE_TZ) + epoch = datetime.datetime(1970, 1, 1, tzinfo=datetime.UTC) + return (aware.astimezone(datetime.UTC) - epoch).total_seconds() -def convert_date_to_ticks(dt, timescale, offset_hours=1): - """Convert datetime to ticks with offset.""" - return int(round(convert_date_to_sec(dt, offset_hours) * timescale)) +def convert_date_to_ticks(dt, timescale): + """Convert France local datetime to ticks.""" + return int(round(convert_date_to_sec(dt) * timescale)) def past(rep, base, duration): @@ -99,9 +102,9 @@ async def bruteforce(track_id, date, batch_size=20000): return valid_ticks -def find_nearest_tick_by_hour(base_tick, dt, timescale, duration, offset_hours=1): +def find_nearest_tick_by_hour(base_tick, dt, timescale, duration): """Find the nearest tick for a given datetime.""" - target_ticks = convert_date_to_ticks(dt, timescale, offset_hours) + target_ticks = convert_date_to_ticks(dt, timescale) diff_ticks = base_tick - target_ticks rep_estimate = diff_ticks / duration diff --git a/utils/utilities.py b/utils/utilities.py index 1645575..d93a623 100644 --- a/utils/utilities.py +++ b/utils/utilities.py @@ -43,11 +43,11 @@ def merge_segments(input_folder: str, track_id: str, output_file: str): logger.info("Merged segments into %s", output_file) -def decrypt(input_file, init_path, output_file, key): - """Decrypt a media file using mp4ff-decrypt. +def decrypt(segment_dir, init_path, output_file, key): + """Decrypt segments in chunks of ~1GB to avoid loading entire file in memory. Args: - input_file: Path to the input encrypted file. + segment_dir: Path to the directory containing .m4s segment files. init_path: Path to the initialization file. output_file: Path to the output decrypted file. key: The decryption key in KID:KEY format. @@ -56,14 +56,90 @@ def decrypt(input_file, init_path, output_file, key): True if decryption succeeded, False otherwise. """ key = key.split(":")[1] - result = subprocess.run( - ["mp4ff-decrypt", "-init", init_path, "-key", key, input_file, output_file], - capture_output=True, - text=True, - check=False, + + segment_files = sorted( + [f for f in os.listdir(segment_dir) if f.endswith(".m4s")], + key=lambda x: int(x.split(".")[0]), ) - if result.returncode == 0: - logger.info("Decrypted %s to %s", input_file, output_file) - return True - logger.error("Decryption failed: %s", result.stderr) - return False + + if not segment_files: + logger.error("No segment files found in %s", segment_dir) + return False + + logger.info( + "Decrypting %d segments from %s to %s", + len(segment_files), + segment_dir, + output_file, + ) + + chunk_num = 0 + temp_files = [] + TARGET_SIZE = 1 * 1024 * 1024 * 1024 # 1GB + + i = 0 + while i < len(segment_files): + chunk_files = [] + chunk_size = 0 + + while i < len(segment_files) and chunk_size < TARGET_SIZE: + fname = segment_files[i] + fpath = os.path.join(segment_dir, fname) + fsize = os.path.getsize(fpath) + chunk_files.append(fname) + chunk_size += fsize + i += 1 + + logger.debug( + "Processing chunk %d: %d segments, %.2f MB", + chunk_num, + len(chunk_files), + chunk_size / (1024 * 1024), + ) + + chunk_merged = os.path.join(segment_dir, f"chunk_{chunk_num}_merged") + chunk_dec = os.path.join(segment_dir, f"chunk_{chunk_num}_dec") + temp_files.extend([chunk_merged, chunk_dec]) + + with open(chunk_merged, "wb") as outfile: + for fname in chunk_files: + with open(os.path.join(segment_dir, fname), "rb") as infile: + outfile.write(infile.read()) + + logger.debug("Decrypting chunk %d", chunk_num) + result = subprocess.run( + [ + "mp4ff-decrypt", + "-init", + init_path, + "-key", + key, + chunk_merged, + chunk_dec, + ], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + logger.error("Decryption failed for chunk %d: %s", chunk_num, result.stderr) + for tf in temp_files: + if os.path.exists(tf): + os.remove(tf) + return False + + chunk_num += 1 + + logger.debug("Concatenating %d decrypted chunks into %s", chunk_num, output_file) + with open(output_file, "wb") as outfile: + for c in range(chunk_num): + chunk_dec = os.path.join(segment_dir, f"chunk_{c}_dec") + with open(chunk_dec, "rb") as infile: + outfile.write(infile.read()) + + for tf in temp_files: + if os.path.exists(tf): + os.remove(tf) + + logger.info("Successfully decrypted to %s", output_file) + return True diff --git a/uv.lock b/uv.lock index 13969cc..eb0561c 100644 --- a/uv.lock +++ b/uv.lock @@ -779,7 +779,7 @@ wheels = [ [[package]] name = "oqee-rewind" -version = "0.1.0" +version = "1.0" source = { editable = "." } dependencies = [ { name = "aiohttp" },