diff --git a/docs/reference/configuration/drm.md b/docs/reference/configuration/drm.md index e01076d6..aef8a7fa 100644 --- a/docs/reference/configuration/drm.md +++ b/docs/reference/configuration/drm.md @@ -156,3 +156,27 @@ unshackle decrypts HLS AES-128 ClearKey in-process and ignores this setting. default: shaka EXAMPLE1: mp4decrypt ``` + +## `decrypt_segments` + +- **Type:** `bool`  ·  **Default:** `false` + +Decrypts each fragmented MP4 segment as it arrives, instead of decrypting the whole track +after the download ends. The decryption runs during the download, so the wait after the last +segment goes away. The output file is the same either way. + +This works with `mp4decrypt` only. Set [`decryption`](#decryption) to `mp4decrypt` for the +service (or as the default); with `shaka` the option has no effect, because shaka-packager +cannot decrypt a segment without its init segment. It applies to DASH tracks and to HLS +tracks that use one init segment, one `EXT-X-KEY` and no `EXT-X-DISCONTINUITY`. Every other +track keeps the whole-file decryption. + +!!! note "No resume" + + A track that decrypts segment by segment always downloads from the start. + [`continue_downloads`](download.md#continue_downloads) cannot reuse its segments, because a reused segment + is already decrypted. + +```yaml +decrypt_segments: true +``` diff --git a/tests/test_segment_decrypt.py b/tests/test_segment_decrypt.py new file mode 100644 index 00000000..3da022eb --- /dev/null +++ b/tests/test_segment_decrypt.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import json +import sys +from pathlib import Path +from uuid import UUID + +import pytest + +from unshackle.core import binaries +from unshackle.core.drm.segment_decrypt import SegmentDecrypter + +pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="the mp4decrypt stub is a shell script") + +ZERO_KID = "00" * 16 +KID = UUID(hex="11111111111111111111111111111111") +KEY = "22222222222222222222222222222222" + + +class FakeDRM: + """Stands in for Widevine/PlayReady: only the mp4decrypt key args matter here.""" + + def __init__(self) -> None: + self.content_keys = {KID: KEY} + + def mp4decrypt_key_args(self) -> list[str]: + return ["--key", f"{KID.hex}:{KEY}", "--key", f"{ZERO_KID}:{KEY}"] + + +@pytest.fixture +def stub_mp4decrypt(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """An mp4decrypt that copies its input to its output and records its argv.""" + argv_log = tmp_path / "argv.jsonl" + script = tmp_path / "mp4decrypt" + script.write_text( + "#!/usr/bin/env python3\n" + "import json, shutil, sys\n" + f"open({str(argv_log)!r}, 'a').write(json.dumps(sys.argv[1:]) + '\\n')\n" + "shutil.copyfile(sys.argv[-2], sys.argv[-1])\n" + ) + script.chmod(0o755) + monkeypatch.setattr(binaries, "Mp4decrypt", script) + return argv_log + + +def read_argv(argv_log: Path) -> list[list[str]]: + return [json.loads(line) for line in argv_log.read_text().splitlines()] + + +def test_segments_and_init_decrypt(tmp_path: Path, stub_mp4decrypt: Path) -> None: + segment_dir = tmp_path / "track_segments" + segment_dir.mkdir() + segments = [] + for i in range(3): + segment = segment_dir / f"{i:03}.mp4" + segment.write_bytes(f"segment-{i}".encode()) + segments.append(segment) + + decrypter = SegmentDecrypter(FakeDRM(), b"init-data", tmp_path / "work", workers=2) + work_dir = decrypter.work_dir + for segment in segments: + decrypter.submit(segment) + init_data = decrypter.finish() + + assert init_data == b"init-data" + # the stub copies, so a replaced segment keeps its bytes under its original name + assert sorted(p.name for p in segment_dir.iterdir()) == ["000.mp4", "001.mp4", "002.mp4"] + for i, segment in enumerate(segments): + assert segment.read_bytes() == f"segment-{i}".encode() + assert not work_dir.exists() + + calls = read_argv(stub_mp4decrypt) + assert len(calls) == 4 + + segment_calls = [call for call in calls if "--fragments-info" in call] + assert len(segment_calls) == 3 + for call in segment_calls: + assert call[call.index("--fragments-info") + 1].endswith("init.mp4") + assert call[-1].endswith(".dec") + assert Path(call[-2]).parent == segment_dir + + init_calls = [call for call in calls if "--fragments-info" not in call] + assert len(init_calls) == 1 + assert init_calls[0][-2].endswith("init.mp4") + + for call in calls: + assert "--key" in call + assert f"{ZERO_KID}:{KEY}" in call + assert f"{KID.hex}:{KEY}" in call + + +def test_failure_reraises_and_leaves_no_temp_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + script = tmp_path / "mp4decrypt" + script.write_text("#!/bin/sh\necho 'no key' >&2\nexit 1\n") + script.chmod(0o755) + monkeypatch.setattr(binaries, "Mp4decrypt", script) + + segment_dir = tmp_path / "track_segments" + segment_dir.mkdir() + segment = segment_dir / "000.mp4" + segment.write_bytes(b"segment-0") + + decrypter = SegmentDecrypter(FakeDRM(), b"init-data", tmp_path / "work", workers=1) + decrypter.submit(segment) + with pytest.raises(RuntimeError, match="no key"): + decrypter.finish() + + assert [p.name for p in segment_dir.iterdir()] == ["000.mp4"] + assert not decrypter.work_dir.exists() diff --git a/unshackle/core/config.py b/unshackle/core/config.py index d2e288e3..2238bfd9 100644 --- a/unshackle/core/config.py +++ b/unshackle/core/config.py @@ -109,6 +109,7 @@ class Config: else: self.decryption_map = {} self.decryption = decryption_cfg or "shaka" + self.decrypt_segments: bool = bool(kwargs.get("decrypt_segments", False)) self.theme: str = kwargs.get("theme") or "default" self.set_terminal_bg: bool = kwargs.get("set_terminal_bg", False) diff --git a/unshackle/core/drm/clearkey_cenc.py b/unshackle/core/drm/clearkey_cenc.py index 0a1dcad4..d851daa8 100644 --- a/unshackle/core/drm/clearkey_cenc.py +++ b/unshackle/core/drm/clearkey_cenc.py @@ -223,6 +223,13 @@ class ClearKeyCENC: output_size=path.stat().st_size if path.exists() else 0, ) + def mp4decrypt_key_args(self) -> list[str]: + """Build the mp4decrypt --key arguments for every content key.""" + key_args = [] + for kid, key in self.content_keys.items(): + key_args.extend(["--key", f"{kid.hex}:{key}"]) + return key_args + def decrypt_with_mp4decrypt(self, path: Path) -> None: """Decrypt using mp4decrypt""" if not binaries.Mp4decrypt: @@ -230,9 +237,7 @@ class ClearKeyCENC: output_path = path.with_stem(f"{path.stem}_decrypted") - key_args = [] - for kid, key in self.content_keys.items(): - key_args.extend(["--key", f"{kid.hex}:{key}"]) + key_args = self.mp4decrypt_key_args() cmd = [ str(binaries.Mp4decrypt), diff --git a/unshackle/core/drm/playready.py b/unshackle/core/drm/playready.py index 2893efd5..c40bb905 100644 --- a/unshackle/core/drm/playready.py +++ b/unshackle/core/drm/playready.py @@ -431,13 +431,8 @@ class PlayReady: output_size=path.stat().st_size if path.exists() else 0, ) - def decrypt_with_mp4decrypt(self, path: Path) -> None: - """Decrypt using mp4decrypt""" - if not binaries.Mp4decrypt: - raise EnvironmentError("mp4decrypt executable not found but is required.") - - output_path = path.with_stem(f"{path.stem}_decrypted") - + def mp4decrypt_key_args(self) -> list[str]: + """Build the mp4decrypt --key arguments for every content key.""" key_args = [] for kid, key in self.content_keys.items(): kid_hex = kid.hex if hasattr(kid, "hex") else str(kid).replace("-", "") @@ -452,6 +447,16 @@ class PlayReady: for key in self.content_keys.values(): key_hex = key if isinstance(key, str) else key.hex() key_args.extend(["--key", f"{zero_kid}:{key_hex}"]) + return key_args + + def decrypt_with_mp4decrypt(self, path: Path) -> None: + """Decrypt using mp4decrypt""" + if not binaries.Mp4decrypt: + raise EnvironmentError("mp4decrypt executable not found but is required.") + + output_path = path.with_stem(f"{path.stem}_decrypted") + + key_args = self.mp4decrypt_key_args() cmd = [ str(binaries.Mp4decrypt), diff --git a/unshackle/core/drm/segment_decrypt.py b/unshackle/core/drm/segment_decrypt.py new file mode 100644 index 00000000..575a6bec --- /dev/null +++ b/unshackle/core/drm/segment_decrypt.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import os +import shutil +import subprocess +import tempfile +from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path +from typing import Any, Optional, Union + +from unshackle.core import binaries +from unshackle.core.config import config +from unshackle.core.drm.clearkey_cenc import ClearKeyCENC +from unshackle.core.drm.playready import PlayReady +from unshackle.core.drm.widevine import Widevine + +SEGMENT_DRM_T = Union[Widevine, PlayReady, ClearKeyCENC] + +MAX_WORKERS = 16 + + +def can_use(drm: Any, decryption: str) -> bool: + """Whether unshackle can decrypt a track's segments one by one as they download.""" + return ( + isinstance(drm, (Widevine, PlayReady, ClearKeyCENC)) + and str(decryption).lower() == "mp4decrypt" + and config.decrypt_segments + and bool(binaries.Mp4decrypt) + ) + + +class SegmentDecrypter: + """Decrypts fMP4 segments as they arrive, via mp4decrypt --fragments-info, on a thread pool. + + mp4decrypt cannot read a bare `moof+mdat` segment on its own, so every call passes the + track's init segment through --fragments-info for the moov it needs. One last call + decrypts the init segment itself without that flag, which rewrites encv/enca back to the + real codec 4CC. Concatenating that init segment with the decrypted segments gives the same + bytes as decrypting the merged file in one pass. + """ + + def __init__(self, drm: SEGMENT_DRM_T, init_data: bytes, work_dir: Path, workers: Optional[int] = None) -> None: + if not binaries.Mp4decrypt: + raise EnvironmentError("mp4decrypt executable not found but is required.") + + self.drm = drm + self.key_args = drm.mp4decrypt_key_args() + + work_dir.mkdir(parents=True, exist_ok=True) + self.work_dir = Path(tempfile.mkdtemp(prefix="segdec-", dir=work_dir)) + self.init_path = self.work_dir / "init.mp4" + self.init_path.write_bytes(init_data) + + self.pool = ThreadPoolExecutor(max_workers=min(workers or os.cpu_count() or 4, MAX_WORKERS)) + self.futures: list[Future] = [] + + def submit(self, segment: Path) -> None: + """Queue one downloaded segment for decryption.""" + self.futures.append(self.pool.submit(self.decrypt_segment, segment)) + + def decrypt_segment(self, segment: Path) -> None: + output = segment.with_suffix(segment.suffix + ".dec") + try: + self.run(["--fragments-info", str(self.init_path), str(segment), str(output)]) + os.replace(output, segment) + finally: + output.unlink(missing_ok=True) + + def finish(self) -> bytes: + """Wait for every queued segment, then return the decrypted init segment bytes.""" + try: + for future in self.futures: + future.result() + decrypted_init = self.work_dir / "init_decrypted.mp4" + self.run([str(self.init_path), str(decrypted_init)]) + return decrypted_init.read_bytes() + finally: + self.close() + + def close(self) -> None: + """Drop any pending work and remove the temporary files. Safe to call twice.""" + for future in self.futures: + future.cancel() + self.futures.clear() + self.pool.shutdown(wait=True) + shutil.rmtree(self.work_dir, ignore_errors=True) + + def run(self, args: list[str]) -> None: + cmd = [str(binaries.Mp4decrypt), *self.key_args, *args] + try: + subprocess.run( + cmd, + check=True, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + ) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"mp4decrypt failed on {args[-2]}: {e.stderr.strip() if e.stderr else f'exit code {e.returncode}'}" + ) from e diff --git a/unshackle/core/drm/widevine.py b/unshackle/core/drm/widevine.py index 6d6690e5..d560705f 100644 --- a/unshackle/core/drm/widevine.py +++ b/unshackle/core/drm/widevine.py @@ -347,13 +347,8 @@ class Widevine: output_size=path.stat().st_size if path.exists() else 0, ) - def decrypt_with_mp4decrypt(self, path: Path) -> None: - """Decrypt using mp4decrypt""" - if not binaries.Mp4decrypt: - raise EnvironmentError("mp4decrypt executable not found but is required.") - - output_path = path.with_stem(f"{path.stem}_decrypted") - + def mp4decrypt_key_args(self) -> list[str]: + """Build the mp4decrypt --key arguments for every content key.""" key_args = [] for kid, key in self.content_keys.items(): kid_hex = kid.hex if hasattr(kid, "hex") else str(kid).replace("-", "") @@ -368,6 +363,16 @@ class Widevine: for key in self.content_keys.values(): key_hex = key if isinstance(key, str) else key.hex() key_args.extend(["--key", f"{zero_kid}:{key_hex}"]) + return key_args + + def decrypt_with_mp4decrypt(self, path: Path) -> None: + """Decrypt using mp4decrypt""" + if not binaries.Mp4decrypt: + raise EnvironmentError("mp4decrypt executable not found but is required.") + + output_path = path.with_stem(f"{path.stem}_decrypted") + + key_args = self.mp4decrypt_key_args() cmd = [ str(binaries.Mp4decrypt), diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 418d4d8a..f59a708c 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -25,6 +25,7 @@ from requests import Session from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine +from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use from unshackle.core.events import events from unshackle.core.session import RnetSession from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume @@ -469,13 +470,18 @@ class DASH: len(init_data) if init_data is not None else None, ) + use_segment_decrypt = not collapse_single_url and bool(init_data) and can_use(drm, config.decryption) + if not collapse_single_url: + # per-segment decryption rewrites each segment in place, so a segment kept from an + # earlier run would either be decrypted twice or never at all + can_resume = config.continue_downloads and not use_segment_decrypt # reuse a prior run's segments only when the fingerprint proves the segmentation unchanged digest = resume.fingerprint(track.url, 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) if collapse_single_url: @@ -524,15 +530,29 @@ class DASH: }, ) - 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) - else: - downloaded = status_update.get("downloaded") - if downloaded and downloaded.endswith("/s"): - status_update["downloaded"] = f"DASH {downloaded}" - progress(**status_update) + decrypter: Optional[SegmentDecrypter] = None + if use_segment_decrypt and init_data: + decrypter = SegmentDecrypter(drm, init_data, save_dir.parent, max_workers) + + try: + for status_update in downloader(**downloader_args): + file_downloaded = status_update.get("file_downloaded") + if file_downloaded: + if decrypter: + 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"DASH {downloaded}" + progress(**status_update) + + if decrypter: + init_data = decrypter.finish() + except Exception: + if decrypter: + decrypter.close() + raise if collapse_single_url: with open(save_path, "r+b") as collapsed: @@ -622,7 +642,8 @@ class DASH: if drm: progress(downloaded="Decrypting", completed=0, total=None) - drm.decrypt(save_path) + if not decrypter: + drm.decrypt(save_path) assert_fragments_decrypted(save_path) track.drm = None events.emit(events.Types.TRACK_DECRYPTED, track=track, drm=drm, segment=None) diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index dcd3b8db..0e2f42a2 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -31,6 +31,7 @@ from unshackle.core.cdm.detect import is_playready_cdm, is_widevine_cdm from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack from unshackle.core.drm import DRM_T, ClearKey, MonaLisa, PlayReady, Widevine +from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use from unshackle.core.events import events from unshackle.core.session import RnetResponse, RnetSession from unshackle.core.tracks import Audio, DownloadContext, Subtitle, Tracks, Video, resume @@ -518,6 +519,77 @@ class HLS: return None + @staticmethod + def resolve_segment_key( + segment_keys: list[Union[m3u8.model.SessionKey, m3u8.model.Key]], + cdm: object, + drm_preference: Optional[str] = None, + ) -> Optional[m3u8.Key]: + """Pick the EXT-X-KEY for a segment, preferring one the CDM can handle.""" + if cdm: + cdm_segment_keys = HLS.filter_keys_for_cdm(segment_keys, cdm, drm_preference) + if cdm_segment_keys: + return HLS.get_supported_key(cdm_segment_keys) + return HLS.get_supported_key(segment_keys) + + @staticmethod + def fetch_init_section( + session: Union[Session, RnetSession], + init_section: m3u8.model.InitializationSection, + range_offset: int, + ) -> tuple[bytes, int]: + """Download an EXT-X-MAP init segment. Returns its bytes and the next range offset.""" + if init_section.byterange: + init_byte_range = HLS.calculate_byte_range(init_section.byterange, range_offset) + range_offset = int(init_byte_range.split("-")[0]) + init_range_header = {"Range": f"bytes={init_byte_range}"} + else: + init_range_header = {} + + res = session.get( + url=urljoin(init_section.base_uri, init_section.uri), + headers=init_range_header, + ) + + if not isinstance(res, (requests.Response, RnetResponse)): + raise TypeError(f"Expected response to be requests.Response or rnet.Response, not {type(res)}") + res.raise_for_status() + return res.content, range_offset + + @staticmethod + def single_init_no_rotation( + wanted_segments: list[Any], + initial_key: Optional[m3u8.Key], + cdm: object, + drm_preference: Optional[str] = None, + ) -> bool: + """Whether every wanted segment shares one init segment and one EXT-X-KEY, with no discontinuity. + + Only then does the track merge into a single file with one init segment and one set of + content keys, which segment-by-segment decryption needs. + """ + if not wanted_segments: + return False + + first = wanted_segments[0] + if not first.init_section: + return False + + first_keys = list(getattr(first, "keys", None) or []) + for segment in wanted_segments[1:]: + if segment.discontinuity or segment.init_section != first.init_section: + return False + # the merge loop treats a segment without its own EXT-X-KEY as "key unchanged" + segment_keys = list(getattr(segment, "keys", None) or []) + if segment_keys and segment_keys != first_keys: + return False + + if not first_keys: + return True + # the merge loop re-resolves the content key per segment; a pick other than the one already + # licensed would re-key mid-track, which one pass of segment decryption cannot follow + return initial_key is not None and HLS.resolve_segment_key(first_keys, cdm, drm_preference) == initial_key + @staticmethod def download_track(track: AnyTrack, ctx: DownloadContext) -> None: session = ctx.ensure_session() @@ -633,6 +705,16 @@ class HLS: segment for segment in master.segments if callable(track.OnSegmentFilter) and track.OnSegmentFilter(segment) ] + # Downloaded segment files are named by post-filter index; map that back to the wanted + # 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] + + 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) + ) + total_segments = len(master.segments) - len(unwanted_segments) progress(total=total_segments) @@ -666,8 +748,12 @@ class HLS: segment_save_dir = save_dir / "segments" # no resume for AES-128/ClearKey: in-place decryption at merge time makes numbered files ambiguous - resumable_drm = (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 []) + # per-segment decryption rewrites each segment in place, so a segment kept from an + # earlier run would either be decrypted twice or never at all + resumable_drm = ( + not use_segment_decrypt + 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 [])) ) # media_sequence feeds AES IVs; total_segments pins the OnSegmentFilter verdict (shifts every index) digest = resume.fingerprint( @@ -708,15 +794,30 @@ class HLS: }, ) - 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) - else: - downloaded = status_update.get("downloaded") - if downloaded and downloaded.endswith("/s"): - status_update["downloaded"] = f"HLS {downloaded}" - progress(**status_update) + map_data: Optional[tuple[m3u8.model.InitializationSection, bytes]] = None + segment_decrypter: Optional[SegmentDecrypter] = None + if use_segment_decrypt: + init_section = wanted_segments[0].init_section + init_content, _ = HLS.fetch_init_section(session, init_section, 0) + map_data = (init_section, init_content) + segment_decrypter = SegmentDecrypter(session_drm, init_content, save_dir.parent, max_workers) + + 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: + if segment_decrypter: + segment_decrypter.close() + raise for control_file in segment_save_dir.glob("*.!dev"): control_file.unlink(missing_ok=True) @@ -732,10 +833,7 @@ class HLS: # First segment's Media Sequence Number - used as the AES-128 IV when EXT-X-KEY has # no explicit IV (RFC 8216 §5.2), where each segment's IV is its sequence number. media_sequence_start = getattr(master, "media_sequence", None) or 0 - # Downloaded segment files are named by post-filter index; map that back to the wanted - # 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] - map_data: Optional[tuple[m3u8.model.InitializationSection, bytes]] = None + decrypted_init: Optional[bytes] = None if session_drm: encryption_data: Optional[tuple[Optional[m3u8.Key], DRM_T]] = (initial_drm_key, session_drm) else: @@ -792,6 +890,7 @@ class HLS: Returns the decrypted path. """ + nonlocal map_data, decrypted_init drm = encryption_data[1] first_segment_i = next( (int(file.stem) for file in sorted(segment_save_dir.iterdir()) if file.stem.isdigit()), None @@ -817,8 +916,15 @@ class HLS: raise ValueError(f"Missing {range_len - len(files)} segment files for {segment_range}...") if isinstance(drm, (Widevine, PlayReady)): + if segment_decrypter: + if not map_data: + raise ValueError("Segment decryption needs an EXT-X-MAP init segment.") + if decrypted_init is None: + decrypted_init = segment_decrypter.finish() + map_data = (map_data[0], decrypted_init) merge(to=merged_path, via=files, delete=True, include_map_data=True) - drm.decrypt(merged_path) + if not segment_decrypter: + drm.decrypt(merged_path) assert_fragments_decrypted(merged_path) merged_path.rename(decrypted_path) else: @@ -896,37 +1002,12 @@ class HLS: map_data = None if segment.init_section and (not map_data or segment.init_section != map_data[0]): - if segment.init_section.byterange: - init_byte_range = HLS.calculate_byte_range(segment.init_section.byterange, range_offset) - range_offset = int(init_byte_range.split("-")[0]) - init_range_header = {"Range": f"bytes={init_byte_range}"} - else: - init_range_header = {} - - res = session.get( - url=urljoin(segment.init_section.base_uri, segment.init_section.uri), - headers=init_range_header, - ) - - if isinstance(res, requests.Response) or isinstance(res, RnetResponse): - res.raise_for_status() - init_content = res.content - else: - raise TypeError(f"Expected response to be requests.Response or rnet.Response, not {type(res)}") - + init_content, range_offset = HLS.fetch_init_section(session, segment.init_section, range_offset) map_data = (segment.init_section, init_content) segment_keys = getattr(segment, "keys", None) if segment_keys and segment not in unwanted_segments: - if cdm: - cdm_segment_keys = HLS.filter_keys_for_cdm(segment_keys, cdm, track.drm_preference) - key = ( - HLS.get_supported_key(cdm_segment_keys) - if cdm_segment_keys - else HLS.get_supported_key(segment_keys) - ) - else: - key = HLS.get_supported_key(segment_keys) + key = HLS.resolve_segment_key(segment_keys, cdm, track.drm_preference) if encryption_data and encryption_data[0] != key and i != 0 and segment not in unwanted_segments: decrypt(include_this_segment=False)