diff --git a/docs/reference/configuration/download.md b/docs/reference/configuration/download.md index cc40f586..4ebc4b7a 100644 --- a/docs/reference/configuration/download.md +++ b/docs/reference/configuration/download.md @@ -135,7 +135,7 @@ continue_downloads: true - **Type:** `bool`  ยท  **Default:** `false` -Appends each DASH segment to the output file as soon as it and every segment before it have +Appends each segment to the output file as soon as it and every segment before it have arrived, instead of joining all the segment files after the last one lands. The join runs during the download, so the "Merging" step after the last segment goes away. The output file is the same either way. @@ -151,8 +151,10 @@ the read usually comes from the page cache. On a hard disk, the append to one fi workers write many segment files can still make the head seek. If a download gets slower with the option on, turn it off. -This applies to DASH video and audio tracks only. Subtitle tracks, HLS tracks and ISM -tracks keep the merge after the download. +This applies to video and audio tracks from DASH, ISM, and HLS. An HLS track qualifies +only when every segment shares one init segment and one `EXT-X-KEY`, with no +`EXT-X-DISCONTINUITY`, and the DRM is Widevine, PlayReady, or none. AES-128 tracks, tracks +with content key rotation or discontinuities, and subtitle tracks keep the merge after the download. !!! note "No resume" diff --git a/tests/core/test_rolling_merge_hls_ism.py b/tests/core/test_rolling_merge_hls_ism.py new file mode 100644 index 00000000..b6cb4501 --- /dev/null +++ b/tests/core/test_rolling_merge_hls_ism.py @@ -0,0 +1,266 @@ +"""Rolling merge through ``HLS.download_track`` and ``ISM.download_track``. + +Same localhost-server shape as the DASH collapse test. Each parser downloads one unencrypted +track twice, with ``merge_segments`` off and on, and the two output files must be byte +identical with no segment directory left behind, and the rolling run must report no Merging +phase (the post-download run must). A second HLS case with a discontinuity checks that an +ineligible playlist keeps the post-download merge and still produces the same bytes. A stalled +drain must raise and leave no output file. +""" + +from __future__ import annotations + +import threading +from concurrent.futures import Future +from functools import partial +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from uuid import UUID + +import pytest +from Cryptodome.Cipher import AES +from Cryptodome.Util.Padding import pad +from pywidevine.pssh import PSSH + +from tests.core.test_ism_init import piff_fragment +from unshackle.core.config import config +from unshackle.core.constants import DOWNLOAD_CANCELLED +from unshackle.core.drm import Widevine +from unshackle.core.manifests import HLS, ISM +from unshackle.core.manifests import hls as hls_module +from unshackle.core.manifests import ism as ism_module +from unshackle.core.manifests.dash import RollingMerge +from unshackle.core.manifests.ism_init import PIFF_SENC_UUID +from unshackle.core.tracks.track import DownloadContext + +INIT = bytes(range(100)) +PARTS = [bytes([0xA0 + i]) * (150 + 37 * i) for i in range(5)] +KID = UUID(hex="11111111111111111111111111111111") +AES_KEY = bytes(range(16)) +DECRYPTED_INIT = b"decrypted-init" * 8 +# one PIFF fragment per index, each with its own mdat payload so a misordered append shows up +PIFF_PARTS = [piff_fragment().replace(b"\xab" * 64, bytes([0xC0 + i]) * 64) for i in range(len(PARTS))] + + +class _Handler(BaseHTTPRequestHandler): + def do_GET(self): + body = self.server.routes.get(self.path.split("?")[0]) + if body is None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + return + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + +@pytest.fixture +def server(): + srv = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + srv.routes = {"/init.mp4": INIT} + for i, part in enumerate(PARTS): + srv.routes[f"/part{i}.m4s"] = srv.routes[f"/x.ism/part{i}.m4s"] = part + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + yield srv + srv.shutdown() + thread.join(timeout=5) + + +@pytest.fixture(autouse=True) +def no_cancel(): + DOWNLOAD_CANCELLED.clear() + yield + DOWNLOAD_CANCELLED.clear() + + +def base(srv) -> str: + return f"http://127.0.0.1:{srv.server_address[1]}/" + + +def make_ctx(tmp_path, name): + statuses: list[str] = [] + + def progress(**kw): + if kw.get("downloaded"): + statuses.append(kw["downloaded"]) + + ctx = DownloadContext( + save_path=tmp_path / name, + save_dir=tmp_path / f"{name}_segments", + progress=partial(progress), + max_workers=4, + ) + ctx.statuses = statuses + return ctx + + +def hls_track(srv, discontinuity: bool): + lines = ["#EXTM3U", "#EXT-X-VERSION:6", "#EXT-X-TARGETDURATION:4", '#EXT-X-MAP:URI="init.mp4"'] + for i in range(len(PARTS)): + if discontinuity and i == 2: + lines += ["#EXT-X-DISCONTINUITY", '#EXT-X-MAP:URI="init.mp4"'] + lines += ["#EXTINF:4.0,", f"part{i}.m4s"] + lines.append("#EXT-X-ENDLIST") + srv.routes["/media.m3u8"] = "\n".join(lines).encode() + master = '#EXTM3U\n#EXT-X-STREAM-INF:BANDWIDTH=1000000,RESOLUTION=1920x1080,CODECS="avc1.640028"\nmedia.m3u8\n' + return HLS.from_text(master, url=f"{base(srv)}master.m3u8").to_tracks(language="en").videos[0] + + +def ism_track(srv): + # audio: the init synthesis needs only CodecPrivateData, no first-segment parse + manifest = f""" + + + + + + +""" + return ISM.from_text(manifest, f"{base(srv)}x.ism/manifest").to_tracks().audio[0] + + +def download_both_ways(parser, make, tmp_path, monkeypatch, expect_rolling: bool) -> tuple[bytes, bytes]: + monkeypatch.setattr(config, "merge_segments", False) + plain = make_ctx(tmp_path, "plain.mp4") + parser.download_track(make(), plain) + assert "Merging" in plain.statuses + + monkeypatch.setattr(config, "merge_segments", True) + rolling = make_ctx(tmp_path, "rolling.mp4") + parser.download_track(make(), rolling) + assert not rolling.save_dir.exists() or not list(rolling.save_dir.glob("*")) + # the rolling path never reports a Merging phase; an ineligible track must still report one + assert ("Merging" not in rolling.statuses) == expect_rolling + return plain.save_path.read_bytes(), rolling.save_path.read_bytes() + + +def test_hls_single_init_rolling_matches_post_merge(server, tmp_path, monkeypatch): + plain, rolling = download_both_ways(HLS, lambda: hls_track(server, False), tmp_path, monkeypatch, True) + assert rolling == plain == INIT + b"".join(PARTS) + + +def test_hls_discontinuity_keeps_post_merge(server, tmp_path, monkeypatch): + plain, rolling = download_both_ways(HLS, lambda: hls_track(server, True), tmp_path, monkeypatch, False) + assert rolling == plain + assert INIT in rolling[len(INIT) :] + + +def test_ism_rolling_matches_post_merge(server, tmp_path, monkeypatch): + plain, rolling = download_both_ways(ISM, lambda: ism_track(server), tmp_path, monkeypatch, True) + assert rolling == plain + assert rolling.endswith(b"".join(PARTS)) + assert len(rolling) > len(b"".join(PARTS)) + + +@pytest.mark.parametrize("parser, make", [(HLS, hls_track), (ISM, lambda srv, _d=None: ism_track(srv))]) +def test_incomplete_rolling_merge_raises_and_leaves_no_output(server, tmp_path, monkeypatch, parser, make): + monkeypatch.setattr(config, "merge_segments", True) + # a drain that never reaches the end must not ship a short file as a finished track + monkeypatch.setattr(RollingMerge, "add", lambda self, *a, **kw: None) + ctx = make_ctx(tmp_path, "short.mp4") + with pytest.raises(FileNotFoundError, match="Rolling merge appended 0 of"): + parser.download_track(make(server, False), ctx) + assert not ctx.save_path.exists() + + +class FakeDRM: + """The attributes the ISM merge reads from a licensed DRM: the KID for the synthesised tenc.""" + + kids = [KID] + + +class FakeSegmentDecrypter: + """Per-segment decryption that hands each segment straight back, with a fixed decrypted init.""" + + def __init__(self, drm, init_content, workdir, max_workers): + self.submitted: list = [] + + def init_bytes(self) -> bytes: + return DECRYPTED_INIT + + def submit(self, segment) -> Future: + self.submitted.append(segment) + future: Future = Future() + future.set_result(segment) + return future + + def finish(self) -> bytes: + return DECRYPTED_INIT + + def close(self) -> None: + pass + + +def no_decrypt(monkeypatch, module): + # the merged file is compared byte for byte, so the decrypt step and its check are stubbed + monkeypatch.setattr(module, "decrypt_track", lambda *a, **kw: None) + monkeypatch.setattr(module, "assert_fragments_decrypted", lambda *a, **kw: None) + + +def test_ism_encrypted_rolling_matches_post_merge(server, tmp_path, monkeypatch): + for i, part in enumerate(PIFF_PARTS): + server.routes[f"/x.ism/part{i}.m4s"] = part + no_decrypt(monkeypatch, ism_module) + + def make(): + track = ism_track(server) + track.drm = [FakeDRM()] + return track + + plain, rolling = download_both_ways(ISM, make, tmp_path, monkeypatch, True) + assert rolling == plain + # every fragment went through the PIFF to CENC rewrite, and the head carries the KID + assert PIFF_SENC_UUID not in rolling + assert rolling.count(b"senc") == len(PIFF_PARTS) + assert bytes.fromhex(KID.hex) in rolling + assert rolling.index(b"moov") < rolling.index(b"moof") + for i in range(len(PIFF_PARTS)): + assert bytes([0xC0 + i]) * 64 in rolling + + +def test_hls_segment_decrypt_rolling_writes_decrypted_init(server, tmp_path, monkeypatch): + monkeypatch.setattr(hls_module, "can_use", lambda *a, **kw: True) + monkeypatch.setattr(hls_module, "SegmentDecrypter", FakeSegmentDecrypter) + no_decrypt(monkeypatch, hls_module) + + def make(): + track = hls_track(server, False) + track.drm = [Widevine(pssh=PSSH.new(system_id=PSSH.SystemId.Widevine, key_ids=[KID]))] + return track + + def download(name, merge_segments): + monkeypatch.setattr(config, "merge_segments", merge_segments) + ctx = make_ctx(tmp_path, name) + ctx.license_widevine = lambda *a, **kw: None + HLS.download_track(make(), ctx) + return ctx + + plain = download("plain.mp4", False) + rolling = download("rolling.mp4", True) + assert "Merging" in plain.statuses and "Merging" not in rolling.statuses + assert "Decrypting" in rolling.statuses + assert rolling.save_path.read_bytes() == plain.save_path.read_bytes() == DECRYPTED_INIT + b"".join(PARTS) + assert not rolling.save_dir.exists() or not list(rolling.save_dir.glob("*")) + + +def test_hls_aes128_keeps_post_merge(server, tmp_path, monkeypatch): + # AES-128 with no IV attribute: each segment's IV is its media sequence number + for i, part in enumerate(PARTS): + server.routes[f"/part{i}.m4s"] = AES.new(AES_KEY, AES.MODE_CBC, i.to_bytes(16, "big")).encrypt(pad(part, 16)) + server.routes["/key.bin"] = AES_KEY + + def make(): + track = hls_track(server, False) + playlist = server.routes["/media.m3u8"].decode() + key_line = f'#EXT-X-KEY:METHOD=AES-128,URI="{base(server)}key.bin"\n' + server.routes["/media.m3u8"] = playlist.replace("#EXTINF", key_line + "#EXTINF", 1).encode() + return track + + plain, rolling = download_both_ways(HLS, make, tmp_path, monkeypatch, False) + assert rolling == plain == INIT + b"".join(PARTS) diff --git a/tests/test_rolling_merge.py b/tests/test_rolling_merge.py index cf089762..740cb03d 100644 --- a/tests/test_rolling_merge.py +++ b/tests/test_rolling_merge.py @@ -34,14 +34,14 @@ def test_out_of_order_segments_match_whole_merge(tmp_path: Path) -> None: width = write_segments(save_dir, payloads) dst = io.BytesIO() dst.write(init) - merger = RollingMerge(dst, save_dir, width) + merger = RollingMerge(dst) order = list(range(len(payloads))) rng.shuffle(order) for index in order: future: Future = Future() future.set_result(None) - merger.add(index, future if index % 2 else None) + merger.add(index, save_dir / f"{index:0{width}d}.mp4", future if index % 2 else None) # the cursor never rests on a segment that is ready: every contiguous run is drained assert merger.merged not in merger.ready diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 20ee9b5f..1a1c18ca 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -62,31 +62,33 @@ class RollingMerge: A cursor holds the next expected index; every `add` records one ready segment and then drains every contiguous ready segment from the cursor onward. A segment that carries a future (per-segment decryption) waits for that future before the append, so the output never holds raw bytes. + + `append(dst, segment_file)` writes one segment; the default is a plain byte copy. A parser + whose merge needs a per-segment transform, or a head that only exists once segment 0 has + landed (ISM), passes its own. """ - def __init__(self, dst: BinaryIO, save_dir: Path, width: int) -> None: + def __init__(self, dst: BinaryIO, append: Optional[Callable[[BinaryIO, Path], None]] = None) -> None: self.dst = dst - self.save_dir = save_dir - self.width = width + self.append = append or (lambda dst, path: append_segment(dst, path, is_text_subtitle=False)) self.cursor = 0 - self.ready: dict[int, Optional[Future]] = {} + self.ready: dict[int, tuple[Path, Optional[Future]]] = {} - def add(self, index: int, future: Optional[Future] = None) -> None: + def add(self, index: int, segment_file: Path, future: Optional[Future] = None) -> None: # ponytail: runs on the downloader's own loop, so a decrypt wait or a catch-up burst # delays new submissions; move the drain to a writer thread if throughput measurably drops - self.ready[index] = future + self.ready[index] = (segment_file, future) while self.cursor in self.ready: - pending = self.ready.pop(self.cursor) + path, pending = self.ready.pop(self.cursor) if pending is not None: pending.result() - segment_file = self.save_dir / f"{self.cursor:0{self.width}d}.mp4" - append_segment(self.dst, segment_file, is_text_subtitle=False) + self.append(self.dst, path) self.cursor += 1 # the bytes are already in the output; a handle another process still holds on # the just-replaced file (Windows) must not fail the track, the segment dir sweep # after the download removes any leftover try: - segment_file.unlink() + path.unlink() except OSError: pass @@ -609,7 +611,7 @@ class DASH: head = decrypter.init_bytes() if decrypter else init_data if head: output.write(head) - merger = RollingMerge(output, save_dir, len(str(len(segments)))) + merger = RollingMerge(output) # an exception raised in this body (a decrypt failure re-raised by the merge) # must close the generator now, so its finally aborts the workers @@ -621,7 +623,7 @@ class DASH: # listeners get the segment while it still exists; the merge unlinks it events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded) if merger: - merger.add(int(file_downloaded.stem), future) + merger.add(int(file_downloaded.stem), file_downloaded, future) else: downloaded = status_update.get("downloaded") if downloaded and downloaded.endswith("/s"): @@ -630,7 +632,7 @@ class DASH: if decrypter: init_data = decrypter.finish() - except Exception: + except BaseException: if decrypter: decrypter.close() if merger: diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 5ffb4469..f3e86e12 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -10,6 +10,7 @@ import shutil import subprocess import sys import time +from contextlib import ExitStack, closing from pathlib import Path from typing import Any, Optional, Union from urllib.parse import urljoin, urlparse @@ -34,6 +35,7 @@ from unshackle.core.drm import DRM_T, ClearKey, MonaLisa, PlayReady, Widevine from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use from unshackle.core.drm.verify import decrypt_track from unshackle.core.events import events +from unshackle.core.manifests.dash import RollingMerge from unshackle.core.session import RnetResponse, RnetSession from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume from unshackle.core.tracks.track import DRM_PREFERENCE_TYPES, assert_fragments_decrypted @@ -752,10 +754,15 @@ class HLS: # segment so the IV uses the absolute media sequence number, not the download index. wanted_segments = [seg for seg in master.segments if seg not in unwanted_segments] + single_init = HLS.single_init_no_rotation(wanted_segments, initial_drm_key, cdm, track.drm_preference) use_segment_decrypt = ( - isinstance(session_drm, (Widevine, PlayReady)) - and can_use(session_drm, config.decryption) - and HLS.single_init_no_rotation(wanted_segments, initial_drm_key, cdm, track.drm_preference) + isinstance(session_drm, (Widevine, PlayReady)) and can_use(session_drm, config.decryption) and single_init + ) + use_rolling_merge = ( + config.merge_segments + and single_init + and not isinstance(track, Subtitle) + and (session_drm is None or isinstance(session_drm, (Widevine, PlayReady))) ) total_segments = len(master.segments) - len(unwanted_segments) @@ -795,6 +802,7 @@ class HLS: # earlier run would either be decrypted twice or never at all resumable_drm = ( not use_segment_decrypt + and not use_rolling_merge and (session_drm is None or isinstance(session_drm, (Widevine, PlayReady))) and not any(key and key.method == "AES-128" for key in (master.keys or [])) ) @@ -839,27 +847,39 @@ class HLS: map_data: Optional[tuple[m3u8.model.InitializationSection, bytes]] = None segment_decrypter: Optional[SegmentDecrypter] = None - if use_segment_decrypt: + if use_segment_decrypt or use_rolling_merge: init_section = wanted_segments[0].init_section init_content, _ = HLS.fetch_init_section(session, init_section, 0) map_data = (init_section, init_content) + if use_segment_decrypt: segment_decrypter = SegmentDecrypter(session_drm, init_content, save_dir.parent, max_workers) + merger: Optional[RollingMerge] = None try: - for status_update in downloader(**downloader_args): - file_downloaded = status_update.get("file_downloaded") - if file_downloaded: - if segment_decrypter: - segment_decrypter.submit(file_downloaded) - events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded) - else: - downloaded = status_update.get("downloaded") - if downloaded and downloaded.endswith("/s"): - status_update["downloaded"] = f"HLS {downloaded}" - progress(**status_update) - except Exception: + with ExitStack() as stack: + if use_rolling_merge: + output = stack.enter_context(open(save_path, "wb")) + output.write(segment_decrypter.init_bytes() if segment_decrypter else init_content) + merger = RollingMerge(output) + + stream = stack.enter_context(closing(downloader(**downloader_args))) + for status_update in stream: + file_downloaded = status_update.get("file_downloaded") + if file_downloaded: + future = segment_decrypter.submit(file_downloaded) if segment_decrypter else None + events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded) + if merger: + merger.add(int(file_downloaded.stem), file_downloaded, future) + else: + downloaded = status_update.get("downloaded") + if downloaded and downloaded.endswith("/s"): + status_update["downloaded"] = f"HLS {downloaded}" + progress(**status_update) + except BaseException: if segment_decrypter: segment_decrypter.close() + if merger: + save_path.unlink(missing_ok=True) raise for control_file in segment_save_dir.glob("*.!dev"): @@ -868,6 +888,52 @@ class HLS: # merge mutates segment files in place from here on; withdraw the reuse proof resume.clear_sidecar(save_dir) + if merger: + if segment_decrypter: + segment_decrypter.finish() + if merger.merged != total_segments: + error_msg = f"Rolling merge appended {merger.merged} of {total_segments} segments: {save_dir}" + log_event( + "manifest_hls_rolling_merge_incomplete", + level="ERROR", + message=error_msg, + context={ + "track_id": getattr(track, "id", None), + "track_type": track.__class__.__name__, + "save_dir": str(save_dir), + "merged": merger.merged, + "total_segments": total_segments, + "downloader": "requests", + }, + ) + save_path.unlink(missing_ok=True) + raise FileNotFoundError(error_msg) + if save_dir.name.endswith("_segments"): + shutil.rmtree(save_dir, ignore_errors=True) + log_event( + "manifest_hls_download_complete", + level="DEBUG", + message="HLS download complete, segments appended during the download", + context={ + "track_id": getattr(track, "id", None), + "track_type": track.__class__.__name__, + "save_dir": str(save_dir), + "segments_found": merger.merged, + "rolling_merge": True, + "downloader": "requests", + }, + ) + if session_drm: + progress(downloaded="Decrypting", completed=0, total=None) + decrypt_track(session_drm, save_path, license_widevine, decrypt=not segment_decrypter) + assert_fragments_decrypted(save_path) + track.drm = None + events.emit(events.Types.TRACK_DECRYPTED, track=track, drm=session_drm, segment=None) + progress(downloaded="Downloaded", completed=100, total=100) + track.path = save_path + events.emit(events.Types.TRACK_DOWNLOADED, track=track) + return + progress(total=total_segments, completed=0, downloaded="Merging") name_len = len(str(total_segments)) diff --git a/unshackle/core/manifests/ism.py b/unshackle/core/manifests/ism.py index 43915b7f..a29f71ff 100644 --- a/unshackle/core/manifests/ism.py +++ b/unshackle/core/manifests/ism.py @@ -8,7 +8,9 @@ import re import shutil import struct import urllib.parse -from typing import Any, Optional, Union +from contextlib import ExitStack, closing +from pathlib import Path +from typing import Any, BinaryIO, Optional, Union import requests from langcodes import Language, tag_is_valid @@ -22,6 +24,7 @@ from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, from unshackle.core.drm import DRM_T, PlayReady, Widevine from unshackle.core.drm.verify import decrypt_track from unshackle.core.events import events +from unshackle.core.manifests.dash import RollingMerge, append_segment from unshackle.core.manifests.ism_init import ( build_init_segment, parse_codec_private_data_vui, @@ -437,11 +440,14 @@ class ISM: downloader = track.downloader + use_rolling_merge = config.merge_segments and not isinstance(track, Subtitle) + can_resume = config.continue_downloads and not use_rolling_merge + digest = resume.fingerprint(track.url, [(url, None) for url in segments]) - if not (config.continue_downloads and resume.reusable(save_dir, digest)): + if not (can_resume and resume.reusable(save_dir, digest)): shutil.rmtree(save_dir, ignore_errors=True) save_dir.mkdir(parents=True, exist_ok=True) - if config.continue_downloads: + if can_resume: resume.write_sidecar(save_dir, digest) downloader_args = dict( @@ -472,15 +478,45 @@ class ISM: }, ) - for status_update in downloader(**downloader_args): - file_downloaded = status_update.get("file_downloaded") - if file_downloaded: - events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded) + iv_size = 8 + + def append_ism(dst: BinaryIO, segment_file: Path) -> None: + """Write one fragment; segment 0 also synthesises the init box the fragments lack.""" + nonlocal iv_size + if int(segment_file.stem) == 0: + segment_data = segment_file.read_bytes() + init_segment = ISM.init_segment(track, session_drm, segment_data) + if init_segment: + dst.write(init_segment) + iv_size = (read_per_sample_iv_size(segment_data) if session_drm else None) or 8 + dst.write(piff_senc_to_cenc(segment_data, iv_size) if session_drm else segment_data) + elif session_drm: + dst.write(piff_senc_to_cenc(segment_file.read_bytes(), iv_size)) else: - downloaded = status_update.get("downloaded") - if downloaded and downloaded.endswith("/s"): - status_update["downloaded"] = f"ISM {downloaded}" - progress(**status_update) + append_segment(dst, segment_file, is_text_subtitle=False) + + merger: Optional[RollingMerge] = None + try: + with ExitStack() as stack: + if use_rolling_merge: + output = stack.enter_context(open(save_path, "wb")) + merger = RollingMerge(output, append_ism) + stream = stack.enter_context(closing(downloader(**downloader_args))) + for status_update in stream: + file_downloaded = status_update.get("file_downloaded") + if file_downloaded: + events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded) + if merger: + merger.add(int(file_downloaded.stem), file_downloaded) + else: + downloaded = status_update.get("downloaded") + if downloaded and downloaded.endswith("/s"): + status_update["downloaded"] = f"ISM {downloaded}" + progress(**status_update) + except BaseException: + if merger: + save_path.unlink(missing_ok=True) + raise if not save_dir.exists(): error_msg = f"Output directory does not exist: {save_dir}" @@ -501,7 +537,25 @@ class ISM: for control_file in save_dir.glob("*.!dev"): control_file.unlink(missing_ok=True) - segments_to_merge = [x for x in sorted(save_dir.iterdir()) if x.is_file()] + if merger and merger.merged != len(segments): + error_msg = f"Rolling merge appended {merger.merged} of {len(segments)} segments: {save_dir}" + log_event( + "manifest_ism_rolling_merge_incomplete", + level="ERROR", + message=error_msg, + context={ + "track_id": getattr(track, "id", None), + "track_type": track.__class__.__name__, + "save_dir": str(save_dir), + "merged": merger.merged, + "total_segments": len(segments), + "downloader": "requests", + }, + ) + save_path.unlink(missing_ok=True) + raise FileNotFoundError(error_msg) + + segments_to_merge = [] if merger else [x for x in sorted(save_dir.iterdir()) if x.is_file()] log_event( "manifest_ism_download_complete", @@ -512,13 +566,14 @@ class ISM: "track_type": track.__class__.__name__, "save_dir": str(save_dir), "save_dir_exists": save_dir.exists(), - "segments_found": len(segments_to_merge), + "segments_found": merger.merged if merger else len(segments_to_merge), "segment_files": [f.name for f in segments_to_merge[:10]], + "rolling_merge": merger is not None, "downloader": "requests", }, ) - if not segments_to_merge: + if not segments_to_merge and not merger: error_msg = f"No segment files found in output directory: {save_dir}" all_contents = list(save_dir.iterdir()) if save_dir.exists() else [] log_event( @@ -535,40 +590,41 @@ class ISM: ) raise FileNotFoundError(error_msg) - is_text_subtitle = ( - not session_drm - and isinstance(track, Subtitle) - and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML) - ) - progress(downloaded="Merging", completed=0, total=len(segments_to_merge)) - with open(save_path, "wb") as f: - first_segment = segments_to_merge[0].read_bytes() if segments_to_merge else None - init_segment = ISM.init_segment(track, session_drm, first_segment) - if init_segment: - f.write(init_segment) - iv_size = (read_per_sample_iv_size(first_segment) if session_drm and first_segment else None) or 8 - for index, segment_file in enumerate(segments_to_merge): - if is_text_subtitle: - # first segment was already read for the init synthesis, reuse it - segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes() - segment_data = try_ensure_utf8(segment_data) - segment_data = ( - segment_data.decode("utf8") - .replace("‎", html.unescape("‎")) - .replace("‏", html.unescape("‏")) - .encode("utf8") - ) - f.write(segment_data) - elif session_drm: - segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes() - f.write(piff_senc_to_cenc(segment_data, iv_size)) - elif index == 0 and first_segment: - f.write(first_segment) - else: - with open(segment_file, "rb") as src: - shutil.copyfileobj(src, f, 1024 * 1024) - segment_file.unlink() - progress(advance=1) + if not merger: + is_text_subtitle = ( + not session_drm + and isinstance(track, Subtitle) + and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML) + ) + progress(downloaded="Merging", completed=0, total=len(segments_to_merge)) + with open(save_path, "wb") as f: + first_segment = segments_to_merge[0].read_bytes() if segments_to_merge else None + init_segment = ISM.init_segment(track, session_drm, first_segment) + if init_segment: + f.write(init_segment) + iv_size = (read_per_sample_iv_size(first_segment) if session_drm and first_segment else None) or 8 + for index, segment_file in enumerate(segments_to_merge): + if is_text_subtitle: + # first segment was already read for the init synthesis, reuse it + segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes() + segment_data = try_ensure_utf8(segment_data) + segment_data = ( + segment_data.decode("utf8") + .replace("‎", html.unescape("‎")) + .replace("‏", html.unescape("‏")) + .encode("utf8") + ) + f.write(segment_data) + elif session_drm: + segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes() + f.write(piff_senc_to_cenc(segment_data, iv_size)) + elif index == 0 and first_segment: + f.write(first_segment) + else: + with open(segment_file, "rb") as src: + shutil.copyfileobj(src, f, 1024 * 1024) + segment_file.unlink() + progress(advance=1) track.path = save_path events.emit(events.Types.TRACK_DOWNLOADED, track=track) diff --git a/unshackle/core/service.py b/unshackle/core/service.py index 49246331..ca7cedc5 100644 --- a/unshackle/core/service.py +++ b/unshackle/core/service.py @@ -710,7 +710,8 @@ class Service(metaclass=ABCMeta): Parameters: track: The decrypted Track object. drm: The DRM object it decrypted with. - segment: The decrypted HLS segment information. + segment: The decrypted HLS segment information. None for DASH and ISM tracks and + for an HLS track that merged during the download (merge_segments). """ def on_track_repacked(self, track: AnyTrack) -> None: diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 73e370f6..6aff7909 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -619,9 +619,7 @@ class Track: else: save_dir = save_path.parent - keep_segments = ( - config.continue_downloads and not config.merge_segments and self.descriptor != self.Descriptor.URL - ) + keep_segments = config.continue_downloads and self.descriptor != self.Descriptor.URL def cleanup() -> None: save_path.unlink(missing_ok=True)