unshackle/tests/remote/unit/test_licence_declared_kwargs.py

"""The REST API licence paths must give a service only the arguments it declares.

`serve` reaches a service's licence function through three call sites that bypass the
`dl` command: server_cdm Widevine, server_cdm PlayReady, and the proxy-license forward.
Each one now filters with `declared_kwargs`, so a service that never asked for `pssh` or
`session_id` keeps working while a session-aware service finally receives `session_id`
on the server_cdm path.
"""

from __future__ import annotations

import base64
from types import SimpleNamespace
from typing import Any, Optional
from uuid import UUID

import pyplayready.system.pssh as pr_pssh_mod
import pytest
import pywidevine.pssh as wv_pssh_mod

import unshackle.core.cdm as cdm_mod
import unshackle.core.cdm.detect as detect_mod
import unshackle.core.drm as drm_mod
from unshackle.core.api import handlers
from unshackle.core.drm import PlayReady, Widevine
from unshackle.core.service import Service

pytestmark = pytest.mark.unit

KID = UUID("11111111111111111111111111111111")
KEY = bytes.fromhex("22222222222222222222222222222222")
PSSH_B64 = base64.b64encode(b"pssh").decode()
SESSION_ID = bytes(range(16))


class _Recorder:
    """Base for the fake services; every licence call lands in `calls`."""

    def __init__(self) -> None:
        self.calls: list[dict[str, Any]] = []

    def get_widevine_service_certificate(self, *, challenge: Any, title: Any, track: Any) -> None:
        return None


class ClassicService(_Recorder):
    """EXAMPLE: the pre-`session_id` signature every older service uses."""

    def get_widevine_license(self, *, challenge: bytes, title: Any, track: Any) -> bytes:
        self.calls.append({"challenge": challenge, "title": title, "track": track})
        return b"LIC"

    def get_playready_license(self, *, challenge: bytes, title: Any, track: Any) -> str:
        self.calls.append({"challenge": challenge, "title": title, "track": track})
        return "<License>ok</License>"


class SessionService(_Recorder):
    """DEMO: a session-aware service, the shape that needs the open CDM session id."""

    def get_widevine_license(
        self, *, challenge: bytes, title: Any, track: Any, session_id: Optional[bytes] = None
    ) -> bytes:
        self.calls.append({"challenge": challenge, "title": title, "track": track, "session_id": session_id})
        return b"LIC"

    def get_playready_license(
        self, *, challenge: bytes, title: Any, track: Any, session_id: Optional[bytes] = None
    ) -> str:
        self.calls.append({"challenge": challenge, "title": title, "track": track, "session_id": session_id})
        return "<License>ok</License>"


class VarKwService(_Recorder):
    """DEMO2: a `**kwargs` signature takes whatever the framework offers."""

    def get_widevine_license(self, *, challenge: bytes, **kwargs: Any) -> bytes:
        self.calls.append({"challenge": challenge, **kwargs})
        return b"LIC"

    def get_playready_license(self, *, challenge: bytes, **kwargs: Any) -> str:
        self.calls.append({"challenge": challenge, **kwargs})
        return "<License>ok</License>"


class StubWidevine(Widevine):
    """Real `Widevine.get_content_keys`, with the PSSH parsing skipped."""

    # pssh and kids are read-only properties on the real class; shadow them so a stub can set them
    pssh: Any = None
    kids: Any = None

    def __init__(self, pssh: Any = None, **kwargs: Any) -> None:
        self.pssh = pssh
        self.kids = [KID]
        self.content_keys: dict[UUID, str] = {}


