feat(remote): decrypt service-owned cbcs samples on the server

The client downloads the file and sends binary sample frames to a new samples route.

docs/dev/rest-api/endpoints.md

@@ -933,6 +933,7 @@ GET  /api/session/{id}/titles   → list titles
 POST /api/session/{id}/tracks   → tracks + chapters (with URLs) for one title
 POST /api/session/{id}/segments → resolve segment URLs / DRM / headers per track
 POST /api/session/{id}/segment_filter → HLS segments the service drops (ads, bumpers)
+POST /api/session/{id}/samples  → decrypt cbcs samples of a track on the server
 POST /api/session/{id}/license  → proxy or server-CDM DRM licensing
 POST /api/session/{id}/keys/bad → report a server-vault content key that did not decrypt
 DELETE /api/session/{id}        → tear down, harvest updated cache
@@ -1140,6 +1141,24 @@ Get the HLS segments a service drops (ads, bumpers, dub cards). The server runs
 | `404` | `TRACK_NOT_FOUND` | An unknown track ID was requested. |
 | `500` | `SERVICE_ERROR` | The playlist request or the filter failed. |
 
+### `POST /api/session/{session_id}/samples`
+
+Decrypt the samples of a track that the service decrypts itself after the download. The tracks response lists those track IDs in `sample_decrypt_tracks`. The client downloads the file, then sends each `stsd` sample entry's samples here, and writes the plaintexts back in place. The query string names the track and the 1-based entry: `?track_id=a-1&entry=2`.
+
+- An empty body asks how the entry decrypts. The JSON response is `{ "key": "<hex>" }`, which the client then uses itself, or `{ "decrypter": true }`.
+- An `application/octet-stream` body is a sample frame: a big-endian `uint32` count, one `uint32` length per sample, then the samples. Each sample is a ciphertext cut to whole AES blocks. The response is a frame of the same shape with the plaintexts in the same order.
+
+The client sends at most 8 MiB of samples per request, and the route accepts a body of up to 16 MiB. Binary frames keep the upload the size of the samples, and the server does not compress them. The route needs the `server_cdm` grant for the service, because the decrypter can use a process or an account of the server.
+
+| Status | Error code | Meaning |
+| --- | --- | --- |
+| `200` | - | The content key, `decrypter`, or the plaintext frame. |
+| `400` | `INVALID_INPUT` | Missing `track_id`, a bad `entry`, a malformed frame, or a track the service does not decrypt. |
+| `413` | - | The body is larger than 16 MiB. |
+| `403` | `FORBIDDEN` | The API key has no `server_cdm` grant for the service. |
+| `404` | `TRACK_NOT_FOUND` | An unknown track ID was requested. |
+| `500` | `SERVICE_ERROR` | The service's decrypter failed. |
+
 ### `POST /api/session/{session_id}/license`
 
 Get the content keys for the DRM. The `mode` field selects one of two modes.

docs/dev/rest-api/index.md

@@ -205,6 +205,7 @@ request and response in full.
 | `POST` | `/api/session/{session_id}/tracks` | Tracks (with URLs) for a chosen title. |
 | `POST` | `/api/session/{session_id}/segments` | Per-track segment/download details. |
 | `POST` | `/api/session/{session_id}/segment_filter` | Server-side HLS segment filter results for one track. |
+| `POST` | `/api/session/{session_id}/samples` | Server-side cbcs sample decryption for one track. |
 | `POST` | `/api/session/{session_id}/license` | Proxy or server-side DRM licensing. |
 | `POST` | `/api/session/{session_id}/keys/bad` | Report a server-vault content key that did not decrypt. |
 | `GET` | `/api/session/{session_id}/logs` | Drain the service's server-side log output. |

docs/dev/rest-api/remote-sessions.md

