feat(manifests): extend merge_segments rolling merge to HLS and ISM

RollingMerge takes the segment path and an append hook. ISM synthesises the init from segment 0 and converts PIFF per fragment; HLS qualifies only with one init, one key and no discontinuity. Both parsers check the merge reached every segment, and resume is decided per track.
This commit is contained in:
imSp4rky
2026-09-21 20:53:37 -06:00
parent 79d109d548
commit 6e681af513
8 changed files with 477 additions and 86 deletions
+5 -3
View File
@@ -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"
+266
View File
@@ -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"""<?xml version="1.0"?>
<SmoothStreamingMedia MajorVersion="2" MinorVersion="0" Duration="{len(PARTS) * 20000000}" TimeScale="10000000">
<StreamIndex Type="audio" Name="audio" Language="en" Url="part{{start_time}}.m4s">
<QualityLevel Index="0" Bitrate="128000" FourCC="AACL" AudioTag="255" SamplingRate="48000"
Channels="2" BitsPerSample="16" PacketSize="4" CodecPrivateData="1190"/>
<c t="0" d="1" r="{len(PARTS)}"/>
</StreamIndex>
</SmoothStreamingMedia>
"""
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)
+2 -2
View File
@@ -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
+15 -13
View File
@@ -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:
+82 -16
View File
@@ -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))
+104 -48
View File
@@ -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("&lrm;", html.unescape("&lrm;"))
.replace("&rlm;", html.unescape("&rlm;"))
.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("&lrm;", html.unescape("&lrm;"))
.replace("&rlm;", html.unescape("&rlm;"))
.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)
+2 -1
View File
@@ -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:
+1 -3
View File
@@ -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)