class StubPlayReady(PlayReady):
    """Real `PlayReady.get_content_keys`, with the WRMHEADER parsing skipped."""

    # pssh, pssh_b64 and kids are read-only properties on the real class
    pssh: Any = None
    pssh_b64: Any = None
    kids: Any = None

    def __init__(self, pssh: Any = None, pssh_b64: Optional[str] = None, **kwargs: Any) -> None:
        self.pssh = SimpleNamespace(wrm_headers=[SimpleNamespace(dumps=lambda: "<WRMHEADER/>")])
        self.pssh_b64 = pssh_b64
        self.kids = [KID]
        self._extra_headers: list[Any] = []
        self._refused_headers: set[str] = set()
        self.content_keys: dict[UUID, str] = {}

    def extract_keys_from_cdm(self, cdm: Any, session_id: Any) -> dict[UUID, str]:
        return {KID: KEY.hex()}


class FakeWidevineCdm:
    service_certificate_challenge = b"cert-chal"

    def open(self) -> bytes:
        return SESSION_ID

    def get_license_challenge(self, session_id: bytes, pssh: Any) -> bytes:
        return b"CHAL"

    def parse_license(self, session_id: bytes, license_res: Any) -> None:
        return None

    def get_keys(self, session_id: bytes, type_: str) -> list[SimpleNamespace]:
        return [SimpleNamespace(kid=KID, key=KEY)]

    def close(self, session_id: bytes) -> None:
        return None


class FakePlayReadyCdm:
    def open(self) -> bytes:
        return SESSION_ID

    def get_license_challenge(self, session_id: bytes, wrm_header: Any) -> str:
        return "CHAL"

    def parse_license(self, session_id: bytes, license_str: str) -> None:
        return None

    def close(self, session_id: bytes) -> None:
        return None


@pytest.fixture
def widevine_env(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(wv_pssh_mod, "PSSH", lambda b64: SimpleNamespace(b64=b64))
    monkeypatch.setattr(drm_mod, "Widevine", StubWidevine)
    monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", lambda *a, **k: FakeWidevineCdm())
    monkeypatch.setattr(detect_mod, "is_widevine_cdm", lambda cdm: True)
    monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None)
    monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev")
    monkeypatch.setattr(handlers, "check_vaults", lambda kids, name, **_: None)
    monkeypatch.setattr(handlers, "cache_to_vaults", lambda keys, name, **_: None)
    monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False)


@pytest.fixture
def playready_env(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(pr_pssh_mod, "PSSH", lambda data: SimpleNamespace(data=data))
    monkeypatch.setattr(drm_mod, "PlayReady", StubPlayReady)
    monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", lambda *a, **k: FakePlayReadyCdm())
    monkeypatch.setattr(detect_mod, "is_playready_cdm", lambda cdm: True)
    monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None)
    monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev")
    monkeypatch.setattr(handlers, "cache_to_vaults", lambda keys, name, **_: None)
    monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False)


def _server_cdm(service: Any, drm_type: str) -> dict[str, str]:
    return handlers.handle_single_server_cdm(
        service=service,
        title=SimpleNamespace(),
        track=SimpleNamespace(),
        pssh_b64=PSSH_B64,
        drm_type=drm_type,
        request=None,
    )


def test_server_cdm_widevine_classic_signature_gets_three_arguments(widevine_env: None) -> None:
    service = ClassicService()
    keys = _server_cdm(service, "widevine")

    assert keys == {KID.hex: KEY.hex()}
    assert len(service.calls) == 1
    assert set(service.calls[0]) == {"challenge", "title", "track"}
    assert service.calls[0]["challenge"] == b"CHAL"


def test_server_cdm_widevine_session_service_receives_the_open_session_id(widevine_env: None) -> None:
    service = SessionService()
    _server_cdm(service, "widevine")

    assert set(service.calls[0]) == {"challenge", "title", "track", "session_id"}
    assert service.calls[0]["session_id"] == SESSION_ID


def test_server_cdm_widevine_var_kwargs_service_receives_pssh_and_session_id(widevine_env: None) -> None:
    service = VarKwService()
    _server_cdm(service, "widevine")

    assert set(service.calls[0]) == {"challenge", "title", "track", "pssh", "session_id"}
    assert service.calls[0]["session_id"] == SESSION_ID