@@ -410,6 +410,7 @@ The server mounts all these routes, even in `--remote-only` mode. Paths use the
 | `POST` | `/api/session/{id}/tracks` | Resolve tracks, manifests, chapters, headers/cookies for a title |
 | `POST` | `/api/session/{id}/segments` | Resolve per-segment/track download descriptors |
 | `POST` | `/api/session/{id}/segment_filter` | Unwanted HLS segment URIs for one track (ads, bumpers) |
+| `POST` | `/api/session/{id}/samples` | Decrypt the cbcs samples of a track the service decrypts itself |
 | `POST` | `/api/session/{id}/license` | DRM licensing (proxy or server CDM) |
 | `POST` | `/api/session/{id}/keys/bad` | Report a server-vault content key that did not decrypt |
 | `GET` | `/api/session/{id}/logs` | Drain the service's server-side log output |

docs/dev/service-architecture.md

@@ -418,14 +418,20 @@ classes drive the CDM handshake and call the decryption backend. Your hook is on
 
 Some tracks cannot use the framework's decryption: a single cbcs file whose sample entries share
 a zero default KID, or a sample entry whose content key no CDM can license and only a separate
-process can use. For these, call `unshackle.core.drm.cbcs.decrypt_file(track.path, key_for)` from
-`on_track_downloaded`. Set no `drm` on the track, so the framework does not try to license it.
+process can use. For these, set no `drm` on the track and override two methods:
 
-`key_for(track, entry)` gets the track layout and one protected `stsd` entry. It returns the
-16-byte content key, or a `cbcs.SampleDecrypter`. A `SampleDecrypter` gets the protected bytes of
-each sample of one fragment, cut to whole AES blocks, and returns plaintexts of the same lengths.
-It handles full-sample cbcs only, not pattern encryption. `decrypt_file` writes the plaintext
-back in place and removes the protection boxes, so the muxer reads a clear file.
+- `decrypts_samples(track)` returns `True` for such a track.
+- `get_sample_decrypter(track, entry)` returns, for the 1-based `stsd` entry, the 16-byte
+  content key or a `cbcs.SampleDecrypter`. A `SampleDecrypter` gets the protected bytes of each sample of
+  one fragment, cut to whole AES blocks, and returns plaintexts of the same lengths. It handles
+  full-sample cbcs only, not pattern encryption.
+
+When the download of the track ends, unshackle calls `unshackle.core.drm.cbcs.decrypt_file`, and
+the progress row shows "Decrypting". Your `on_track_downloaded` runs before this and gets the
+encrypted file. `decrypt_file` writes the plaintext back in place and removes the
+protection boxes, so the muxer reads a clear file. In a remote session `get_sample_decrypter`
+runs on the server: the client sends the samples through the `samples` route, so a decrypter
+that needs a process on the server also works for a remote client.
 
 ### How unshackle chooses a CDM and finds keys
 

unshackle/commands/dl.py

@@ -55,7 +55,7 @@ from unshackle.core.constants import (
 )
 from unshackle.core.credential import Credential
 from unshackle.core.downloaders import default_max_workers, format_speed, parse_speed_limit, set_speed_limit
-from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, own_kids, verify
+from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, cbcs, own_kids, verify
 from unshackle.core.events import events
 from unshackle.core.export_name import ExportNamer, move_export, plural, season_episodes, title_label
 from unshackle.core.providers.anilist import parse_anilist_ref
@@ -3156,6 +3156,19 @@ class dl:
                             download_processes=download_processes,
                             progress=tracks_progress_callables[i],
                         )
+                        if (
+                            track.path
+                            and service.decrypts_samples(track)
+                            and not DOWNLOAD_CANCELLED.is_set()
+                            and not DOWNLOAD_LICENCE_ONLY.is_set()
+                        ):
+                            progress = tracks_progress_callables[i]
+                            progress(downloaded="Decrypting", completed=0, total=None)
+                            cbcs.decrypt_file(
+                                track.path, lambda _l, entry: service.get_sample_decrypter(track, entry.index)
+                            )
+                            assert_fragments_decrypted(track.path)
+                            progress(downloaded="Decrypted", completed=100, total=100)
                         # DRM-free and HLS-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.

