unshackle/tests/core/test_export.py

"""Tests for ``dl.write_export``, the ``--export`` mediaexport 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
import logging
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Optional, cast
from uuid import UUID

import click
import pytest

from unshackle.commands.dl import dl
from unshackle.core.drm import drm_from_dict
from unshackle.core.drm.clearkey import ClearKey
from unshackle.core.drm.clearkey_cenc import ClearKeyCENC
from unshackle.core.import_service import ImportService
from unshackle.core.titles import Movie
from unshackle.core.tracks import Audio, Chapter, 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` and `log`.
    instance = dl.__new__(dl)
    instance.service = "EXAMPLE"
    instance.log = logging.getLogger("download")
    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 import_ctx() -> click.Context:
    # ImportService only touches ctx.parent.params (proxy flags) when building its session.
    return cast(click.Context, SimpleNamespace(parent=None, params={}))


def read_export(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf8"))


def entry(path: Path) -> dict:
    return read_export(path)["titles"][0]


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["kind"] == "mediaexport" and doc["version"] == 1
    assert doc["service"]["tag"] == "EXAMPLE"
    tinfo = doc["titles"][0]
    assert tinfo["title"] == "Example Movie" and tinfo["kind"] == "movie"
    assert {row["id"] for row in tinfo["tracks"]} == {"v1", "a1", "s1"}
    assert tinfo["x-unshackle"]["meta"]["name"] == "Example Movie"
    assert set(tinfo["x-unshackle"]["tracks"]) == {"v1"}
    assert "keys" not in tinfo and "drm" not in tinfo


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))

    tinfo = entry(export)
    assert "keys" not in tinfo and "drm" not in tinfo
    assert "v1" in tinfo["x-unshackle"]["tracks"]


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

    tinfo = entry(export)
    assert tinfo["drm"] == [{"system": "widevine", "pssh": "AAAA"}]
    assert tinfo["keys"] == {KID.hex: "aa" * 16}
    assert tinfo["x-unshackle"]["drm"] == [{"system": "Widevine", "pssh_b64": "AAAA"}]


def test_drm_free_export_roundtrips_through_import_service(tmp_path: Path) -> None:
    """A DRM-free export must rebuild via ImportService: titles, tracks (no DRM),
    empty key pool (resolve_server_keys no-op) and chapters."""
    export = tmp_path / "export.json"
    title = make_title()
    title.tracks.chapters.add(Chapter("00:00:10.000", "Intro"))
    runner = make_dl()
    for track in [*title.tracks.videos, *title.tracks.audio, *title.tracks.subtitles]:
        runner.write_export(export, title, track)

    svc = ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export))

    titles = list(svc.get_titles())
    assert len(titles) == 1
    movie = titles[0]
    assert isinstance(movie, Movie)
    assert movie.name == "Example Movie"
    assert movie.year == 2024

    tracks = svc.get_tracks(movie)
    assert {t.id for t in tracks} == {"v1", "a1", "s1"}
    assert all(not t.drm for t in tracks)

    assert svc.key_pool() == {}
    movie.tracks = tracks
    svc.resolve_server_keys(movie)
    assert all(not t.drm for t in movie.tracks)

    # Chapters.add auto-inserts a nameless 00:00:00 baseline chapter; it round-trips too.
    assert [c.name for c in svc.get_chapters(movie)] == [None, "Intro"]


def test_clearkey_cenc_exports_drm_and_keys(tmp_path: Path) -> None:
    """A licensed ClearKeyCENC exports its system dict and KID:KEY map, and the
    exported DRM dict plus keys rebuild a decrypt-ready instance via drm_from_dict."""
    export = tmp_path / "export.json"
    title = make_title()
    video = title.tracks.videos[0]
    drm = ClearKeyCENC(kids=[KID], laurl="https://license.example.test/ck", content_keys={KID: "cc" * 16})

    make_dl().write_export(export, title, video, drm)

    tinfo = entry(export)
    assert tinfo["drm"] == [{"system": "clearkey"}]
    assert tinfo["keys"] == {KID.hex: "cc" * 16}
    own = tinfo["x-unshackle"]["drm"]
    assert own == [{"system": "ClearKeyCENC", "kids": [KID.hex], "laurl": "https://license.example.test/ck"}]

    rebuilt = drm_from_dict({**own[0], "content_keys": tinfo["keys"]})
    assert isinstance(rebuilt, ClearKeyCENC)
    assert rebuilt.content_keys == {KID: "cc" * 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)

    tinfo = entry(export)
    assert tinfo["drm"] == [{"system": "widevine", "pssh": "AAAA"}]
    assert "keys" not in tinfo


def test_unidl_export_imports_without_track_dicts(tmp_path: Path) -> None:
    """A file from another tool has no x-unshackle block: titles, keys and chapters still rebuild."""
    export = tmp_path / "unidl.json"
    export.write_text(
        json.dumps(
            {
                "kind": "unidl-export",
                "version": 1,
                "service": "example",
                "titles": [
                    {
                        "save_name": "Show.S01E02",
                        "title": {
                            "id": "ep-2",
                            "kind": "episode",
                            "name": "Show",
                            "episode_name": "Two",
                            "season": 1,
                            "episode": 2,
                            "year": "2024",
                        },
                        "keys": [f"{KID.hex}:{'dd' * 16}"],
                        "manifest_url": "https://example.test/m.mpd",
                        "headers": {"User-Agent": "unidl-ua"},
                        "drm": {"system": "widevine", "pssh": "AAAA"},
                        "chapters": [{"start_ms": 5000, "title": "Recap"}],
                    }
                ],
            }
        ),
        encoding="utf8",
    )
    svc = ImportService(import_ctx(), "EXAMPLE", "ep-2", str(export))

    ep = next(iter(svc.get_titles()))
    assert (ep.title, ep.season, ep.number) == ("Show", 1, 2)
    assert svc.session.headers["User-Agent"] == "unidl-ua"
    assert svc.key_pool() == {KID: "dd" * 16}
    assert svc.exported_drm_system() == "widevine"
    assert [c.name for c in svc.get_chapters(ep)] == [None, "Recap"]
    assert svc.titles_data["ep-2"]["manifest_type"] == "DASH"


def crit_export(path: Path, *crit_ids: str) -> Path:
    """A mediaexport file with one usable title per id, except ``crit_ids``, which need ``segments``."""
    titles = [
        {
            "id": tid,
            "kind": "movie",
            "title": tid,
            "manifests": [{"url": f"https://example.test/{tid}.mpd", "type": "dash"}],
            **({"crit": ["segments"]} if tid in crit_ids else {}),
        }
        for tid in ("movie-1", "movie-2")
    ]
    doc = {"kind": "mediaexport", "version": 1, "service": {"tag": "EXAMPLE"}, "titles": titles}
    path.write_text(json.dumps(doc), encoding="utf8")
    return path


def test_import_skips_a_title_it_cannot_use(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
    """A title another tool marked as needing a capability unshackle lacks is named and skipped."""
    export = crit_export(tmp_path / "export.json", "movie-2")

    with caplog.at_level(logging.WARNING, logger="EXAMPLE"):
        svc = ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export))

    assert list(svc.titles_data) == ["movie-1"]
    assert "Skipping exported title 'movie-2'" in caplog.text


def test_import_with_no_usable_title_says_so(tmp_path: Path) -> None:
    """A file whose every title is refused stops with a reason, not an empty download."""
    export = crit_export(tmp_path / "export.json", "movie-1", "movie-2")

    with pytest.raises(click.ClickException, match="no title unshackle can use"):
        ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export))


@pytest.mark.parametrize(("system", "expected"), [("playready", "playready"), ("widevine", "widevine")])
def test_import_takes_its_drm_type_from_the_export(tmp_path: Path, system: str, expected: str) -> None:
    """The type stub routes the DRM objects, so it follows the exported system, not a widevine default."""
    export = crit_export(tmp_path / "export.json")
    doc = read_export(export)
    doc["titles"][0]["drm"] = [{"system": system, "pssh": "AAAA"}]
    export.write_text(json.dumps(doc), encoding="utf8")

    assert ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export)).server_cdm_type == expected


class KeyDRM:
    """A licensed DRM system that holds only ``content_keys``."""

    def __init__(self, key: str) -> None:
        self.content_keys = {KID: key}


def test_export_drops_session_headers(tmp_path: Path) -> None:
    """The file travels between people: the session cookie and bearer token must not go with it."""
    export = tmp_path / "export.json"
    title = make_title()
    title.tracks.manifest_url = "https://example.test/m.mpd"
    runner = make_dl()
    headers = {"User-Agent": "example-ua", "Cookie": "session=secret", "Authorization": "Bearer secret"}
    runner.export_service = cast(Any, SimpleNamespace(session=SimpleNamespace(headers=headers)))

    runner.write_export(export, title, title.tracks.videos[0])

    assert entry(export)["manifests"][0]["headers"] == {"User-Agent": "example-ua"}
    assert "secret" not in export.read_text(encoding="utf8")


def test_second_key_for_one_kid_keeps_the_first(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
    """Two titles that license one KID to different keys: the file keeps the first key and stays readable."""
    export = tmp_path / "export.json"
    first, second = make_title(), make_title()
    second.id = "movie-2"
    runner = make_dl()

    runner.write_export(export, first, first.tracks.videos[0], KeyDRM("aa" * 16))
    with caplog.at_level(logging.WARNING, logger="download"):
        runner.write_export(export, second, second.tracks.videos[0], KeyDRM("bb" * 16))

    titles = read_export(export)["titles"]
    assert titles[0]["keys"] == {KID.hex: "aa" * 16}
    assert "keys" not in titles[1]
    assert f"KID {KID.hex} already has a different key" in caplog.text


def test_uppercase_content_key_is_stored_lowercase(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
    """A licence server that sends uppercase hex still gives one key, not a conflict with itself."""
    export = tmp_path / "export.json"
    first, second = make_title(), make_title()
    second.id = "movie-2"
    runner = make_dl()

    runner.write_export(export, first, first.tracks.videos[0], KeyDRM("AA" * 16))
    with caplog.at_level(logging.WARNING, logger="download"):
        runner.write_export(export, second, second.tracks.videos[0], KeyDRM("AA" * 16))

    titles = read_export(export)["titles"]
    assert titles[0]["keys"] == {KID.hex: "aa" * 16}
    assert "different key" not in caplog.text


@pytest.mark.parametrize("key", ["zz" * 16, "aa" * 8, "aa" * 32])
def test_malformed_content_key_skips_the_key_not_the_download(
    tmp_path: Path, caplog: pytest.LogCaptureFixture, key: str
) -> None:
    """A key that is not 16 bytes of hex, as a new DRM system's tool might return, must not stop licensing."""
    export = tmp_path / "export.json"
    title = make_title()
    title.tracks.manifest_url = "https://example.test/m.mpd"
    runner = make_dl()

    with caplog.at_level(logging.WARNING, logger="download"):
        runner.write_export(export, title, title.tracks.videos[0], KeyDRM(key))

    assert "keys" not in entry(export)
    assert f"Not exporting the key for KID {KID.hex}" in caplog.text