def test_server_cdm_widevine_lambda_is_never_itself_filtered(widevine_env: None) -> None:
    # The `(**kw)` lambda would look like a VAR_KEYWORD service if declared_kwargs were ever
    # applied to it, and a classic service would then be handed `pssh` and `session_id`.
    service = ClassicService()
    _server_cdm(service, "widevine")

    assert "pssh" not in service.calls[0]
    assert "session_id" not in service.calls[0]


def test_server_cdm_playready_classic_signature_gets_three_arguments(playready_env: None) -> None:
    service = ClassicService()
    keys = _server_cdm(service, "playready")

    assert keys == {KID.hex: KEY.hex()}
    assert set(service.calls[0]) == {"challenge", "title", "track"}


def test_server_cdm_playready_offers_pssh_b64_and_session_id_not_pssh(playready_env: None) -> None:
    service = VarKwService()
    _server_cdm(service, "playready")

    assert set(service.calls[0]) == {"challenge", "title", "track", "pssh_b64", "session_id"}
    assert service.calls[0]["pssh_b64"] == PSSH_B64
    assert service.calls[0]["session_id"] == SESSION_ID


def test_server_cdm_playready_session_service_receives_session_id(playready_env: None) -> None:
    service = SessionService()
    _server_cdm(service, "playready")

    assert set(service.calls[0]) == {"challenge", "title", "track", "session_id"}
    assert service.calls[0]["session_id"] == SESSION_ID


@pytest.mark.parametrize(
    ("drm_type", "rejects", "last_call"),
    [
        ("widevine", ("pssh", "session_id"), {"challenge", "title", "track"}),
        ("playready", ("session_id",), {"challenge", "title", "track", "pssh_b64"}),
        ("playready", ("pssh_b64", "session_id"), {"challenge", "title", "track"}),
    ],
)
def test_server_cdm_unnarrowable_signature_retries_without_the_rejected_extras(
    drm_type: str, rejects: tuple[str, ...], last_call: set[str], request: pytest.FixtureRequest
) -> None:
    # declared_kwargs cannot narrow a signature it cannot read (a compiled service) so it offers
    # everything. Each TypeError names one rejected extra, and call_licence retries without it.
    request.getfixturevalue(f"{drm_type}_env")
    service = ClassicService()
    classic = getattr(service, f"get_{drm_type}_license")
    calls: list[dict[str, Any]] = []

    def compiled(**kw: Any) -> Any:
        calls.append(kw)
        rejected = next((k for k in rejects if k in kw), None)
        if rejected:
            raise TypeError(f"compiled() got an unexpected keyword argument '{rejected}'")
        return classic(challenge=kw["challenge"], title=None, track=None)

    setattr(service, f"get_{drm_type}_license", compiled)
    _server_cdm(service, drm_type)

    assert len(calls) == len(rejects) + 1
    assert set(calls[0]) == {
        "challenge",
        "title",
        "track",
        "session_id",
        "pssh" if drm_type == "widevine" else "pssh_b64",
    }
    assert set(calls[-1]) == last_call


@pytest.mark.parametrize("drm_type", ["widevine", "playready"])
def test_server_cdm_typeerror_inside_the_service_is_not_retried(drm_type: str, request: pytest.FixtureRequest) -> None:
    request.getfixturevalue(f"{drm_type}_env")
    service = SessionService()
    calls: list[dict[str, Any]] = []

    def broken(**kw: Any) -> bytes:
        calls.append(kw)
        raise TypeError("'NoneType' object is not subscriptable")

    setattr(service, f"get_{drm_type}_license", broken)
    with pytest.raises(TypeError, match="not subscriptable"):
        _server_cdm(service, drm_type)

    assert len(calls) == 1


