mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 03:55:05 +00:00
imSp4rky
06f87dafa3
BlowfishStripe decrypts on the client after download, so a --remote client gets clear audio.
113 lines
3.9 KiB
Python
113 lines
3.9 KiB
Python
"""Tests for the Blowfish CBC stripe cipher (``unshackle.core.drm.blowfish_stripe``)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from collections.abc import Iterator
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import requests
|
|
from Cryptodome.Cipher import Blowfish
|
|
|
|
from unshackle.core.api.handlers import serialize_audio_track, serialize_drm
|
|
from unshackle.core.config import config
|
|
from unshackle.core.constants import DOWNLOAD_CANCELLED
|
|
from unshackle.core.drm import BlowfishStripe
|
|
from unshackle.core.remote_service import build_tracks, reconstruct_drm
|
|
from unshackle.core.tracks import Audio, Track
|
|
from unshackle.core.tracks import track as track_module
|
|
|
|
KEY = bytes.fromhex("00112233445566778899aabbccddeeff")
|
|
CHUNK = BlowfishStripe.CHUNK
|
|
|
|
|
|
def stripe_encrypt(data: bytes, key: bytes = KEY) -> bytes:
|
|
"""Reference encryption: every third full chunk, each with a fresh CBC run."""
|
|
out = bytearray(data)
|
|
for offset in range(0, len(data), CHUNK * 3):
|
|
chunk = data[offset : offset + CHUNK]
|
|
if len(chunk) == CHUNK:
|
|
out[offset : offset + CHUNK] = Blowfish.new(key, Blowfish.MODE_CBC, iv=BlowfishStripe.IV).encrypt(chunk)
|
|
return bytes(out)
|
|
|
|
|
|
def test_decrypt_round_trip(tmp_path: Path) -> None:
|
|
plain = os.urandom(CHUNK * 10 + 123)
|
|
path = tmp_path / "track.flac"
|
|
path.write_bytes(stripe_encrypt(plain))
|
|
assert path.read_bytes() != plain
|
|
|
|
BlowfishStripe(KEY).decrypt(path)
|
|
|
|
assert path.read_bytes() == plain
|
|
assert list(tmp_path.iterdir()) == [path]
|
|
|
|
|
|
def test_short_trailing_chunk_untouched(tmp_path: Path) -> None:
|
|
# the last stripe starts at CHUNK * 3 and holds only 100 bytes, so it stays clear
|
|
plain = os.urandom(CHUNK * 3 + 100)
|
|
encrypted = stripe_encrypt(plain)
|
|
assert encrypted[CHUNK * 3 :] == plain[CHUNK * 3 :]
|
|
path = tmp_path / "track.mp3"
|
|
path.write_bytes(encrypted)
|
|
|
|
BlowfishStripe(KEY).decrypt(path)
|
|
|
|
assert path.read_bytes() == plain
|
|
|
|
|
|
def test_serialize_round_trip_keeps_key() -> None:
|
|
serialized = serialize_drm([BlowfishStripe(KEY)])
|
|
assert serialized == [{"type": "blowfishstripe", "key": KEY.hex()}]
|
|
|
|
rebuilt = reconstruct_drm(serialized)
|
|
|
|
assert len(rebuilt) == 1
|
|
assert isinstance(rebuilt[0], BlowfishStripe)
|
|
assert rebuilt[0].key == KEY
|
|
|
|
|
|
def test_remote_tracks_keep_the_stripe_drm() -> None:
|
|
track = Audio(
|
|
id_="a1", url="https://example.test/a1", language="en", codec=Audio.Codec.FLAC, drm=[BlowfishStripe(KEY)]
|
|
)
|
|
|
|
rebuilt = build_tracks({"audio": [serialize_audio_track(track, include_url=True)]})
|
|
|
|
assert isinstance(rebuilt.audio[0].drm[0], BlowfishStripe)
|
|
assert rebuilt.audio[0].drm[0].key == KEY
|
|
|
|
|
|
def test_url_download_decrypts_once_without_a_kid_probe(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(config.directories, "temp", tmp_path)
|
|
monkeypatch.setattr(Track, "get_key_id", lambda *a, **k: pytest.fail("a stripe-cipher file has no KID to probe"))
|
|
monkeypatch.setattr(
|
|
track_module, "assert_fragments_decrypted", lambda path: pytest.fail("a stripe-cipher file has no fragments")
|
|
)
|
|
plain = os.urandom(CHUNK * 7 + 5)
|
|
|
|
def downloader(**kwargs: Any) -> Iterator[dict[str, Any]]:
|
|
(Path(kwargs["output_dir"]) / kwargs["filename"]).write_bytes(stripe_encrypt(plain))
|
|
yield {"file_downloaded": Path(kwargs["output_dir"]) / kwargs["filename"]}
|
|
|
|
track = Audio(
|
|
id_="a1",
|
|
url="https://example.test/a1.flac",
|
|
language="en",
|
|
codec=Audio.Codec.FLAC,
|
|
drm=[BlowfishStripe(KEY)],
|
|
downloader=downloader,
|
|
)
|
|
try:
|
|
track.download(
|
|
session=requests.Session(), prepare_drm=partial(lambda *a, **kw: None), progress=partial(lambda **kw: None)
|
|
)
|
|
finally:
|
|
DOWNLOAD_CANCELLED.clear()
|
|
|
|
assert track.path is not None and track.path.read_bytes() == plain
|
|
assert track.drm is None
|