unshackle/core/api/handlers.py

@@ -28,6 +28,7 @@ from unshackle.core.cacher import Cacher
 from unshackle.core.cdm.detect import cdm_type_stub
 from unshackle.core.config import config
 from unshackle.core.constants import AUDIO_CODEC_MAP, DYNAMIC_RANGE_MAP, VIDEO_CODEC_MAP
+from unshackle.core.drm import cbcs
 from unshackle.core.providers.anilist import parse_anilist_ref
 from unshackle.core.proxies.resolve import initialize_proxy_providers, resolve_proxy
 from unshackle.core.services import Services
@@ -3834,6 +3835,9 @@ async def session_tracks_handler(
                 "session_cookies": session_cookies,
                 "server_cdm": session.server_cdm,
                 "server_cdm_type": server_cdm_type,
+                "sample_decrypt_tracks": [
+                    str(t.id) for t in [*video_tracks, *audio_tracks] if service_instance.decrypts_samples(t)
+                ],
             }
         )
 
@@ -4826,6 +4830,59 @@ async def session_segment_filter_handler(
     return web.json_response({"unwanted": unwanted})
 
 
+async def session_samples_handler(
+    session_id: str, track_id: str, entry: str, frame: bytes, request: Optional[web.Request] = None
+) -> web.Response:
+    """Run the service's sample decrypter for one `stsd` entry of a track that `decrypts_samples` names.
+
+    An empty `frame` asks how the entry decrypts: its content key, or `decrypter`. A `cbcs.pack_samples`
+    frame of block-aligned ciphertexts gets a frame of their plaintexts back. The decrypter can use a
+    process or an account of the server, so the `server_cdm` grant gates it.
+    """
+    session = await get_validated_session(session_id, request)
+    require_authenticated(session)
+
+    index = int(entry) if entry.isascii() and entry.isdigit() else 0
+    if not track_id or index < 1:
+        raise APIError(APIErrorCode.INVALID_INPUT, "Required: track_id, and entry as a 1-based stsd index")
+    track = session.tracks.get(track_id)
+    if not track:
+        raise APIError(
+            APIErrorCode.TRACK_NOT_FOUND,
+            f"Track not found in session: {track_id}",
+            details={"track_id": track_id, "session_id": session_id},
+        )
+    service = session.service_instance
+    if not service.decrypts_samples(track):
+        raise APIError(APIErrorCode.INVALID_INPUT, f"The service does not decrypt the samples of track {track_id}")
+    if not server_cdm_allowed(request, session.service_tag):
+        raise APIError(
+            APIErrorCode.FORBIDDEN, f"This API key may not decrypt {session.service_tag} samples on the server"
+        )
+
+    def run() -> web.Response:
+        decrypter = service.get_sample_decrypter(track, index)
+        if isinstance(decrypter, (bytes, bytearray)):
+            return web.json_response({"key": bytes(decrypter).hex()})
+        if not frame:
+            return web.json_response({"decrypter": True})
+        try:
+            samples = cbcs.unpack_samples(frame)
+        except ValueError as e:
+            raise APIError(APIErrorCode.INVALID_INPUT, str(e)) from None
+        return web.Response(body=cbcs.pack_samples(decrypter(samples)), content_type="application/octet-stream")
+
+    try:
+        async with session.lock:
+            return await asyncio.to_thread(run)
+    except APIError:
+        raise
+    except (Exception, SystemExit):
+        # `from None`: a debug-mode traceback would otherwise carry the service's exception chain
+        log.exception(f"Error decrypting samples for track {sanitize_log(str(track_id))}")
+        raise APIError(APIErrorCode.SERVICE_ERROR, "Could not decrypt the samples of this track") from None
+
+
 async def session_license_handler(
     data: Dict[str, Any], session_id: str, request: Optional[web.Request] = None
 ) -> web.Response:

unshackle/core/api/routes.py

@@ -58,11 +58,13 @@ from unshackle.core.api.handlers import (
     session_logs_handler,
     session_prompt_get_handler,
     session_prompt_post_handler,
+    session_samples_handler,
     session_segment_filter_handler,
     session_segments_handler,
     session_titles_handler,
     session_tracks_handler,
 )
+from unshackle.core.drm import cbcs
 from unshackle.core.services import Services
 from unshackle.core.update_checker import UpdateChecker
 from unshackle.core.utils.redact import redact_path
@@ -2122,6 +2124,54 @@ async def session_segment_filter(request: web.Request) -> web.Response:
         )
 
 