def _proxy(service: Any, drm_type: str) -> Any:
    return handlers.handle_proxy_license(
        service=service,
        title=SimpleNamespace(),
        track=SimpleNamespace(),
        challenge_b64=base64.b64encode(b"CHAL").decode(),
        drm_type=drm_type,
    )


# A PlayReady challenge is SOAP XML, and pyplayready returns it as str on the local path too.
@pytest.mark.parametrize("drm_type, challenge", [("widevine", b"CHAL"), ("playready", "CHAL")])
def test_proxy_license_classic_signature(drm_type: str, challenge: Any) -> None:
    service = ClassicService()
    response = _proxy(service, drm_type)

    assert response.status == 200
    assert set(service.calls[0]) == {"challenge", "title", "track"}
    assert service.calls[0]["challenge"] == challenge


@pytest.mark.parametrize("drm_type", ["widevine", "playready"])
def test_proxy_license_var_kwargs_service_gets_no_session_id(drm_type: str) -> None:
    # The client holds the CDM, so there is no server-side session to name. A `**kwargs`
    # service must not see a session_id key at all.
    service = VarKwService()
    _proxy(service, drm_type)

    assert set(service.calls[0]) == {"challenge", "title", "track"}


@pytest.mark.parametrize("drm_type", ["widevine", "playready"])
def test_proxy_license_session_service_sees_its_own_default(drm_type: str) -> None:
    # A session-aware service is called without session_id, so its signature default applies.
    # A service that then uses session_id unconditionally fails here, before and after the
    # declared_kwargs change alike.
    service = SessionService()
    _proxy(service, drm_type)

    assert set(service.calls[0]) == {"challenge", "title", "track", "session_id"}
    assert service.calls[0]["session_id"] is None


def test_proxy_license_rejects_an_unknown_drm_type() -> None:
    from unshackle.core.api.errors import APIError

    with pytest.raises(APIError):
        _proxy(ClassicService(), "monalisa")


class _DelegatingService(Service):
    """A service with no PlayReady implementation: the base class delegates to Widevine."""

    def __init__(self) -> None:  # deliberately skips Service.__init__
        self.calls: list[dict[str, Any]] = []

    def get_titles(self) -> Any: ...  # pragma: no cover - abstract stub

    def get_tracks(self, title: Any) -> Any: ...  # pragma: no cover - abstract stub

    def get_chapters(self, title: Any) -> Any: ...  # pragma: no cover - abstract stub

    def get_widevine_license(self, *, challenge: bytes, title: Any, track: Any) -> bytes:
        self.calls.append({"challenge": challenge, "title": title, "track": track})
        return b"LIC"


class _DelegatingVarKwService(_DelegatingService):
    def get_widevine_license(self, *, challenge: bytes, **kwargs: Any) -> bytes:  # type: ignore[override]
        self.calls.append({"challenge": challenge, **kwargs})
        return b"LIC"


def test_base_playready_delegate_filters_to_the_widevine_signature() -> None:
    service = _DelegatingService()
    assert service.get_playready_license(challenge=b"C", title="T", track="TR") == b"LIC"
    assert set(service.calls[0]) == {"challenge", "title", "track"}


def test_base_playready_delegate_passes_everything_to_a_var_kwargs_service() -> None:
    service = _DelegatingVarKwService()
    service.get_playready_license(challenge=b"C", title="T", track="TR")
    assert set(service.calls[0]) == {"challenge", "title", "track"}


def test_base_playready_delegate_cannot_forward_session_id() -> None:
    # Service.get_playready_license declares (challenge, title, track) only, so the server_cdm
    # filter drops session_id before the delegate ever runs. A session-aware service that wants
    # PlayReady must implement get_playready_license itself.
    from unshackle.core.utilities import declared_kwargs

    offered = {"challenge": b"C", "title": "T", "track": "TR", "session_id": SESSION_ID}
    assert "session_id" not in declared_kwargs(Service.get_playready_license, offered)