mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-14 21:15: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`
|
- **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
|
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
|
during the download, so the "Merging" step after the last segment goes away. The output file
|
||||||
is the same either way.
|
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
|
workers write many segment files can still make the head seek. If a download gets slower with
|
||||||
the option on, turn it off.
|
the option on, turn it off.
|
||||||
|
|
||||||
This applies to DASH video and audio tracks only. Subtitle tracks, HLS tracks and ISM
|
This applies to video and audio tracks from DASH, ISM, and HLS. An HLS track qualifies
|
||||||
tracks keep the merge after the download.
|
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"
|
!!! 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)
|
width = write_segments(save_dir, payloads)
|
||||||
dst = io.BytesIO()
|
dst = io.BytesIO()
|
||||||
dst.write(init)
|
dst.write(init)
|
||||||
merger = RollingMerge(dst, save_dir, width)
|
merger = RollingMerge(dst)
|
||||||
|
|
||||||
order = list(range(len(payloads)))
|
order = list(range(len(payloads)))
|
||||||
rng.shuffle(order)
|
rng.shuffle(order)
|
||||||
for index in order:
|
for index in order:
|
||||||
future: Future = Future()
|
future: Future = Future()
|
||||||
future.set_result(None)
|
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
|
# the cursor never rests on a segment that is ready: every contiguous run is drained
|
||||||
assert merger.merged not in merger.ready
|
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
|
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
|
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.
|
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.dst = dst
|
||||||
self.save_dir = save_dir
|
self.append = append or (lambda dst, path: append_segment(dst, path, is_text_subtitle=False))
|
||||||
self.width = width
|
|
||||||
self.cursor = 0
|
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
|
# 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
|
# 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:
|
while self.cursor in self.ready:
|
||||||
pending = self.ready.pop(self.cursor)
|
path, pending = self.ready.pop(self.cursor)
|
||||||
if pending is not None:
|
if pending is not None:
|
||||||
pending.result()
|
pending.result()
|
||||||
segment_file = self.save_dir / f"{self.cursor:0{self.width}d}.mp4"
|
self.append(self.dst, path)
|
||||||
append_segment(self.dst, segment_file, is_text_subtitle=False)
|
|
||||||
self.cursor += 1
|
self.cursor += 1
|
||||||
# the bytes are already in the output; a handle another process still holds on
|
# 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
|
# the just-replaced file (Windows) must not fail the track, the segment dir sweep
|
||||||
# after the download removes any leftover
|
# after the download removes any leftover
|
||||||
try:
|
try:
|
||||||
segment_file.unlink()
|
path.unlink()
|
||||||
except OSError:
|
except OSError:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@@ -609,7 +611,7 @@ class DASH:
|
|||||||
head = decrypter.init_bytes() if decrypter else init_data
|
head = decrypter.init_bytes() if decrypter else init_data
|
||||||
if head:
|
if head:
|
||||||
output.write(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)
|
# 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
|
# 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
|
# listeners get the segment while it still exists; the merge unlinks it
|
||||||
events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded)
|
events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded)
|
||||||
if merger:
|
if merger:
|
||||||
merger.add(int(file_downloaded.stem), future)
|
merger.add(int(file_downloaded.stem), file_downloaded, future)
|
||||||
else:
|
else:
|
||||||
downloaded = status_update.get("downloaded")
|
downloaded = status_update.get("downloaded")
|
||||||
if downloaded and downloaded.endswith("/s"):
|
if downloaded and downloaded.endswith("/s"):
|
||||||
@@ -630,7 +632,7 @@ class DASH:
|
|||||||
|
|
||||||
if decrypter:
|
if decrypter:
|
||||||
init_data = decrypter.finish()
|
init_data = decrypter.finish()
|
||||||
except Exception:
|
except BaseException:
|
||||||
if decrypter:
|
if decrypter:
|
||||||
decrypter.close()
|
decrypter.close()
|
||||||
if merger:
|
if merger:
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ import shutil
|
|||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
import time
|
import time
|
||||||
|
from contextlib import ExitStack, closing
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Optional, Union
|
from typing import Any, Optional, Union
|
||||||
from urllib.parse import urljoin, urlparse
|
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.segment_decrypt import SegmentDecrypter, can_use
|
||||||
from unshackle.core.drm.verify import decrypt_track
|
from unshackle.core.drm.verify import decrypt_track
|
||||||
from unshackle.core.events import events
|
from unshackle.core.events import events
|
||||||
|
from unshackle.core.manifests.dash import RollingMerge
|
||||||
from unshackle.core.session import RnetResponse, RnetSession
|
from unshackle.core.session import RnetResponse, RnetSession
|
||||||
from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume
|
from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume
|
||||||
from unshackle.core.tracks.track import DRM_PREFERENCE_TYPES, assert_fragments_decrypted
|
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.
|
# 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]
|
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 = (
|
use_segment_decrypt = (
|
||||||
isinstance(session_drm, (Widevine, PlayReady))
|
isinstance(session_drm, (Widevine, PlayReady)) and can_use(session_drm, config.decryption) and single_init
|
||||||
and can_use(session_drm, config.decryption)
|
)
|
||||||
and HLS.single_init_no_rotation(wanted_segments, initial_drm_key, cdm, track.drm_preference)
|
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)
|
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
|
# earlier run would either be decrypted twice or never at all
|
||||||
resumable_drm = (
|
resumable_drm = (
|
||||||
not use_segment_decrypt
|
not use_segment_decrypt
|
||||||
|
and not use_rolling_merge
|
||||||
and (session_drm is None or isinstance(session_drm, (Widevine, PlayReady)))
|
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 []))
|
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
|
map_data: Optional[tuple[m3u8.model.InitializationSection, bytes]] = None
|
||||||
segment_decrypter: Optional[SegmentDecrypter] = 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_section = wanted_segments[0].init_section
|
||||||
init_content, _ = HLS.fetch_init_section(session, init_section, 0)
|
init_content, _ = HLS.fetch_init_section(session, init_section, 0)
|
||||||
map_data = (init_section, init_content)
|
map_data = (init_section, init_content)
|
||||||
|
if use_segment_decrypt:
|
||||||
segment_decrypter = SegmentDecrypter(session_drm, init_content, save_dir.parent, max_workers)
|
segment_decrypter = SegmentDecrypter(session_drm, init_content, save_dir.parent, max_workers)
|
||||||
|
|
||||||
|
merger: Optional[RollingMerge] = None
|
||||||
try:
|
try:
|
||||||
for status_update in downloader(**downloader_args):
|
with ExitStack() as stack:
|
||||||
file_downloaded = status_update.get("file_downloaded")
|
if use_rolling_merge:
|
||||||
if file_downloaded:
|
output = stack.enter_context(open(save_path, "wb"))
|
||||||
if segment_decrypter:
|
output.write(segment_decrypter.init_bytes() if segment_decrypter else init_content)
|
||||||
segment_decrypter.submit(file_downloaded)
|
merger = RollingMerge(output)
|
||||||
events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded)
|
|
||||||
else:
|
stream = stack.enter_context(closing(downloader(**downloader_args)))
|
||||||
downloaded = status_update.get("downloaded")
|
for status_update in stream:
|
||||||
if downloaded and downloaded.endswith("/s"):
|
file_downloaded = status_update.get("file_downloaded")
|
||||||
status_update["downloaded"] = f"HLS {downloaded}"
|
if file_downloaded:
|
||||||
progress(**status_update)
|
future = segment_decrypter.submit(file_downloaded) if segment_decrypter else None
|
||||||
except Exception:
|
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:
|
if segment_decrypter:
|
||||||
segment_decrypter.close()
|
segment_decrypter.close()
|
||||||
|
if merger:
|
||||||
|
save_path.unlink(missing_ok=True)
|
||||||
raise
|
raise
|
||||||
|
|
||||||
for control_file in segment_save_dir.glob("*.!dev"):
|
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
|
# merge mutates segment files in place from here on; withdraw the reuse proof
|
||||||
resume.clear_sidecar(save_dir)
|
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")
|
progress(total=total_segments, completed=0, downloaded="Merging")
|
||||||
|
|
||||||
name_len = len(str(total_segments))
|
name_len = len(str(total_segments))
|
||||||
|
|||||||
+104
-48
@@ -8,7 +8,9 @@ import re
|
|||||||
import shutil
|
import shutil
|
||||||
import struct
|
import struct
|
||||||
import urllib.parse
|
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
|
import requests
|
||||||
from langcodes import Language, tag_is_valid
|
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 import DRM_T, PlayReady, Widevine
|
||||||
from unshackle.core.drm.verify import decrypt_track
|
from unshackle.core.drm.verify import decrypt_track
|
||||||
from unshackle.core.events import events
|
from unshackle.core.events import events
|
||||||
|
from unshackle.core.manifests.dash import RollingMerge, append_segment
|
||||||
from unshackle.core.manifests.ism_init import (
|
from unshackle.core.manifests.ism_init import (
|
||||||
build_init_segment,
|
build_init_segment,
|
||||||
parse_codec_private_data_vui,
|
parse_codec_private_data_vui,
|
||||||
@@ -437,11 +440,14 @@ class ISM:
|
|||||||
|
|
||||||
downloader = track.downloader
|
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])
|
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)
|
shutil.rmtree(save_dir, ignore_errors=True)
|
||||||
save_dir.mkdir(parents=True, exist_ok=True)
|
save_dir.mkdir(parents=True, exist_ok=True)
|
||||||
if config.continue_downloads:
|
if can_resume:
|
||||||
resume.write_sidecar(save_dir, digest)
|
resume.write_sidecar(save_dir, digest)
|
||||||
|
|
||||||
downloader_args = dict(
|
downloader_args = dict(
|
||||||
@@ -472,15 +478,45 @@ class ISM:
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
for status_update in downloader(**downloader_args):
|
iv_size = 8
|
||||||
file_downloaded = status_update.get("file_downloaded")
|
|
||||||
if file_downloaded:
|
def append_ism(dst: BinaryIO, segment_file: Path) -> None:
|
||||||
events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded)
|
"""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:
|
else:
|
||||||
downloaded = status_update.get("downloaded")
|
append_segment(dst, segment_file, is_text_subtitle=False)
|
||||||
if downloaded and downloaded.endswith("/s"):
|
|
||||||
status_update["downloaded"] = f"ISM {downloaded}"
|
merger: Optional[RollingMerge] = None
|
||||||
progress(**status_update)
|
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():
|
if not save_dir.exists():
|
||||||
error_msg = f"Output directory does not exist: {save_dir}"
|
error_msg = f"Output directory does not exist: {save_dir}"
|
||||||
@@ -501,7 +537,25 @@ class ISM:
|
|||||||
for control_file in save_dir.glob("*.!dev"):
|
for control_file in save_dir.glob("*.!dev"):
|
||||||
control_file.unlink(missing_ok=True)
|
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(
|
log_event(
|
||||||
"manifest_ism_download_complete",
|
"manifest_ism_download_complete",
|
||||||
@@ -512,13 +566,14 @@ class ISM:
|
|||||||
"track_type": track.__class__.__name__,
|
"track_type": track.__class__.__name__,
|
||||||
"save_dir": str(save_dir),
|
"save_dir": str(save_dir),
|
||||||
"save_dir_exists": save_dir.exists(),
|
"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]],
|
"segment_files": [f.name for f in segments_to_merge[:10]],
|
||||||
|
"rolling_merge": merger is not None,
|
||||||
"downloader": "requests",
|
"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}"
|
error_msg = f"No segment files found in output directory: {save_dir}"
|
||||||
all_contents = list(save_dir.iterdir()) if save_dir.exists() else []
|
all_contents = list(save_dir.iterdir()) if save_dir.exists() else []
|
||||||
log_event(
|
log_event(
|
||||||
@@ -535,40 +590,41 @@ class ISM:
|
|||||||
)
|
)
|
||||||
raise FileNotFoundError(error_msg)
|
raise FileNotFoundError(error_msg)
|
||||||
|
|
||||||
is_text_subtitle = (
|
if not merger:
|
||||||
not session_drm
|
is_text_subtitle = (
|
||||||
and isinstance(track, Subtitle)
|
not session_drm
|
||||||
and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML)
|
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:
|
progress(downloaded="Merging", completed=0, total=len(segments_to_merge))
|
||||||
first_segment = segments_to_merge[0].read_bytes() if segments_to_merge else None
|
with open(save_path, "wb") as f:
|
||||||
init_segment = ISM.init_segment(track, session_drm, first_segment)
|
first_segment = segments_to_merge[0].read_bytes() if segments_to_merge else None
|
||||||
if init_segment:
|
init_segment = ISM.init_segment(track, session_drm, first_segment)
|
||||||
f.write(init_segment)
|
if init_segment:
|
||||||
iv_size = (read_per_sample_iv_size(first_segment) if session_drm and first_segment else None) or 8
|
f.write(init_segment)
|
||||||
for index, segment_file in enumerate(segments_to_merge):
|
iv_size = (read_per_sample_iv_size(first_segment) if session_drm and first_segment else None) or 8
|
||||||
if is_text_subtitle:
|
for index, segment_file in enumerate(segments_to_merge):
|
||||||
# first segment was already read for the init synthesis, reuse it
|
if is_text_subtitle:
|
||||||
segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes()
|
# first segment was already read for the init synthesis, reuse it
|
||||||
segment_data = try_ensure_utf8(segment_data)
|
segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes()
|
||||||
segment_data = (
|
segment_data = try_ensure_utf8(segment_data)
|
||||||
segment_data.decode("utf8")
|
segment_data = (
|
||||||
.replace("‎", html.unescape("‎"))
|
segment_data.decode("utf8")
|
||||||
.replace("‏", html.unescape("‏"))
|
.replace("‎", html.unescape("‎"))
|
||||||
.encode("utf8")
|
.replace("‏", html.unescape("‏"))
|
||||||
)
|
.encode("utf8")
|
||||||
f.write(segment_data)
|
)
|
||||||
elif session_drm:
|
f.write(segment_data)
|
||||||
segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes()
|
elif session_drm:
|
||||||
f.write(piff_senc_to_cenc(segment_data, iv_size))
|
segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes()
|
||||||
elif index == 0 and first_segment:
|
f.write(piff_senc_to_cenc(segment_data, iv_size))
|
||||||
f.write(first_segment)
|
elif index == 0 and first_segment:
|
||||||
else:
|
f.write(first_segment)
|
||||||
with open(segment_file, "rb") as src:
|
else:
|
||||||
shutil.copyfileobj(src, f, 1024 * 1024)
|
with open(segment_file, "rb") as src:
|
||||||
segment_file.unlink()
|
shutil.copyfileobj(src, f, 1024 * 1024)
|
||||||
progress(advance=1)
|
segment_file.unlink()
|
||||||
|
progress(advance=1)
|
||||||
|
|
||||||
track.path = save_path
|
track.path = save_path
|
||||||
events.emit(events.Types.TRACK_DOWNLOADED, track=track)
|
events.emit(events.Types.TRACK_DOWNLOADED, track=track)
|
||||||
|
|||||||
@@ -710,7 +710,8 @@ class Service(metaclass=ABCMeta):
|
|||||||
Parameters:
|
Parameters:
|
||||||
track: The decrypted Track object.
|
track: The decrypted Track object.
|
||||||
drm: The DRM object it decrypted with.
|
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:
|
def on_track_repacked(self, track: AnyTrack) -> None:
|
||||||
|
|||||||
@@ -619,9 +619,7 @@ class Track:
|
|||||||
else:
|
else:
|
||||||
save_dir = save_path.parent
|
save_dir = save_path.parent
|
||||||
|
|
||||||
keep_segments = (
|
keep_segments = config.continue_downloads and self.descriptor != self.Descriptor.URL
|
||||||
config.continue_downloads and not config.merge_segments and self.descriptor != self.Descriptor.URL
|
|
||||||
)
|
|
||||||
|
|
||||||
def cleanup() -> None:
|
def cleanup() -> None:
|
||||||
save_path.unlink(missing_ok=True)
|
save_path.unlink(missing_ok=True)
|
||||||
|
|||||||
Reference in New Issue
Block a user