AD_MPD_URL = "https://example.test/ads/m.mpd"


def ad_mpd(ad_id: str) -> str:
    """An MPD whose pre-roll carries only a marker that a service filter, not core, knows."""

    def period(pid: str, marker: str) -> str:
        prop = f'<SupplementalProperty schemeIdUri="urn:example:ad" value="{marker}"/>' if marker else ""
        return f"""<Period id="{pid}" duration="PT30S">{prop}
          <AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
            <Representation id="v1" bandwidth="5000000" width="1920" height="1080" codecs="avc1.640028">
              <SegmentTemplate media="{pid}/$Number$.m4s" initialization="{pid}/i.mp4" duration="10" startNumber="1"/>
            </Representation>
          </AdaptationSet></Period>"""

    return f"""<?xml version="1.0"?>
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT60S">
{period(ad_id, "PRE_ROLL")}{period("main", "")}</MPD>"""


def export_service_filtered_title(tmp_path: Path, manifest_url: Optional[str]) -> Path:
    """Export a DASH title as a service would, its own period_filter dropping the pre-roll, with no manifest snapshot.

    Without a snapshot the import re-fetches the manifest, the path ``content_periods`` serves.
    """
    from unshackle.core.manifests import DASH

    title = Movie(id_="movie-1", service=StubService, name="Example Movie", year=2024, language="en")
    parsed = DASH.from_text(ad_mpd("ad-1"), AD_MPD_URL).to_tracks(
        language="en",
        period_filter=lambda p: any(x.get("value") == "PRE_ROLL" for x in p.findall("SupplementalProperty")),
    )
    video = parsed.videos[0]
    assert video.data["dash"]["filtered_period_ids"] == ["ad-1"]
    title.tracks.add(video)
    title.tracks.manifest_url = manifest_url
    export = tmp_path / "export.json"
    make_dl().write_export(export, title, video)
    doc = read_export(export)
    doc["titles"][0]["x-unshackle"].pop("snapshot", None)
    export.write_text(json.dumps(doc), encoding="utf8")
    return export


