From 57a5d4269a310e0e72c0bebb728c4ac66db5cf99 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Thu, 11 Jun 2026 09:36:58 -0600 Subject: [PATCH] fix(dl): export DRM-free, ClearKey, MonaLisa and server-CDM tracks write_export now tolerates drm=None; every downloaded track is written to the --export sidecar, not just Widevine/PlayReady-licensed ones. --- tests/core/test_export.py | 141 ++++++++++++++++++++++++++++++++++++++ unshackle/commands/dl.py | 29 ++++++-- 2 files changed, 163 insertions(+), 7 deletions(-) create mode 100644 tests/core/test_export.py diff --git a/tests/core/test_export.py b/tests/core/test_export.py new file mode 100644 index 0000000..8a49af5 --- /dev/null +++ b/tests/core/test_export.py @@ -0,0 +1,141 @@ +"""Tests for ``dl.write_export`` — the ``--export`` JSON sidecar. + +Regression: DRM-free tracks never pass through ``prepare_drm``, so ``write_export`` +must accept ``drm=None`` (and DRM systems without ``to_dict``/``content_keys`` such +as ClearKey) and still record the track/manifest/chapter/attachment info that +``unshackle import`` rebuilds a download from. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from uuid import UUID + +from unshackle.commands.dl import dl +from unshackle.core.drm.clearkey import ClearKey +from unshackle.core.titles import Movie +from unshackle.core.tracks import Audio, Subtitle, Video + +KID = UUID(hex="00000000000000000000000000000001") + + +class StubService: + """Stands in for the service class slot on Movie; never instantiated.""" + + +class StubDRM: + """Minimal licensed-DRM shape: ``to_dict`` plus filled ``content_keys``.""" + + def __init__(self) -> None: + self.content_keys = {KID: "aa" * 16} + + def to_dict(self) -> dict: + return {"system": "Widevine", "pssh_b64": "AAAA"} + + +def make_dl() -> dl: + # __new__ skips the CLI-driven __init__; write_export only needs `service`. + instance = dl.__new__(dl) + instance.service = "EXAMPLE" + return instance + + +def make_title() -> Movie: + title = Movie(id_="movie-1", service=StubService, name="Example Movie", year=2024, language="en") + title.tracks.add( + Video( + id_="v1", + url="https://example.test/v1.mp4", + language="en", + codec=Video.Codec.AVC, + range_=Video.Range.SDR, + width=1920, + height=1080, + bitrate=5_000_000, + ) + ) + title.tracks.add( + Audio( + id_="a1", + url="https://example.test/a1.mp4", + language="en", + codec=Audio.Codec.AAC, + bitrate=128_000, + ) + ) + title.tracks.add( + Subtitle( + id_="s1", + url="https://example.test/s1.vtt", + language="en", + codec=Subtitle.Codec.WebVTT, + ) + ) + return title + + +def read_export(path: Path) -> dict: + return json.loads(path.read_text(encoding="utf8")) + + +def test_drm_free_track_exports(tmp_path: Path) -> None: + """The reported bug: DRM-free downloads produced no usable export.""" + export = tmp_path / "export.json" + title = make_title() + video = title.tracks.videos[0] + + make_dl().write_export(export, title, video) + + doc = read_export(export) + assert doc["version"] == 2 + assert doc["service"] == "EXAMPLE" + tinfo = doc["titles"]["movie-1"] + assert tinfo["meta"]["name"] == "Example Movie" + assert set(tinfo["tracks"]) == {"v1", "a1", "s1"} + assert "keys" not in tinfo["tracks"]["v1"] + assert "drm" not in tinfo["tracks"]["v1"] + + +def test_clearkey_drm_exports_track_without_keys(tmp_path: Path) -> None: + """ClearKey has no to_dict/content_keys; the track info must still export.""" + export = tmp_path / "export.json" + title = make_title() + video = title.tracks.videos[0] + + make_dl().write_export(export, title, video, ClearKey(key="bb" * 16)) + + doc = read_export(export) + track = doc["titles"]["movie-1"]["tracks"]["v1"] + assert "keys" not in track + assert "drm" not in track + + +def test_post_download_write_keeps_licensed_keys(tmp_path: Path) -> None: + """The drm=None write after download must not clobber prepare_drm's DRM/keys.""" + export = tmp_path / "export.json" + title = make_title() + video = title.tracks.videos[0] + runner = make_dl() + + runner.write_export(export, title, video, StubDRM()) # prepare_drm + runner.write_export(export, title, video) # post-download hook + + track = read_export(export)["titles"]["movie-1"]["tracks"]["v1"] + assert track["drm"] == [{"system": "Widevine", "pssh_b64": "AAAA"}] + assert track["keys"] == {KID.hex: "aa" * 16} + + +def test_keyless_content_keys_writes_no_keys_entry(tmp_path: Path) -> None: + """A DRM object with empty content_keys must not create an empty keys map.""" + export = tmp_path / "export.json" + title = make_title() + video = title.tracks.videos[0] + drm = StubDRM() + drm.content_keys = {} + + make_dl().write_export(export, title, video, drm) + + track = read_export(export)["titles"]["movie-1"]["tracks"]["v1"] + assert track["drm"] == [{"system": "Widevine", "pssh_b64": "AAAA"}] + assert "keys" not in track diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index f670629..be6477a 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -2328,6 +2328,11 @@ class dl: max_workers=workers, progress=tracks_progress_callables[i], ) + # DRM-free and ClearKey tracks never reach prepare_drm, so export here. + # drm=None on purpose: licensed tracks already recorded their DRM/keys + # in prepare_drm, and write_export merges via setdefault. + if export_path: + self.write_export(export_path, title, track) def on_subtitle_skipped(track: Subtitle) -> None: lang = str(track.language) @@ -2933,12 +2938,14 @@ class dl: meta.update(type="movie", name=str(title)) return meta - def write_export(self, export: Path, title: Title_T, track: AnyTrack, drm: Any) -> None: + def write_export(self, export: Path, title: Title_T, track: AnyTrack, drm: Any = None) -> None: """Write a shareable v2 export usable by ``unshackle import``. Carries no session/cookies/dl-flags. Region (country code) is stored only when the export used ``--proxy``, as an import geofence. Each track records only the licensed - DRM system; content keys live once under the track's ``keys``. + DRM system; content keys live once under the track's ``keys``. ``drm`` may be None + (DRM-free track) or a DRM system without ``to_dict``/``content_keys`` (e.g. ClearKey) — + the track, manifest, chapter and attachment info is still exported. """ with self.EXPORT_LOCK: doc: dict[str, Any] = {} @@ -2971,11 +2978,14 @@ class dl: tracks_map[str(t.id)] = t.to_dict() track_data = tracks_map.setdefault(str(track.id), track.to_dict()) - if hasattr(drm, "to_dict"): - track_data["drm"] = [drm.to_dict()] - keys = track_data.setdefault("keys", {}) - for kid, key in drm.content_keys.items(): - keys[kid.hex] = key + if drm is not None: + if hasattr(drm, "to_dict"): + track_data["drm"] = [drm.to_dict()] + content_keys = getattr(drm, "content_keys", None) or {} + if content_keys: + keys = track_data.setdefault("keys", {}) + for kid, key in content_keys.items(): + keys[kid.hex] = key if "chapters" not in tinfo: tinfo["chapters"] = [ @@ -3057,6 +3067,8 @@ class dl: cek_tree.add(f"[text2]{kid.hex}:{key}") if not any(isinstance(x, Tree) and x.label == cek_tree.label for x in table.columns[0].cells): table.add_row(cek_tree) + if export: + self.write_export(export, title, track, drm) return track_quality = None @@ -3464,6 +3476,9 @@ class dl: table.add_row() table.add_row(cek_tree) + if export: + self.write_export(export, title, track, drm) + @staticmethod def get_cookie_path(service: str, profile: Optional[str]) -> Optional[Path]: """Get Service Cookie File Path for Profile."""