From ae5035c6a8fe5797d0917245bc553c82138a0911 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=88=9A=28noham=29=C2=B2?= <100566912+NohamR@users.noreply.github.com> Date: Thu, 13 Aug 2026 17:40:06 +0200 Subject: [PATCH] optimize DASH probing and segment download paths - Switch bruteforce probes from GET to HEAD to avoid downloading full segment bodies (only the status code is needed) - Cap bruteforce concurrency with a semaphore (default 100) instead of issuing one task per tick, and cancel leftover tasks on early exit - Reuse a single aiohttp session / event loop for the manifest mode's AVC/HEVC/audio bruteforces to keep TCP connections alive - Hoist the per-request segment headers into a module-level constant - Use a set for the already-downloaded resume check in save_segments - Replace read()-into-memory copies with shutil.copyfileobj in merge_segments and decrypt --- main.py | 79 ++++++++++++++++++++++++++++------------------ utils/stream.py | 51 ++++++++++-------------------- utils/times.py | 71 ++++++++++++++++++++++++++--------------- utils/utilities.py | 4 +-- 4 files changed, 114 insertions(+), 91 deletions(-) diff --git a/main.py b/main.py index 3f34381..d04cd82 100644 --- a/main.py +++ b/main.py @@ -9,6 +9,7 @@ import shutil import logging from datetime import datetime, timedelta from dotenv import load_dotenv +import aiohttp from utils.input import ( stream_selection, get_date_input, @@ -190,45 +191,21 @@ if __name__ == "__main__": mpd_content = get_manifest(dash_id) manifest_info = parse_mpd_manifest(mpd_content) - avc_tick = None + avc_track_id = 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.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: logger.warning( "No AVC video track available. AVC tracks will be removed from manifest." ) - hevc_tick = None + hevc_track_id = 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." @@ -239,10 +216,52 @@ if __name__ == "__main__": audio_init_segment = audio_sel["segments"]["initialization"] audio_track_id = audio_init_segment.split("/")[-1].split("_init")[0] - logger.info("Bruteforcing audio track %s...", audio_track_id) - audio_valid_ticks = asyncio.run( - bruteforce(audio_track_id, start_tick_user, batch_size) - ) + async def bruteforce_all(): + async with aiohttp.ClientSession() as session: + avc_ticks = None + hevc_ticks = None + audio_ticks = [] + + if avc_track_id: + logger.info( + "Bruteforcing AVC video track %s...", avc_track_id + ) + avc_valid = await bruteforce( + avc_track_id, start_tick_user, batch_size, session=session + ) + if len(avc_valid) == 0: + logger.warning( + "No valid ticks found for AVC video. AVC tracks will be removed from manifest." + ) + else: + avc_ticks = avc_valid[0] + logger.info("AVC video tick found: %s", avc_ticks) + + if hevc_track_id: + logger.info( + "Bruteforcing HEVC video track %s...", hevc_track_id + ) + hevc_valid = await bruteforce( + hevc_track_id, start_tick_user, batch_size, session=session + ) + if len(hevc_valid) == 0: + logger.warning( + "No valid ticks found for HEVC video. HEVC tracks will be removed from manifest." + ) + else: + hevc_ticks = hevc_valid[0] + logger.info("HEVC video tick found: %s", hevc_ticks) + + logger.info("Bruteforcing audio track %s...", audio_track_id) + audio_valid = await bruteforce( + audio_track_id, start_tick_user, batch_size, session=session + ) + if len(audio_valid) > 0: + audio_ticks = audio_valid + + return avc_ticks, hevc_ticks, audio_ticks + + avc_tick, hevc_tick, audio_valid_ticks = asyncio.run(bruteforce_all()) if len(audio_valid_ticks) == 0: logger.error("No valid ticks found for audio.") sys.exit(1) diff --git a/utils/stream.py b/utils/stream.py index 06a1950..c344bb6 100644 --- a/utils/stream.py +++ b/utils/stream.py @@ -18,6 +18,16 @@ ET.register_namespace("", "urn:mpeg:dash:schema:mpd:2011") ET.register_namespace("xsi", "http://www.w3.org/2001/XMLSchema-instance") ET.register_namespace("cenc", "urn:mpeg:cenc:2013") +SEGMENT_HEADERS = { + "Accept": "*/*", + "Referer": "https://tv.free.fr/", + "User-Agent": ( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/143.0.0.0 Safari/537.36" + ), +} + def parse_mpd_manifest(mpd_content: str) -> Dict[str, Any]: """Parse an MPD manifest and extract metadata. @@ -340,7 +350,7 @@ def get_manifest(manifest_id): async def fetch_segment(session, ticks, track_id): - """Fetch a media segment asynchronously. + """Probe whether a media segment exists using a lightweight HEAD request. Args: session: The aiohttp ClientSession. @@ -348,20 +358,11 @@ async def fetch_segment(session, ticks, track_id): track_id: The track identifier. Returns: - The tick value if successful, None otherwise. + The tick value if the segment exists, None otherwise. """ url = f"https://media4.stream.proxad.net/media/{track_id}_{ticks}" - headers = { - "Accept": "*/*", - "Referer": "https://tv.free.fr/", - "User-Agent": ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/143.0.0.0 Safari/537.36" - ), - } try: - async with session.get(url, headers=headers) as resp: + async with session.head(url, headers=SEGMENT_HEADERS) as resp: if resp.status == 200: return ticks return None @@ -377,16 +378,7 @@ def get_init(output_folder, track_id): track_id: The track identifier. """ url = f"https://media4.stream.proxad.net/media/{track_id}_init" - headers = { - "Accept": "*/*", - "Referer": "https://tv.free.fr/", - "User-Agent": ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/143.0.0.0 Safari/537.36" - ), - } - response = requests.get(url, headers=headers, timeout=10) + response = requests.get(url, headers=SEGMENT_HEADERS, timeout=10) if response.status_code == 200: os.makedirs(f"{output_folder}/segments_{track_id}", exist_ok=True) init_path = f"{output_folder}/segments_{track_id}/init.mp4" @@ -418,17 +410,8 @@ async def save_segments( Tuple of (success: bool, tick: int, rep: int) """ url = f"https://media4.stream.proxad.net/media/{track_id}_{tick}" - headers = { - "Accept": "*/*", - "Referer": "https://tv.free.fr/", - "User-Agent": ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/143.0.0.0 Safari/537.36" - ), - } try: - async with session.get(url, headers=headers) as resp: + async with session.get(url, headers=SEGMENT_HEADERS) as resp: if resp.status == 200: content = await resp.read() filename = f"{output_folder}/segments_{track_id}/{tick}.m4s" @@ -458,11 +441,11 @@ async def save_segments( 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 = [ + 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 ] diff --git a/utils/times.py b/utils/times.py index 8c9dd47..0a77c36 100644 --- a/utils/times.py +++ b/utils/times.py @@ -55,45 +55,66 @@ def future(rep, base, duration): return base + rep * duration -async def bruteforce(track_id, date, batch_size=20000): - """Bruteforce segments to find valid ticks.""" +async def bruteforce(track_id, date, batch_size=20000, concurrency=100, session=None): + """Bruteforce segments to find valid ticks. + + Args: + track_id: The track identifier. + date: The base tick to start probing from. + batch_size: Number of ticks to probe per batch. + concurrency: Maximum number of in-flight requests at once. + session: Optional shared aiohttp session. A new one is created if omitted. + + Returns: + List of valid ticks found (empty if none). + """ valid_ticks = [] total_requests = 288000 logger.debug("Starting bruteforce for %s near (%s) %s", track_id, date, convert_sec_to_date(convert_ticks_to_sec(date, TIMESCALE))) start_time = time.time() + owns_session = session is None + if owns_session: + session = aiohttp.ClientSession() try: - async with aiohttp.ClientSession() as session: - for batch_start in range(0, total_requests, batch_size): - batch_end = min(batch_start + batch_size, total_requests) - tasks = [ - fetch_segment(session, t + date, track_id) - for t in range(batch_start, batch_end) - ] + semaphore = asyncio.Semaphore(concurrency) - results = [] - for coro in tqdm( - asyncio.as_completed(tasks), - total=len(tasks), - desc="Bruteforce", - unit="req", - ): - result = await coro - results.append(result) + async def probe(t): + async with semaphore: + return await fetch_segment(session, t + date, track_id) - valid_ticks.extend( - [r for r in results if r and not isinstance(r, Exception)] - ) + for batch_start in range(0, total_requests, batch_size): + batch_end = min(batch_start + batch_size, total_requests) + tasks = [probe(t) for t in range(batch_start, batch_end)] - # Stop if we found valid ticks - if valid_ticks: - logger.debug("Found valid ticks: %s, stopping bruteforce.", valid_ticks) - break + results = [] + for coro in tqdm( + asyncio.as_completed(tasks), + total=len(tasks), + desc="Bruteforce", + unit="req", + ): + result = await coro + results.append(result) + + valid_ticks.extend( + [r for r in results if r and not isinstance(r, Exception)] + ) + + # Stop if we found valid ticks + if valid_ticks: + logger.debug("Found valid ticks: %s, stopping bruteforce.", valid_ticks) + for task in tasks: + task.cancel() + break except KeyboardInterrupt: logger.error("Interrupted by user (Ctrl+C)") + finally: + if owns_session: + await session.close() elapsed = time.time() - start_time logger.debug("Completed in %.2fs", elapsed) diff --git a/utils/utilities.py b/utils/utilities.py index d93a623..c806aed 100644 --- a/utils/utilities.py +++ b/utils/utilities.py @@ -39,7 +39,7 @@ def merge_segments(input_folder: str, track_id: str, output_file: str): with open(output_file, "wb") as outfile: for fname in segment_files: with open(f"{segment_folder}/{fname}", "rb") as infile: - outfile.write(infile.read()) + shutil.copyfileobj(infile, outfile) logger.info("Merged segments into %s", output_file) @@ -104,7 +104,7 @@ def decrypt(segment_dir, init_path, output_file, key): 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()) + shutil.copyfileobj(infile, outfile) logger.debug("Decrypting chunk %d", chunk_num) result = subprocess.run(