@pytest.mark.parametrize("manifest_url", [AD_MPD_URL, None])
@pytest.mark.parametrize("ad_id_on_refetch", ["ad-1", "ad-9"])
def test_import_keeps_out_the_periods_the_service_filtered(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, manifest_url: Optional[str], ad_id_on_refetch: str
) -> None:
    """The import re-fetches the MPD without the service, so its filter must travel in the export.

    A re-fetch can name the ad period anew, so only the exported content periods are kept.
    """
    from unshackle.core.manifests import DASH

    export = export_service_filtered_title(tmp_path, manifest_url)
    assert [m.get("content_periods") for m in entry(export)["manifests"]] == [["main"]]
    monkeypatch.setattr(DASH, "from_url", lambda url, session=None, **kw: DASH.from_text(ad_mpd(ad_id_on_refetch), url))

    svc = ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export))
    video = svc.get_tracks(next(iter(svc.get_titles()))).videos[0]

    assert video.data["dash"]["filtered_period_ids"] == [ad_id_on_refetch]


def test_an_export_without_content_periods_imports_every_period(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """An export written before content_period_ids, or from a service with no filter, keeps every period."""
    from unshackle.core.manifests import DASH

    export = export_service_filtered_title(tmp_path, AD_MPD_URL)
    doc = read_export(export)
    for manifest in doc["titles"][0]["manifests"]:
        del manifest["content_periods"]
    export.write_text(json.dumps(doc), encoding="utf8")
    monkeypatch.setattr(DASH, "from_url", lambda url, session=None, **kw: DASH.from_text(ad_mpd("ad-1"), url))

    svc = ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export))
    video = svc.get_tracks(next(iter(svc.get_titles()))).videos[0]

    assert video.data["dash"]["filtered_period_ids"] == []


