diff --git a/docs/reference/configuration/drm.md b/docs/reference/configuration/drm.md index b2322b72..4dd1f0f0 100644 --- a/docs/reference/configuration/drm.md +++ b/docs/reference/configuration/drm.md @@ -187,6 +187,12 @@ service tag case-insensitively and uses the `default` entry as the fallback. unshackle decrypts HLS AES-128 ClearKey in-process and ignores this setting. +mp4decrypt finds a content key by the KID in the first `sinf` box of the track. When that box +has no CENC `tenc` (for example, a PIFF `sinf` before the CENC one), unshackle also gives +mp4decrypt the content key of that KID by track ID. When mp4decrypt copies a track through +still encrypted, unshackle decrypts it again with shaka-packager and logs a warning. This +retry needs shaka-packager installed. + === "Global" ```yaml @@ -213,7 +219,8 @@ This works with `mp4decrypt` only. Set [`decryption`](#decryption) to `mp4decryp 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. +track keeps the whole-file decryption, and so does a track whose init segment mp4decrypt +cannot key by KID when no content key has that KID. !!! note "No resume" diff --git a/tests/conftest.py b/tests/conftest.py index c30deac5..e258ad36 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -91,3 +91,19 @@ def mux_args(monkeypatch: pytest.MonkeyPatch) -> Callable[..., list[str]]: return captured[0] return run + + +@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 diff --git a/tests/core/test_kid_windows.py b/tests/core/test_kid_windows.py index 25218c48..0a94a1e9 100644 --- a/tests/core/test_kid_windows.py +++ b/tests/core/test_kid_windows.py @@ -19,10 +19,14 @@ def full(kind: bytes, payload: bytes, version: int = 0, flags: int = 0) -> bytes return box(kind, bytes([version]) + flags.to_bytes(3, "big") + payload) -def encv(kid: UUID, piff: bool = False) -> bytes: +def sinf(kid: UUID, piff: bool = False) -> bytes: tenc_body = b"\x00\x00\x01\x08" + kid.bytes tenc = box(b"uuid", PIFF_TENC, bytes(4) + tenc_body) if piff else full(b"tenc", tenc_body) - return box(b"encv", bytes(78), box(b"sinf", box(b"frma", b"avc1"), box(b"schi", tenc))) + return box(b"sinf", box(b"frma", b"avc1"), box(b"schi", tenc)) + + +def encv(kid: UUID, piff: bool = False) -> bytes: + return box(b"encv", bytes(78), sinf(kid, piff)) def enca(kid: UUID) -> bytes: diff --git a/tests/core/test_mp4decrypt_track_key.py b/tests/core/test_mp4decrypt_track_key.py new file mode 100644 index 00000000..afc052e4 --- /dev/null +++ b/tests/core/test_mp4decrypt_track_key.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path +from typing import Any, Callable +from uuid import UUID + +import pytest + +from tests.core.test_kid_windows import A, B, box, clear_audio, encv, moof, movie, sinf +from tests.core.test_zero_kid_fallback import _playready, _widevine +from tests.test_segment_decrypt import read_argv +from unshackle.core import binaries +from unshackle.core.config import config +from unshackle.core.drm import ClearKeyCENC +from unshackle.core.drm.key_args import shaka_noise, track_key_args +from unshackle.core.drm.segment_decrypt import SegmentDecrypter, can_use +from unshackle.core.drm.verify import kidless_track_kids, mp4decrypt_or_shaka + +KEY = "b5df2ae3f63b8dda9049d6e41422e075" +OTHER_KEY = "11" * 16 +POSIX_STUBS = pytest.mark.skipif(sys.platform == "win32", reason="the decrypter stubs are shell scripts") + + +def two_sinf(kid: UUID, piff_first: bool) -> bytes: + """A sample entry with a PIFF and a CENC sinf for the same KID.""" + return box(b"encv", bytes(78), sinf(kid, piff_first), sinf(kid, not piff_first)) + + +def track_zero(entry: bytes) -> bytes: + """A movie whose tkhd carries track_ID 0.""" + data = movie(entry) + tkhd = data.index(b"tkhd") + return data[: tkhd + 16] + bytes(4) + data[tkhd + 20 :] + + +@pytest.mark.parametrize( + ("init", "kids"), + [ + pytest.param(movie(two_sinf(A, piff_first=True)), {1: A}, id="piff-then-cenc"), + pytest.param(movie(encv(A, piff=True)), {1: A}, id="piff-only"), + pytest.param(movie(two_sinf(A, piff_first=False)), {}, id="cenc-then-piff"), + pytest.param(movie(encv(A)), {}, id="cenc-only"), + pytest.param(clear_audio(), {}, id="clear"), + pytest.param(track_zero(two_sinf(A, piff_first=True)), {}, id="track-id-0"), + pytest.param(b"\x00\x00\x00\x01moov\x00\x00", {}, id="truncated-large-size"), + pytest.param(movie(two_sinf(A, piff_first=True))[:-20], {}, id="cut-short"), + ], +) +def test_kidless_track_kids(init: bytes, kids: dict[int, UUID]) -> None: + assert kidless_track_kids(init) == kids + + +def test_kidless_track_kids_survives_deep_nesting() -> None: + nested = b"" + for _ in range(3000): + nested = box(b"trak", nested) + assert kidless_track_kids(nested) == {} + + +def test_track_key_needs_the_track_kid() -> None: + assert track_key_args({A: KEY, B: OTHER_KEY}, {2: A}) == ["--key", f"2:{KEY}"] + assert track_key_args({B: OTHER_KEY}, {2: A}) == [] + + +DRMS: list[Any] = [ + pytest.param(lambda: _widevine(A), id="widevine"), + pytest.param(lambda: _playready(A), id="playready"), + pytest.param(lambda: ClearKeyCENC(kids=[A]), id="clearkey"), +] + + +def write_movie(path: Path, entry: bytes, fragment: bytes = b"") -> Path: + path.write_bytes(movie(entry) + fragment + box(b"mdat", bytes(16))) + return path + + +@POSIX_STUBS +@pytest.mark.parametrize("make", DRMS) +def test_whole_file_decrypt_keys_a_piff_first_track_by_id( + tmp_path: Path, stub_mp4decrypt: Path, make: Callable[[], Any] +) -> None: + path = write_movie(tmp_path / "track.mp4", two_sinf(A, piff_first=True)) + drm = make() + drm.content_keys[A] = KEY + + drm.decrypt_with_mp4decrypt(path) + + assert f"1:{KEY}" in read_argv(stub_mp4decrypt)[0] + + +@POSIX_STUBS +@pytest.mark.parametrize("make", DRMS) +def test_decrypt_falls_back_to_shaka_on_a_passthrough( + tmp_path: Path, stub_mp4decrypt: Path, monkeypatch: pytest.MonkeyPatch, make: Callable[[], Any] +) -> None: + path = write_movie(tmp_path / "track.mp4", two_sinf(A, piff_first=True), moof(1, 0)) + drm = make() + drm.content_keys[B] = OTHER_KEY + shaka: list[Path] = [] + monkeypatch.setattr(type(drm), "decrypt_with_shaka_packager", lambda self, p: shaka.append(p)) + monkeypatch.setattr(config, "decryption", "mp4decrypt", raising=False) + + drm.decrypt(path) + + assert shaka == [path] + + +@POSIX_STUBS +def test_segment_decrypter_keys_from_the_init_segment(tmp_path: Path, stub_mp4decrypt: Path) -> None: + drm = _widevine(A) + drm.content_keys[A] = KEY + + decrypter = SegmentDecrypter(drm, movie(two_sinf(A, piff_first=True)), tmp_path, workers=1) + decrypter.pool.shutdown() + + assert f"1:{KEY}" in decrypter.key_args + + +@pytest.mark.parametrize(("key_kid", "usable"), [(A, True), (B, False)], ids=["track-kid-keyed", "other-kid-only"]) +def test_segment_decrypt_needs_a_key_for_a_kidless_track( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, key_kid: UUID, usable: bool +) -> None: + monkeypatch.setattr(config, "decrypt_segments", True, raising=False) + monkeypatch.setattr(binaries, "Mp4decrypt", tmp_path / "mp4decrypt") + drm = _widevine(A) + drm.content_keys[key_kid] = KEY + + assert can_use(drm, "mp4decrypt", movie(two_sinf(A, piff_first=True))) is usable + assert can_use(drm, "mp4decrypt", movie(encv(A))) is True + + +class StubDRM: + """Records which decrypter ran; mp4decrypt writes ``mp4decrypt_output`` over the file.""" + + def __init__(self, mp4decrypt_output: bytes) -> None: + self.output = mp4decrypt_output + self.calls: list[str] = [] + + def decrypt_with_mp4decrypt(self, path: Path) -> None: + self.calls.append("mp4decrypt") + path.write_bytes(self.output) + + def decrypt_with_shaka_packager(self, path: Path) -> None: + self.calls.append("shaka") + + +ENCRYPTED = movie(encv(A)) + moof(1, 0) + + +@pytest.mark.parametrize( + ("output", "calls"), + [ + pytest.param(ENCRYPTED, ["mp4decrypt", "shaka"], id="copied-through"), + pytest.param(movie(encv(A)) + moof(1, 1000), ["mp4decrypt"], id="stale-encrypted-moov"), + pytest.param(movie(box(b"avc1", bytes(78))) + moof(1, 1000), ["mp4decrypt"], id="decrypted"), + pytest.param(b"not an mp4", ["mp4decrypt"], id="unreadable"), + ], +) +def test_mp4decrypt_or_shaka(tmp_path: Path, output: bytes, calls: list[str]) -> None: + drm = StubDRM(output) + path = tmp_path / "track.mp4" + path.write_bytes(ENCRYPTED) + + mp4decrypt_or_shaka(drm, path) + + assert drm.calls == calls + + +@pytest.fixture +def stub_packager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Callable[[int], None]: + """A shaka-packager that writes its output, logs an ERROR line and exits with the code given.""" + + def install(code: int) -> None: + script = tmp_path / "packager" + script.write_text( + "#!/usr/bin/env python3\n" + "import sys\n" + "output = sys.argv[1].split('output=')[1].split(',')[0]\n" + "open(output, 'wb').write(b'decrypted')\n" + "sys.stderr.write('[1103/153702:ERROR:box_definitions.cc(641)] Not implemented reached\\n')\n" + f"sys.exit({code})\n" + ) + script.chmod(0o755) + monkeypatch.setattr(binaries, "ShakaPackager", script) + + return install + + +@POSIX_STUBS +@pytest.mark.parametrize("make", DRMS) +def test_shaka_error_line_with_exit_0_is_a_decrypt( + tmp_path: Path, stub_packager: Callable[[int], None], make: Callable[[], Any] +) -> None: + stub_packager(0) + path = tmp_path / "track.mp4" + path.write_bytes(b"encrypted") + drm = make() + drm.content_keys[A] = KEY + + drm.decrypt_with_shaka_packager(path) + + assert path.read_bytes() == b"decrypted" + + +@POSIX_STUBS +@pytest.mark.parametrize("make", DRMS) +def test_shaka_non_zero_exit_raises( + tmp_path: Path, stub_packager: Callable[[int], None], make: Callable[[], Any] +) -> None: + stub_packager(2) + path = tmp_path / "track.mp4" + path.write_bytes(b"encrypted") + drm = make() + drm.content_keys[A] = KEY + + with pytest.raises(subprocess.CalledProcessError): + drm.decrypt_with_shaka_packager(path) + + +@pytest.mark.parametrize( + ("line", "noise"), + [ + pytest.param( + "I1103 21:41:56.557273 16212 packager_main.cc:591] Excluding 'pssh' boxes", True, id="glog-info" + ), + pytest.param("I0503 01:02:03.000000 1 demuxer.cc:105] Demuxer::Run()", True, id="glog-info-early-month"), + pytest.param("W1103 21:41:56.000000 1 x.cc:1] warning", True, id="glog-warning"), + pytest.param( + "[1103/214121:WARNING:track_run_iterator.cc(699)] Seeing non-zero composition offset", True, id="warning" + ), + pytest.param("[1103/153402:INFO:demuxer.cc(170)] Initialize Demuxer", True, id="info"), + pytest.param("[1103/153402:ERROR:box_definitions.cc(641)] Not implemented reached", False, id="error"), + pytest.param("E1103 21:41:56.000000 1 x.cc:1] failed", False, id="glog-error"), + pytest.param("Packaging completed successfully.", False, id="plain"), + ], +) +def test_shaka_noise(line: str, noise: bool) -> None: + assert shaka_noise(line) is noise diff --git a/tests/test_segment_decrypt.py b/tests/test_segment_decrypt.py index 3da022eb..09944aa6 100644 --- a/tests/test_segment_decrypt.py +++ b/tests/test_segment_decrypt.py @@ -23,26 +23,10 @@ class FakeDRM: def __init__(self) -> None: self.content_keys = {KID: KEY} - def mp4decrypt_key_args(self) -> list[str]: + def mp4decrypt_key_args(self, track_ids: object = ()) -> 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()] diff --git a/unshackle/core/drm/clearkey_cenc.py b/unshackle/core/drm/clearkey_cenc.py index c62a0726..77aaa6c0 100644 --- a/unshackle/core/drm/clearkey_cenc.py +++ b/unshackle/core/drm/clearkey_cenc.py @@ -16,6 +16,8 @@ from rich.text import Text from unshackle.core import binaries from unshackle.core.config import config from unshackle.core.console import console +from unshackle.core.drm import key_args as drm_key_args +from unshackle.core.drm.verify import kidless_track_kids, mp4decrypt_or_shaka, read_top_level_box from unshackle.core.utilities import log_event from unshackle.core.utils import run_proxy @@ -209,7 +211,7 @@ class ClearKeyCENC: decrypt_start = time.monotonic() if decrypter == "mp4decrypt": - self.decrypt_with_mp4decrypt(path) + mp4decrypt_or_shaka(self, path) else: self.decrypt_with_shaka_packager(path) @@ -224,12 +226,12 @@ 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.""" + def mp4decrypt_key_args(self, track_kids: Optional[dict[int, UUID]] = None) -> list[str]: + """Build the mp4decrypt --key arguments for every content key, and by track ID for each ``track_kids`` KID.""" key_args = [] for kid, key in self.content_keys.items(): key_args.extend(["--key", f"{kid.hex}:{key}"]) - return key_args + return key_args + drm_key_args.track_key_args(self.content_keys, track_kids) def decrypt_with_mp4decrypt(self, path: Path) -> None: """Decrypt using mp4decrypt""" @@ -238,7 +240,7 @@ class ClearKeyCENC: output_path = path.with_stem(f"{path.stem}_decrypted") - key_args = self.mp4decrypt_key_args() + key_args = self.mp4decrypt_key_args(kidless_track_kids(read_top_level_box(path, b"moov") or b"")) cmd = [ str(binaries.Mp4decrypt), @@ -301,7 +303,6 @@ class ClearKeyCENC: ) stream_skipped = False - had_error = False shaka_log_buffer = "" assert p.stderr is not None @@ -312,27 +313,19 @@ class ClearKeyCENC: if "Skip stream" in line: # file/segment was so small that it didn't have any actual data, ignore stream_skipped = True - if ":INFO:" in line: - continue - if "I0" in line or "W0" in line: - continue - if ":ERROR:" in line: - had_error = True - if "Insufficient bits in bitstream for given AVC profile" in line: - # this is a warning and is something we don't have to worry about + if drm_key_args.shaka_noise(line): continue shaka_log_buffer += f"{line.strip()}\n" - if shaka_log_buffer: - # wrap to console width - padding - '[ClearKey]: ' - shaka_log_buffer = "\n ".join( - textwrap.wrap(shaka_log_buffer.rstrip(), width=console.width - 22, initial_indent="") - ) - console.log(Text.from_ansi("\n[ClearKey]: " + shaka_log_buffer)) - p.wait() - if p.returncode != 0 or had_error: + if p.returncode != 0: + if shaka_log_buffer: + # wrap to console width - padding - '[ClearKey]: ' + shaka_log_buffer = "\n ".join( + textwrap.wrap(shaka_log_buffer.rstrip(), width=console.width - 22, initial_indent="") + ) + console.log(Text.from_ansi("\n[ClearKey]: " + shaka_log_buffer)) raise subprocess.CalledProcessError(p.returncode, [binaries.ShakaPackager, *arguments]) path.unlink() diff --git a/unshackle/core/drm/key_args.py b/unshackle/core/drm/key_args.py index 7b310e15..528e69a5 100644 --- a/unshackle/core/drm/key_args.py +++ b/unshackle/core/drm/key_args.py @@ -1,10 +1,12 @@ -"""Decrypter key arguments shared by the Widevine and PlayReady DRM systems.""" +"""Decrypter key arguments and output filtering shared by the CENC DRM systems.""" from __future__ import annotations -from typing import Any, Iterable, Optional +import re +from typing import Any, Iterable, Mapping, Optional ZERO_KID = "00" * 16 +SHAKA_NOISE = re.compile(r"^(?:[IWV]\d{4} |\[[^\]]*:(?:INFO|WARNING|VERBOSE\d*):)") def _hex(value: Any) -> str: @@ -42,15 +44,27 @@ def zero_kid_key(content_keys: dict[Any, Any], own_kids: Iterable[Any]) -> Optio return None -def mp4decrypt_key_args(content_keys: dict[Any, Any], own_kids: Iterable[Any]) -> list[str]: - """Build the mp4decrypt --key arguments: every content key, then the zero-KID fallback.""" +def track_key_args(content_keys: dict[Any, Any], track_kids: Optional[Mapping[int, Any]]) -> list[str]: + """Build the mp4decrypt --key arguments by track ID for each track whose own KID has a content key.""" + keys = {_hex(kid): _hex(key) for kid, key in content_keys.items()} + key_args: list[str] = [] + for track_id, kid in (track_kids or {}).items(): + if _hex(kid) in keys: + key_args.extend(["--key", f"{track_id}:{keys[_hex(kid)]}"]) + return key_args + + +def mp4decrypt_key_args( + content_keys: dict[Any, Any], own_kids: Iterable[Any], track_kids: Optional[Mapping[int, Any]] = None +) -> list[str]: + """Build the mp4decrypt --key arguments: every content key, the zero-KID fallback and the ``track_kids`` keys.""" key_args: list[str] = [] for kid, key in content_keys.items(): key_args.extend(["--key", f"{_hex(kid)}:{_hex(key)}"]) fallback = zero_kid_key(content_keys, own_kids) if fallback: key_args.extend(["--key", f"{ZERO_KID}:{fallback}"]) - return key_args + return key_args + track_key_args(content_keys, track_kids) def shaka_keys(content_keys: dict[Any, Any], own_kids: Iterable[Any]) -> str: @@ -60,3 +74,8 @@ def shaka_keys(content_keys: dict[Any, Any], own_kids: Iterable[Any]) -> str: if fallback: entries.append(f"label={len(entries)}:key_id={ZERO_KID}:key={fallback}") return ",".join(entries) + + +def shaka_noise(line: str) -> bool: + """True for a shaka-packager log line below ERROR, in its glog (I1103 ...) or bracket ([1103/...:WARNING:...]) form.""" + return bool(SHAKA_NOISE.match(line)) diff --git a/unshackle/core/drm/playready.py b/unshackle/core/drm/playready.py index 9fa2f08b..20c880ab 100644 --- a/unshackle/core/drm/playready.py +++ b/unshackle/core/drm/playready.py @@ -25,6 +25,7 @@ from unshackle.core.config import config from unshackle.core.console import console from unshackle.core.constants import AnyTrack from unshackle.core.drm import key_args as drm_key_args +from unshackle.core.drm.verify import kidless_track_kids, mp4decrypt_or_shaka, read_top_level_box from unshackle.core.session import AnySession, utf8_text from unshackle.core.utilities import call_licence, get_boxes, log_event from unshackle.core.utils import run_proxy @@ -571,7 +572,7 @@ class PlayReady: decrypt_start = time.monotonic() if decrypter == "mp4decrypt": - self.decrypt_with_mp4decrypt(path) + mp4decrypt_or_shaka(self, path) else: self.decrypt_with_shaka_packager(path) @@ -596,9 +597,9 @@ class PlayReady: return [self._kid] return self.kids - def mp4decrypt_key_args(self) -> list[str]: - """Build the mp4decrypt --key arguments for every content key and the zero-KID fallback.""" - return drm_key_args.mp4decrypt_key_args(self.content_keys, self.own_kids()) + def mp4decrypt_key_args(self, track_kids: Optional[dict[int, UUID]] = None) -> list[str]: + """Build the mp4decrypt --key arguments for every content key, the zero-KID and the track ID fallbacks.""" + return drm_key_args.mp4decrypt_key_args(self.content_keys, self.own_kids(), track_kids) def decrypt_with_mp4decrypt(self, path: Path) -> None: """Decrypt using mp4decrypt""" @@ -607,7 +608,7 @@ class PlayReady: output_path = path.with_stem(f"{path.stem}_decrypted") - key_args = self.mp4decrypt_key_args() + key_args = self.mp4decrypt_key_args(kidless_track_kids(read_top_level_box(path, b"moov") or b"")) cmd = [ str(binaries.Mp4decrypt), @@ -667,7 +668,6 @@ class PlayReady: ) stream_skipped = False - had_error = False shaka_log_buffer = "" assert p.stderr is not None for line in iter(p.stderr.readline, ""): @@ -676,25 +676,18 @@ class PlayReady: continue if "Skip stream" in line: stream_skipped = True - if ":INFO:" in line: - continue - if "I0" in line or "W0" in line: - continue - if ":ERROR:" in line: - had_error = True - if "Insufficient bits in bitstream for given AVC profile" in line: + if drm_key_args.shaka_noise(line): continue shaka_log_buffer += f"{line.strip()}\n" - if shaka_log_buffer: - shaka_log_buffer = "\n ".join( - textwrap.wrap(shaka_log_buffer.rstrip(), width=console.width - 22, initial_indent="") - ) - console.log(Text.from_ansi("\n[PlayReady]: " + shaka_log_buffer)) - p.wait() - if p.returncode != 0 or had_error: + if p.returncode != 0: + if shaka_log_buffer: + shaka_log_buffer = "\n ".join( + textwrap.wrap(shaka_log_buffer.rstrip(), width=console.width - 22, initial_indent="") + ) + console.log(Text.from_ansi("\n[PlayReady]: " + shaka_log_buffer)) raise subprocess.CalledProcessError(p.returncode, [binaries.ShakaPackager, *arguments]) path.unlink() diff --git a/unshackle/core/drm/segment_decrypt.py b/unshackle/core/drm/segment_decrypt.py index 93f315db..6d1471f2 100644 --- a/unshackle/core/drm/segment_decrypt.py +++ b/unshackle/core/drm/segment_decrypt.py @@ -12,6 +12,7 @@ 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.verify import kidless_track_kids from unshackle.core.drm.widevine import Widevine SEGMENT_DRM_T = Union[Widevine, PlayReady, ClearKeyCENC] @@ -19,13 +20,17 @@ 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.""" +def can_use(drm: Any, decryption: str, init: Optional[bytes] = None) -> bool: + """Whether unshackle can decrypt a track's segments one by one as they download. + + A track whose ``init`` mp4decrypt cannot key goes to the whole-file decrypt, which has the shaka fallback. + """ return ( isinstance(drm, (Widevine, PlayReady, ClearKeyCENC)) and str(decryption).lower() == "mp4decrypt" and config.decrypt_segments and bool(binaries.Mp4decrypt) + and not (init and set(kidless_track_kids(init).values()) - set(drm.content_keys)) ) @@ -44,7 +49,7 @@ class SegmentDecrypter: raise EnvironmentError("mp4decrypt executable not found but is required.") self.drm = drm - self.key_args = drm.mp4decrypt_key_args() + self.key_args = drm.mp4decrypt_key_args(kidless_track_kids(init_data)) work_dir.mkdir(parents=True, exist_ok=True) self.work_dir = Path(tempfile.mkdtemp(prefix="segdec-", dir=work_dir)) diff --git a/unshackle/core/drm/verify.py b/unshackle/core/drm/verify.py index 8e54388b..9dfb3ab7 100644 --- a/unshackle/core/drm/verify.py +++ b/unshackle/core/drm/verify.py @@ -1,9 +1,12 @@ +import logging import os import struct from pathlib import Path from typing import Any, BinaryIO, Callable, Iterator, NamedTuple, Optional from uuid import UUID +log = logging.getLogger(__name__) + DECRYPT_HOOK: Optional[Callable[..., None]] = None CONTAINERS = {b"moov", b"trak", b"mdia", b"minf", b"stbl", b"mvex", b"moof", b"traf", b"sinf", b"schi"} @@ -142,6 +145,94 @@ def has_encrypted_entry(moov: bytes) -> bool: ) +def read_top_level_box(path: Path, box_type: bytes) -> Optional[bytes]: + """Read a single top-level box from an MP4 file. + + This function walks only the box headers, so the wanted box is the only payload it reads into memory. + Returns None if the file has no such top-level box, or if a box header is truncated or declares an + impossible size. + """ + file_size = path.stat().st_size + with path.open("rb") as f: + start = 0 + while start + 8 <= file_size: + header = f.read(8) + if len(header) < 8: + return None + size = int.from_bytes(header[:4], "big") + if size == 1: + large = f.read(8) + if len(large) < 8: + return None + size = int.from_bytes(large, "big") + elif size == 0: + size = file_size - start + if size < 8: + return None + if header[4:8] == box_type: + f.seek(start) + return f.read(size) + start += size + f.seek(start) + return None + + +def has_encrypted_sample_entry(path: Path) -> bool: + """True if the MP4's moov still carries an encrypted sample entry (``enc*`` or ``drm*``). + + A faithful decrypt rewrites the entry back to the real codec 4CC through frma, so a + survivor means the decrypt tool skipped/failed the track. Only the moov is read, and its + stsd boxes are parsed, so byte runs in mdat or in metadata cannot match. Any read or + parse error -> False. + """ + try: + moov = read_top_level_box(path, b"moov") + return moov is not None and has_encrypted_entry(moov) + except (OSError, struct.error, RecursionError): + return False + + +def kidless_kid(data: bytes, kind: bytes, entry: int, end: int) -> Optional[UUID]: + """Return the KID of an encrypted sample entry whose first sinf holds no CENC tenc (a PIFF sinf, say).""" + for fields in ENTRY_FIELDS.get(kind, ()): + sinfs = [(body, box_end) for box, body, box_end in children(data, entry + fields, end) if box == b"sinf"] + if sinfs: + if any(box == b"tenc" for box, _, _ in walk(data, *sinfs[0])): + return None + return entry_kid(data, kind, entry, end) + return None + + +def kidless_track_kids(data: bytes) -> dict[int, UUID]: + """Map each track of a moov or init segment whose first sinf has no tenc for mp4decrypt to read to its KID.""" + kids: dict[int, UUID] = {} + track = 0 + try: + for kind, body, end in walk(data): + if kind == b"tkhd": + track = struct.unpack_from(">I", data, body + (20 if data[body] == 1 else 12))[0] + elif kind == b"stsd" and track: + for entry in children(data, body + 8, end): + kid = kidless_kid(data, *entry) + if kid: + kids[track] = kid + except (struct.error, IndexError, ValueError, RecursionError): + return {} + return kids + + +def mp4decrypt_or_shaka(drm: Any, path: Path) -> None: + """Decrypt ``path`` with mp4decrypt, and again with shaka-packager when mp4decrypt copied it through. + + mp4decrypt copies a track it has no content key for unchanged, so its first moof stays the same. + """ + first_moof = read_top_level_box(path, b"moof") + drm.decrypt_with_mp4decrypt(path) + if has_encrypted_sample_entry(path) and read_top_level_box(path, b"moof") == first_moof: + log.warning(f"mp4decrypt left {path.name} encrypted, decrypting it with shaka-packager") + drm.decrypt_with_shaka_packager(path) + + def senc_protects_samples(buf: bytes, body: int, end: int) -> bool: """True if a senc/PIFF-uuid box describes protected samples. diff --git a/unshackle/core/drm/widevine.py b/unshackle/core/drm/widevine.py index 3bf28b0d..6e8b3a93 100644 --- a/unshackle/core/drm/widevine.py +++ b/unshackle/core/drm/widevine.py @@ -21,6 +21,7 @@ from unshackle.core.config import config from unshackle.core.console import console from unshackle.core.constants import AnyTrack from unshackle.core.drm import key_args as drm_key_args +from unshackle.core.drm.verify import kidless_track_kids, mp4decrypt_or_shaka, read_top_level_box from unshackle.core.session import AnySession, utf8_text from unshackle.core.utilities import call_licence, get_boxes, log_event from unshackle.core.utils import run_proxy @@ -300,7 +301,7 @@ class Widevine: decrypt_start = time.monotonic() if decrypter == "mp4decrypt": - self.decrypt_with_mp4decrypt(path) + mp4decrypt_or_shaka(self, path) else: self.decrypt_with_shaka_packager(path) @@ -325,9 +326,9 @@ class Widevine: return [self._kid] return self.kids - def mp4decrypt_key_args(self) -> list[str]: - """Build the mp4decrypt --key arguments for every content key and the zero-KID fallback.""" - return drm_key_args.mp4decrypt_key_args(self.content_keys, self.own_kids()) + def mp4decrypt_key_args(self, track_kids: Optional[dict[int, UUID]] = None) -> list[str]: + """Build the mp4decrypt --key arguments for every content key, the zero-KID and the track ID fallbacks.""" + return drm_key_args.mp4decrypt_key_args(self.content_keys, self.own_kids(), track_kids) def decrypt_with_mp4decrypt(self, path: Path) -> None: """Decrypt using mp4decrypt""" @@ -336,7 +337,7 @@ class Widevine: output_path = path.with_stem(f"{path.stem}_decrypted") - key_args = self.mp4decrypt_key_args() + key_args = self.mp4decrypt_key_args(kidless_track_kids(read_top_level_box(path, b"moov") or b"")) cmd = [ str(binaries.Mp4decrypt), @@ -396,7 +397,6 @@ class Widevine: ) stream_skipped = False - had_error = False shaka_log_buffer = "" assert p.stderr is not None @@ -407,26 +407,19 @@ class Widevine: if "Skip stream" in line: # file/segment was so small that it didn't have any actual data, ignore stream_skipped = True - if ":INFO:" in line: - continue - if "I0" in line or "W0" in line: - continue - if ":ERROR:" in line: - had_error = True - if "Insufficient bits in bitstream for given AVC profile" in line: + if drm_key_args.shaka_noise(line): continue shaka_log_buffer += f"{line.strip()}\n" - if shaka_log_buffer: - # wrap to console width - padding - '[Widevine]: ' - shaka_log_buffer = "\n ".join( - textwrap.wrap(shaka_log_buffer.rstrip(), width=console.width - 22, initial_indent="") - ) - console.log(Text.from_ansi("\n[Widevine]: " + shaka_log_buffer)) - p.wait() - if p.returncode != 0 or had_error: + if p.returncode != 0: + if shaka_log_buffer: + # wrap to console width - padding - '[Widevine]: ' + shaka_log_buffer = "\n ".join( + textwrap.wrap(shaka_log_buffer.rstrip(), width=console.width - 22, initial_indent="") + ) + console.log(Text.from_ansi("\n[Widevine]: " + shaka_log_buffer)) raise subprocess.CalledProcessError(p.returncode, [binaries.ShakaPackager, *arguments]) path.unlink() diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 88e8584b..02dc425b 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -631,7 +631,7 @@ class DASH: collapse_span = None collapse_single_url = collapse_span is not None - use_segment_decrypt = not collapse_single_url and bool(init_data) and can_use(drm, config.decryption) + use_segment_decrypt = not collapse_single_url and bool(init_data) and can_use(drm, config.decryption, init_data) # subtitle segments are tiny, the post-download merge costs nothing there use_rolling_merge = not collapse_single_url and config.merge_segments and not isinstance(track, Subtitle) diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 221f3468..ccb535ab 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -867,6 +867,7 @@ class HLS: init_section = wanted_segments[0].init_section init_content, _ = HLS.fetch_init_section(session, init_section, 0) map_data = (init_section, init_content) + use_segment_decrypt = use_segment_decrypt and can_use(session_drm, config.decryption, init_content) if use_segment_decrypt: assert isinstance(session_drm, (Widevine, PlayReady, ClearKeyCENC)) segment_decrypter = SegmentDecrypter(session_drm, init_content, save_dir.parent, max_workers) diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 464aec49..9968c034 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -28,7 +28,12 @@ from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_ALL_DRM, DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, DownloadCancelled from unshackle.core.downloaders.requests import requests from unshackle.core.drm import DRM_T, BlowfishStripe, ClearKeyCENC, PlayReady, Widevine -from unshackle.core.drm.verify import decrypt_track, has_encrypted_entry, senc_protects_samples +from unshackle.core.drm.verify import ( + decrypt_track, + has_encrypted_sample_entry, + read_top_level_box, + senc_protects_samples, +) from unshackle.core.events import events from unshackle.core.session import MAX_BACKOFF, AnySession, RnetSession, utf8_text from unshackle.core.temp import restore_segments, stash_segments @@ -85,38 +90,6 @@ def direct_session(session: AnySession, proxy: Optional[str] = None) -> Session: return new -def read_top_level_box(path: Path, box_type: bytes) -> Optional[bytes]: - """Read a single top-level box from an MP4 file. - - This function walks only the box headers, so the wanted box is the only payload it reads into memory. - Returns None if the file has no such top-level box, or if a box header is truncated or declares an - impossible size. - """ - file_size = path.stat().st_size - with path.open("rb") as f: - start = 0 - while start + 8 <= file_size: - header = f.read(8) - if len(header) < 8: - return None - size = int.from_bytes(header[:4], "big") - if size == 1: - large = f.read(8) - if len(large) < 8: - return None - size = int.from_bytes(large, "big") - elif size == 0: - size = file_size - start - if size < 8: - return None - if header[4:8] == box_type: - f.seek(start) - return f.read(size) - start += size - f.seek(start) - return None - - def iter_top_level_boxes(path: Path) -> Iterable[tuple[bytes, int, int]]: """Yield (box_type, offset, size) for each top-level box, stopping at the first bad header.""" file_size = path.stat().st_size @@ -238,21 +211,6 @@ def has_mp4_only_audio_entry(path: Path) -> bool: return False -def has_encrypted_sample_entry(path: Path) -> bool: - """True if the MP4's moov still carries an encrypted sample entry (``enc*`` or ``drm*``). - - A faithful decrypt rewrites the entry back to the real codec 4CC through frma, so a - survivor means the decrypt tool skipped/failed the track. Only the moov is read, and its - stsd boxes are parsed, so byte runs in mdat or in metadata cannot match. Any read or - parse error -> False. - """ - try: - moov = read_top_level_box(path, b"moov") - return moov is not None and has_encrypted_entry(moov) - except (OSError, struct.error, RecursionError): - return False - - def moof_still_encrypted(moof: bytes) -> bool: # senc only: a decrypter detaches just the atom it consumed, so a PIFF uuid can # outlive a good decrypt. ISM arrives as senc via piff_senc_to_cenc.