+@api_handler
+async def session_samples(request: web.Request) -> web.Response:
+    """
+    Decrypt the samples of a track on the server.
+    ---
+    summary: Decrypt samples
+    description: >-
+      For a track the tracks response lists in sample_decrypt_tracks. An empty body asks how the
+      stsd entry decrypts: the JSON answer is its content key (hex) in key, or decrypter true. A
+      body of samples gets their plaintexts back in the same frame: a big-endian uint32 count,
+      one uint32 length per sample, then the samples, each ciphertext cut to whole AES blocks. The
+      body is application/octet-stream and is not in this schema, so the docs page cannot send it.
+      Needs the server_cdm grant for the service.
+    parameters:
+      - name: session_id
+        in: path
+        required: true
+        schema:
+          type: string
+      - name: track_id
+        in: query
+        required: true
+        schema:
+          type: string
+      - name: entry
+        in: query
+        required: true
+        description: "1-based stsd sample entry index"
+        schema:
+          type: integer
+    responses:
+      '200':
+        description: "JSON key (hex) or decrypter: true, or a sample frame of plaintexts in order"
+      '403':
+        description: The API key has no server_cdm grant for the service
+      '404':
+        description: Remote session or track not found
+    """
+    frame = await request.clone(client_max_size=2 * cbcs.SAMPLE_BATCH_BYTES).read()
+    return await session_samples_handler(
+        request.match_info["session_id"],
+        request.query.get("track_id", ""),
+        request.query.get("entry", ""),
+        frame,
+        request,
+    )
+
+
 @api_handler
 @heartbeat
 async def session_license(request: web.Request) -> web.Response:
