mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-09-16 06:02:19 +02:00
imSp4rky
89f09f8e84
Opt-in decrypt_segments runs mp4decrypt --fragments-info per segment on a thread pool as segments land (DASH, single-init HLS), then decrypts the init alone; hides the post-download decrypt step with byte-identical output.
110 lines
3.7 KiB
Python
110 lines
3.7 KiB
Python
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()
|