mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 12:05:07 +00:00
imSp4rky
9f054269db
kid_windows returns an empty map for an MP4 with no encrypted sample entry or protecting senc, so the key check judges, flags and copies nothing for it.
238 lines
9.5 KiB
Python
238 lines
9.5 KiB
Python
import struct
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
from uuid import UUID
|
|
|
|
import pytest
|
|
|
|
from unshackle.core.drm.verify import PIFF_TENC, PIFF_TFXD, KidMap, kid_windows, pick_windows
|
|
|
|
A, B, C = UUID(int=0xA), UUID(int=0xB), UUID(int=0xC)
|
|
|
|
|
|
def box(kind: bytes, *parts: bytes) -> bytes:
|
|
body = b"".join(parts)
|
|
return struct.pack(">I4s", 8 + len(body), kind) + body
|
|
|
|
|
|
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:
|
|
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)))
|
|
|
|
|
|
def enca(kid: UUID) -> bytes:
|
|
return box(b"enca", bytes(28), box(b"sinf", box(b"schi", full(b"tenc", b"\x00\x00\x01\x08" + kid.bytes))))
|
|
|
|
|
|
def seig(protected: bool, kid: UUID) -> bytes:
|
|
return bytes([0, 0, int(protected), 8]) + kid.bytes
|
|
|
|
|
|
def sgpd(*entries: bytes, version: int = 1) -> bytes:
|
|
head = b"seig"
|
|
if version >= 1:
|
|
head += struct.pack(">I", 20)
|
|
if version >= 2:
|
|
head += struct.pack(">I", 0)
|
|
return full(b"sgpd", head + struct.pack(">I", len(entries)) + b"".join(entries), version=version)
|
|
|
|
|
|
def sbgp(*runs: tuple[int, int]) -> bytes:
|
|
return full(b"sbgp", b"seig" + struct.pack(">I", len(runs)) + b"".join(struct.pack(">II", *run) for run in runs))
|
|
|
|
|
|
def trak(*sample_entries: bytes, track: int = 1, handler: bytes = b"vide", stbl: bytes = b"") -> bytes:
|
|
stsd = full(b"stsd", struct.pack(">I", len(sample_entries)) + b"".join(sample_entries))
|
|
hdlr = full(b"hdlr", bytes(4) + handler + bytes(12))
|
|
return box(
|
|
b"trak",
|
|
full(b"tkhd", bytes(8) + struct.pack(">I", track)),
|
|
box(b"mdia", full(b"mdhd", bytes(8) + struct.pack(">I", 1000)), hdlr, box(b"minf", box(b"stbl", stsd, stbl))),
|
|
)
|
|
|
|
|
|
def movie(*sample_entries: bytes, handler: bytes = b"vide", stbl: bytes = b"", others: bytes = b"") -> bytes:
|
|
moov = box(
|
|
b"moov",
|
|
trak(*sample_entries, handler=handler, stbl=stbl),
|
|
others,
|
|
box(b"mvex", full(b"trex", struct.pack(">III", 1, 1, 0))),
|
|
)
|
|
return box(b"ftyp", b"iso6") + moov
|
|
|
|
|
|
def clear_audio(others: bytes = b"") -> bytes:
|
|
return movie(box(b"mp4a", bytes(28)), handler=b"soun", others=others)
|
|
|
|
|
|
def moof(
|
|
entry: int,
|
|
time: int,
|
|
*groups: bytes,
|
|
samples: int = 48,
|
|
duration: Optional[int] = None,
|
|
tfxd: bool = False,
|
|
) -> bytes:
|
|
flags = 0x02 | (0x08 if duration is not None else 0)
|
|
tfhd = struct.pack(">II", 1, entry) + (struct.pack(">I", duration) if duration is not None else b"")
|
|
if tfxd:
|
|
timing = box(b"uuid", PIFF_TFXD, bytes([1, 0, 0, 0]) + struct.pack(">QQ", time, 0))
|
|
else:
|
|
timing = full(b"tfdt", struct.pack(">Q", time), version=1)
|
|
trun = full(b"trun", struct.pack(">I", samples))
|
|
return box(b"moof", box(b"traf", full(b"tfhd", tfhd, flags=flags), timing, trun, *groups)) + box(b"mdat", bytes(64))
|
|
|
|
|
|
def local(kid: UUID, protected: bool = True) -> tuple[bytes, bytes]:
|
|
return sgpd(seig(protected, kid)), sbgp((48, 0x10001))
|
|
|
|
|
|
def test_each_kid_maps_to_its_own_fragments(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(
|
|
movie(encv(A), encv(B))
|
|
+ moof(1, 5000) # lead under A
|
|
+ moof(2, 7000) # the rest under B
|
|
+ moof(2, 9000)
|
|
+ moof(1, 11000, *local(C)) # a seig group overrides the sample entry's KID
|
|
+ moof(1, 13000, *local(A, protected=False)) # a clear seig group names no KID
|
|
)
|
|
assert kid_windows(path) == KidMap({A: [(0.0, 2.0)], B: [(2.0, 6.0), (4.0, 6.0)], C: [(6.0, 8.0)]}, True)
|
|
|
|
|
|
def test_windows_run_through_a_kid_but_stop_at_another_or_after_the_longest() -> None:
|
|
long_run = [(t * 2.0, t * 2.0 + 2) for t in range(20)]
|
|
assert pick_windows(long_run, longest=10.0) == [(0.0, 10.0), (20.0, 30.0), (30.0, 40.0)]
|
|
assert pick_windows(long_run) == [(0.0, 4.0), (20.0, 24.0), (36.0, 40.0)]
|
|
# a gap is another KID or a clear stretch: no window crosses it, and the last one ends the last run
|
|
assert pick_windows([(0, 2), (2, 4), (10, 12), (12, 14), (14, 16)], longest=10.0) == [(0, 4), (10, 16)]
|
|
assert pick_windows([(0, 2), (2, 4), (10, 12), (12, 14), (14, 16)]) == [(0, 4), (10, 14.0), (12, 16)]
|
|
|
|
|
|
def test_fragments_longer_than_the_longest_and_empty_spans() -> None:
|
|
assert pick_windows([(0.0, 60.0)]) == [(0.0, 4.0)]
|
|
assert pick_windows([(0, 15), (15, 30), (30, 45), (45, 60)]) == [(0, 4), (30, 34), (45, 49)]
|
|
assert pick_windows([(0.0, 5.0), (5.0, 5.0)]) == [(0.0, 4.0)]
|
|
assert pick_windows([(5.0, 5.0)]) == []
|
|
|
|
|
|
def test_nested_box_bomb_and_oversized_sbgp(tmp_path: Path) -> None:
|
|
body = b""
|
|
for _ in range(5000):
|
|
body = box(b"trak", body)
|
|
(tmp_path / "bomb.mp4").write_bytes(box(b"moov", body))
|
|
assert kid_windows(tmp_path / "bomb.mp4") is None
|
|
huge = full(b"sbgp", b"seig" + struct.pack(">I", 0x00FFFFFF) + bytes(16))
|
|
(tmp_path / "sbgp.mp4").write_bytes(movie(encv(A)) + moof(1, 0, huge))
|
|
assert kid_windows(tmp_path / "sbgp.mp4") is None
|
|
|
|
|
|
def test_first_middle_and_last_fragment(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(encv(A)) + b"".join(moof(1, t * 1000) for t in range(5)))
|
|
assert kid_windows(path) == KidMap({A: [(0.0, 4.0), (2.0, 6.0), (3.0, 7.0)]}, True)
|
|
|
|
|
|
def test_last_fragment_ends_after_its_sample_durations(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(encv(A)) + moof(1, 0, samples=10, duration=250))
|
|
assert kid_windows(path) == KidMap({A: [(0.0, 2.5)]}, True)
|
|
|
|
|
|
def test_audio_track_is_not_video(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.m4a"
|
|
path.write_bytes(movie(enca(A), handler=b"soun") + moof(1, 0))
|
|
assert kid_windows(path) == KidMap({A: [(0.0, 3.0)]}, False)
|
|
|
|
|
|
def test_moov_seig_table_survives_another_grouping(tmp_path: Path) -> None:
|
|
roll = full(b"sgpd", b"roll" + struct.pack(">II", 2, 1) + b"\x00\x01", version=1)
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(encv(A), stbl=sgpd(seig(True, C)) + roll) + moof(1, 0, sbgp((48, 1))))
|
|
assert kid_windows(path) == KidMap({C: [(0.0, 3.0)]}, True)
|
|
|
|
|
|
def test_sgpd_version_2(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(encv(A)) + moof(1, 0, sgpd(seig(True, C), version=2), sbgp((48, 0x10001))))
|
|
assert kid_windows(path) == KidMap({C: [(0.0, 3.0)]}, True)
|
|
|
|
|
|
def test_sgpd_version_0_with_constant_iv(tmp_path: Path) -> None:
|
|
constant_iv = bytes([0, 0, 1, 0]) + B.bytes + b"\x08" + bytes(8)
|
|
groups = sgpd(constant_iv, seig(True, C), version=0)
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(encv(A)) + moof(1, 0, groups, sbgp((24, 0x10001), (24, 0x10002))))
|
|
assert kid_windows(path) == KidMap({B: [(0.0, 3.0)], C: [(0.0, 3.0)]}, True)
|
|
|
|
|
|
def test_samples_past_the_sbgp_runs_take_the_sample_entry_kid(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(encv(A)) + moof(1, 0, sgpd(seig(True, C)), sbgp((10, 0x10001))))
|
|
assert kid_windows(path) == KidMap({A: [(0.0, 3.0)], C: [(0.0, 3.0)]}, True)
|
|
|
|
|
|
def test_piff_tenc_and_tfxd(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.ismv"
|
|
path.write_bytes(movie(encv(A, piff=True)) + moof(1, 0, tfxd=True) + moof(1, 2000, tfxd=True))
|
|
assert kid_windows(path) == KidMap({A: [(0.0, 4.0), (2.0, 5.0)]}, True)
|
|
|
|
|
|
def test_oversized_box(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(encv(A)) + moof(1, 0) + struct.pack(">I4s", 0x7FFFFFFF, b"moof") + bytes(16))
|
|
assert kid_windows(path) is None
|
|
|
|
|
|
def test_not_a_fragmented_mp4(tmp_path: Path) -> None:
|
|
path = tmp_path / "video.ts"
|
|
path.write_bytes(b"\x47" * 188 * 4)
|
|
assert kid_windows(path) is None
|
|
|
|
|
|
def test_clear_mp4_maps_no_kid(tmp_path: Path) -> None:
|
|
path = tmp_path / "clear.m4a"
|
|
path.write_bytes(clear_audio() + moof(1, 0))
|
|
assert kid_windows(path) == KidMap({}, False)
|
|
|
|
|
|
def test_encrypted_entry_without_a_tenc_is_not_clear(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(box(b"encv", bytes(78), box(b"sinf", box(b"frma", b"avc1")))) + moof(1, 0))
|
|
assert kid_windows(path) is None
|
|
|
|
|
|
@pytest.mark.parametrize("kind", [b"enct", b"encs", b"drmi"])
|
|
def test_other_encrypted_entry_kinds_are_not_clear(tmp_path: Path, kind: bytes) -> None:
|
|
path = tmp_path / "enc.mp4"
|
|
path.write_bytes(movie(box(kind, bytes(16)), handler=b"text") + moof(1, 0))
|
|
assert kid_windows(path) is None
|
|
|
|
|
|
def test_a_clear_entry_with_a_protecting_senc_is_not_clear(tmp_path: Path) -> None:
|
|
path = tmp_path / "enc.m4a"
|
|
senc = full(b"senc", struct.pack(">I", 1) + bytes(8))
|
|
path.write_bytes(clear_audio() + moof(1, 0, senc))
|
|
assert kid_windows(path) is None
|
|
path.write_bytes(clear_audio() + moof(1, 0, full(b"senc", bytes(4))))
|
|
assert kid_windows(path) == KidMap({}, False)
|
|
|
|
|
|
def test_only_the_tracks_with_fragments_decide_clear(tmp_path: Path) -> None:
|
|
path = tmp_path / "clear.m4a"
|
|
path.write_bytes(clear_audio(others=trak(encv(A), track=2)) + moof(1, 0))
|
|
assert kid_windows(path) == KidMap({}, False)
|
|
path.write_bytes(clear_audio(others=trak(encv(A), track=2)))
|
|
assert kid_windows(path) is None
|
|
|
|
|
|
def test_a_clear_lead_beside_an_encrypted_entry_is_not_clear(tmp_path: Path) -> None:
|
|
path = tmp_path / "lead.mp4"
|
|
path.write_bytes(movie(box(b"avc1", bytes(78)), encv(A)) + moof(1, 0))
|
|
assert kid_windows(path) is None
|