def test_import_keeps_every_period_when_no_exported_period_is_left(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
    """A re-fetch that names every period anew has nothing to match, so the import warns and filters nothing."""
    from unshackle.core.manifests import DASH

    export = export_service_filtered_title(tmp_path, AD_MPD_URL)
    renamed = ad_mpd("ad-1").replace('id="main"', 'id="main-2"')
    monkeypatch.setattr(DASH, "from_url", lambda url, session=None, **kw: DASH.from_text(renamed, url))

    svc = ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export))
    with caplog.at_level(logging.WARNING):
        video = svc.get_tracks(next(iter(svc.get_titles()))).videos[0]

    assert video.data["dash"]["filtered_period_ids"] == []
    assert "keeping every period" in caplog.text


def test_a_3d_video_stays_3d_through_export_and_import(tmp_path: Path) -> None:
    export = tmp_path / "export.json"
    title = make_title()
    title.tracks.videos[0].stereo = True
    runner = make_dl()
    for track in [*title.tracks.videos, *title.tracks.audio]:
        runner.write_export(export, title, track)
    assert [row.get("stereo") for row in entry(export)["tracks"]][0] is True
    svc = ImportService(import_ctx(), "EXAMPLE", "movie-1", str(export))
    movie = next(iter(svc.get_titles()))
    assert svc.get_tracks(movie).videos[0].stereo