@@ -3071,6 +3121,7 @@ ROUTES: list[tuple[str, str, Handler, bool]] = [
     ("POST", "/api/session/{session_id}/tracks", session_tracks, True),
     ("POST", "/api/session/{session_id}/segments", session_segments, True),
     ("POST", "/api/session/{session_id}/segment_filter", session_segment_filter, True),
+    ("POST", "/api/session/{session_id}/samples", session_samples, True),
     ("POST", "/api/session/{session_id}/license", session_license, True),
     ("POST", "/api/session/{session_id}/keys/bad", session_bad_key, True),
     ("GET", "/api/session/{session_id}/logs", session_logs, True),

unshackle/core/drm/cbcs.py

@@ -431,6 +431,28 @@ SampleDecrypter = Callable[[list[bytes]], list[bytes]]
 
 KeyResolver = Callable[[TrackLayout, SampleEntry], Union[bytes, SampleDecrypter]]
 
+SAMPLE_BATCH_BYTES = 8 * 1024 * 1024
+"""The most sample bytes one remote `pack_samples` frame carries; a larger single sample travels alone."""
+
+
+def pack_samples(samples: list[bytes]) -> bytes:
+    """Frame samples for transport: a big-endian count, one length per sample, then the samples."""
+    return struct.pack(f">I{len(samples)}I", len(samples), *map(len, samples)) + b"".join(samples)
+
+
+def unpack_samples(frame: bytes) -> list[bytes]:
+    """The samples of a `pack_samples` frame. Raises ValueError for a frame whose lengths do not add up."""
+    count = struct.unpack_from(">I", frame)[0] if len(frame) >= 4 else -1
+    offset = 4 + 4 * count
+    lengths = struct.unpack_from(f">{count}I", frame, 4) if 0 <= count and offset <= len(frame) else None
+    if lengths is None or offset + sum(lengths) != len(frame):
+        raise ValueError("The sample frame is truncated or its lengths do not match its size")
+    samples = []
+    for length in lengths:
+        samples.append(frame[offset : offset + length])
+        offset += length
+    return samples
+
 
 def decrypt_file(path: Path, key_for: KeyResolver) -> None:
     """Decrypt every cbcs fragment of `path` in place, with the content key `key_for` names per sample entry.
@@ -546,7 +568,7 @@ def _decrypt_with(decrypter: SampleDecrypter, buf: bytearray, span_start: int, f
             samples.append(runs)
     if not samples:
         return
-    ciphertexts = [b"".join(bytes(buf[start : start + length]) for start, length in runs) for runs in samples]
+    ciphertexts = [b"".join(buf[start : start + length] for start, length in runs) for runs in samples]
     plaintexts = decrypter(ciphertexts)
     if len(plaintexts) != len(ciphertexts):
         raise ValueError(f"The sample decrypter returned {len(plaintexts)} samples for {len(ciphertexts)}")

unshackle/core/remote_service.py

@@ -23,6 +23,7 @@ from http.cookiejar import CookieJar
 from pathlib import Path
 from threading import Event, Lock, Thread
 from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Union
+from urllib.parse import urlencode
 from uuid import UUID
 
 import click
@@ -37,6 +38,7 @@ from unshackle.core.config import config
 from unshackle.core.console import console, prompt_user
 from unshackle.core.constants import AnyTrack
 from unshackle.core.credential import Credential
+from unshackle.core.drm import cbcs
 from unshackle.core.proxies.resolve import is_loopback
 from unshackle.core.session import MAX_BACKOFF
 from unshackle.core.titles import Title_T, Titles_T, remap_titles
@@ -110,13 +112,20 @@ class RemoteClient:
         data: Optional[Dict[str, Any]] = None,
         optional: bool = False,
         expect: tuple[str, ...] = (),
+        body: Optional[bytes] = None,
     ) -> Dict[str, Any]:
-        """Send one request and return its JSON body; an error outside ``expect`` exits, one inside it returns."""
+        """Send one request and return its JSON body; an error outside ``expect`` exits, one inside it returns.
+
+        A raw ``body`` goes as ``application/octet-stream``, and a binary answer comes back under ``"body"``.
+        """
         url = f"{self.server_url}{endpoint}"
+        payload: Dict[str, Any] = (
+            {"json": data} if body is None else {"data": body, "headers": {"Content-Type": "application/octet-stream"}}
+        )
         while True:
             try:
                 with self.in_flight(endpoint):
-                    resp = getattr(self.session, method)(url, json=data, timeout=self.timeout)
+                    resp = getattr(self.session, method)(url, timeout=self.timeout, **payload)
             except requests.ConnectionError as e:
                 if isinstance(e.args[0] if e.args else None, urllib3.exceptions.ReadTimeoutError):
                     log.error(
@@ -137,6 +146,8 @@ class RemoteClient:
             if resp.status_code == 401 and self.next_auth_header():
                 continue
             break
+        if resp.status_code < 400 and "octet-stream" in resp.headers.get("Content-Type", ""):
+            return {"body": resp.content}
         try:
             detail = resp.json()
         except ValueError:
@@ -186,6 +197,20 @@ def enum_get(enum_cls: type[Enum], name: Optional[str], default: Any = None) ->
         return default
 
 
+def sample_batches(samples: list[bytes]) -> Iterator[list[bytes]]:
+    """Split samples into consecutive batches of at most cbcs.SAMPLE_BATCH_BYTES (a larger sample goes alone)."""
+    batch: list[bytes] = []
+    size = 0
+    for sample in samples:
+        if batch and size + len(sample) > cbcs.SAMPLE_BATCH_BYTES:
+            yield batch
+            batch, size = [], 0
+        batch.append(sample)
+        size += len(sample)
+    if batch:
+        yield batch
+
+
 def base_track_kwargs(data: Dict[str, Any]) -> Dict[str, Any]:
     """The shared Track constructor kwargs a serialized track carries.
 
@@ -913,6 +938,7 @@ class RemoteService:
         self.server_vault_keys: Dict[UUID, str] = {}
         self.server_vault = ServerVault(self)
         self._segment_filters: Dict[str, Callable[[Any], bool]] = {}
+        self._sample_decrypt_tracks: set[str] = set()
         self._log_seq = 0
         self._log_drain_lock = Lock()
         self._keepalive_stop = Event()
@@ -1308,7 +1334,9 @@ class RemoteService:
 
         One poller serves every concurrent request, and the local CDM sessions close when the last one ends.
         """
-        if not endpoint.startswith(f"/api/session/{self._session_id}/") or endpoint.endswith(("/prompt", "/logs")):
+        if not endpoint.startswith(f"/api/session/{self._session_id}/") or endpoint.partition("?")[0].endswith(
+            ("/prompt", "/logs", "/samples")
+        ):
             yield
             return
         with self._cdm_lock:
@@ -1407,6 +1435,7 @@ class RemoteService:
                 "falling back to the local CDM"
             )
         self._server_cdm_type = result.get("server_cdm_type", "widevine")
+        self._sample_decrypt_tracks.update(result.get("sample_decrypt_tracks") or [])
 
         for track in [*tracks.videos, *tracks.audio, *tracks.subtitles]:
             if track.descriptor == Track.Descriptor.HLS:
@@ -1417,6 +1446,25 @@ class RemoteService:
 
         return tracks
 
+    def decrypts_samples(self, track: AnyTrack) -> bool:
+        return str(track.id) in self._sample_decrypt_tracks
+
+    def get_sample_decrypter(self, track: AnyTrack, entry: int) -> Union[bytes, cbcs.SampleDecrypter]:
+        """Ask the server how `entry` decrypts. A decrypter sends the samples to the server in batches."""
+        endpoint = f"/api/session/{self._session_id}/samples?{urlencode({'track_id': str(track.id), 'entry': entry})}"
+        answer = self.client.request("post", endpoint, body=b"")
+        if answer.get("key"):
+            return bytes.fromhex(answer["key"])
+
+        def decrypt(samples: list[bytes]) -> list[bytes]:
+            plaintexts: list[bytes] = []
+            for batch in sample_batches(samples):
+                frame = self.client.request("post", endpoint, body=cbcs.pack_samples(batch))["body"]
+                plaintexts.extend(cbcs.unpack_samples(frame))
+            return plaintexts
+
+        return decrypt
+
     def remote_segment_filter(self, track_id: str) -> Callable[[Any], bool]:
         """The server's ``OnSegmentFilter`` result for one HLS track, fetched on first use.
 

unshackle/core/service.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 
 if TYPE_CHECKING:
     from unshackle.core.api.input_bridge import InputBridge
+    from unshackle.core.drm.cbcs import SampleDecrypter
 
 import click
 import m3u8
@@ -720,6 +721,20 @@ class Service(metaclass=ABCMeta):
         option `chapter_fallback_name`. For example, `"Chapter {i:02}"` for "Chapter 01".
         """
 
+    # Optional sample decryption
+
+    def decrypts_samples(self, track: AnyTrack) -> bool:
+        """Whether this service decrypts the cbcs file of `track` after the download, through `get_sample_decrypter`."""
+        return False
+
+    def get_sample_decrypter(self, track: AnyTrack, entry: int) -> Union[bytes, "SampleDecrypter"]:
+        """The content key, or a `cbcs.SampleDecrypter`, for `stsd` entry `entry` (1-based) of `track`.
+
+        unshackle calls this only for a track `decrypts_samples` names. A remote session runs it on the
+        server, so a decrypter that needs a local process works for remote clients too.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not decrypt samples")
+
     # Optional Event methods
 
     def on_segment_downloaded(self, track: AnyTrack, segment: Path) -> None: