mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-11 19:45:06 +00:00
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:
@@ -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"
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user