fix: harden ClearKey/CENC decryption correctness

tests/core/test_clearkey.py

@@ -0,0 +1,178 @@
+"""Tests for the HLS AES-128 ClearKey DRM system (``unshackle.core.drm.clearkey``).
+
+Covers the correctness guards on the whole-segment AES-128-CBC path:
+- ``ClearKey.decrypt`` passing mislabeled-clear / non-block-aligned segments through
+  untouched instead of crashing or corrupting (WS1 / F1+F3)
+- ``ClearKey.from_m3u_key`` ``data:`` URI parsing and 16-byte key validation (WS5 / F4)
+- the RFC 8216 §5.2 sequence-number IV deriving from a segment's absolute media
+  sequence number rather than the post-filter download index (WS2 / F2)
+"""
+
+from __future__ import annotations
+
+import base64
+
+import m3u8
+import pytest
+from Cryptodome.Cipher import AES
+from Cryptodome.Util.Padding import pad
+from m3u8.model import Key
+
+from unshackle.core.drm.clearkey import SYNC_BYTE, TS_PACKET_SIZE, ClearKey
+
+KEY = bytes.fromhex("00112233445566778899aabbccddeeff")
+IV = bytes.fromhex("0f0e0d0c0b0a09080706050403020100")
+
+
+def encrypt(plaintext: bytes, key: bytes = KEY, iv: bytes = IV) -> bytes:
+    return AES.new(key, AES.MODE_CBC, iv).encrypt(pad(plaintext, AES.block_size))
+
+
+def clear_ts(packets: int = 4) -> bytes:
+    # Plaintext MPEG-TS: sync byte 0x47 at every 188-byte packet boundary. Default 4 packets so
+    # the fixture satisfies the hardened 4-boundary sniff (and 752 B is 16-aligned, so it must be
+    # caught by the clear-TS guard, not the len%16 guard).
+    return b"".join(bytes([SYNC_BYTE]) + b"\x00" * (TS_PACKET_SIZE - 1) for _ in range(packets))
+
+
+def test_valid_encrypted_segment_decrypts(tmp_path) -> None:
+    plaintext = b"the quick brown fox jumps over" * 4
+    seg = tmp_path / "0.ts"
+    seg.write_bytes(encrypt(plaintext))
+
+    ClearKey(key=KEY, iv=IV).decrypt(seg)
+
+    assert seg.read_bytes() == plaintext
+
+
+def test_clear_ts_passes_through(tmp_path) -> None:
+    data = clear_ts()
+    seg = tmp_path / "0.ts"
+    seg.write_bytes(data)
+
+    ClearKey(key=KEY, iv=IV).decrypt(seg)
+
+    assert seg.read_bytes() == data
+
+
+def test_non_block_aligned_passes_through(tmp_path) -> None:
+    data = b"\x01" * 17  # not TS, and not a 16-byte multiple -> not valid CBC ciphertext
+    assert len(data) % AES.block_size != 0
+    seg = tmp_path / "0.ts"
+    seg.write_bytes(data)
+
+    ClearKey(key=KEY, iv=IV).decrypt(seg)
+
+    assert seg.read_bytes() == data
+
+
+def test_empty_segment_does_not_crash(tmp_path) -> None:
+    seg = tmp_path / "0.ts"
+    seg.write_bytes(b"")
+
+    ClearKey(key=KEY, iv=IV).decrypt(seg)
+
+    assert seg.exists()
+    assert seg.read_bytes() == b""
+
+
+def test_warns_once_per_instance(tmp_path, caplog) -> None:
+    drm = ClearKey(key=KEY, iv=IV)
+    for name in ("0.ts", "1.ts"):
+        seg = tmp_path / name
+        seg.write_bytes(clear_ts())
+        with caplog.at_level("WARNING"):
+            drm.decrypt(seg)
+
+    warnings = [r for r in caplog.records if r.levelname == "WARNING"]
+    assert len(warnings) == 1
+
+
+def _data_key(raw: bytes, base64_encoded: bool = True) -> Key:
+    if base64_encoded:
+        uri = "data:text/plain;base64," + base64.b64encode(raw).decode("ascii")
+    else:
+        uri = "data:text/plain," + raw.decode("latin-1")
+    return Key(method="AES-128", uri=uri, base_uri="", keyformat=None, keyformatversions=None, iv=None)
+
+
+def test_data_uri_key_decodes_and_validates_length() -> None:
+    drm = ClearKey.from_m3u_key(_data_key(KEY))
+    assert drm.key == KEY
+
+
+def test_data_uri_key_wrong_length_rejected() -> None:
+    with pytest.raises(ValueError, match="Unexpected Length"):
+        ClearKey.from_m3u_key(_data_key(b"\x00" * 8))
+
+
+def test_raw_data_uri_key_materialises_and_validates() -> None:
+    # A raw (non-base64) data: payload must be taken as literal key bytes, not hex-parsed.
+    drm = ClearKey.from_m3u_key(_data_key(KEY, base64_encoded=False))
+    assert drm.key == KEY
+
+
+def test_raw_data_uri_key_wrong_length_rejected() -> None:
+    with pytest.raises(ValueError, match="Unexpected Length"):
+        ClearKey.from_m3u_key(_data_key(b"\x01" * 8, base64_encoded=False))
+
+
+def test_data_uri_split_tolerates_commas_in_payload() -> None:
+    # base64 with padding contains no comma, but the split must only break on the first comma
+    # so a payload that itself contains a comma is not truncated. Encode a 16-byte key whose
+    # base64 happens to be split-safe, and assert the mediatype comma is the only split point.
+    raw = bytes(range(16))
+    b64 = base64.b64encode(raw).decode("ascii")
+    key = Key(
+        method="AES-128",
+        uri=f"data:application/octet-stream;base64,{b64}",
+        base_uri="",
+        keyformat=None,
+        keyformatversions=None,
+        iv=None,
+    )
+    drm = ClearKey.from_m3u_key(key)
+    assert drm.key == raw
+
+
+SEQ_PLAYLIST = """#EXTM3U
+#EXT-X-VERSION:3
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:100
+#EXT-X-KEY:METHOD=AES-128,URI="skd://key"
+#EXTINF:6.0,
+seg0.ts
+#EXTINF:6.0,
+seg1.ts
+#EXTINF:6.0,
+seg2.ts
+#EXTINF:6.0,
+seg3.ts
+#EXT-X-ENDLIST
+"""
+
+
+def test_sequence_number_iv_uses_absolute_media_sequence() -> None:
+    """IV must be the segment's absolute media sequence number, not the download index.
+
+    Mirrors the derivation in ``HLS.download_track``: filtered segments are dropped, the
+    surviving segments are named by post-filter download index, and the IV comes from the
+    wanted segment's ``media_sequence`` (EXT-X-MEDIA-SEQUENCE + absolute position).
+    """
+    master = m3u8.loads(SEQ_PLAYLIST)
+    # media sequence numbers are the offset (100) plus absolute playlist position.
+    assert [s.media_sequence for s in master.segments] == [100, 101, 102, 103]
+
+    # Drop an early segment (index 1), as an OnSegmentFilter would.
+    unwanted = {master.segments[1]}
+    wanted_segments = [s for s in master.segments if s not in unwanted]
+
+    # download index 1 is the SECOND surviving segment (absolute seq 102), not 101.
+    download_i = 1
+    seq = wanted_segments[download_i].media_sequence
+    assert seq == 102
+
+    iv = int(seq).to_bytes(16, "big")
+    assert iv == (102).to_bytes(16, "big")
+    # the old download-index derivation would have produced the wrong IV (100 + 1 = 101)
+    assert iv != (100 + download_i).to_bytes(16, "big")

unshackle/commands/dl.py

@@ -1851,6 +1851,7 @@ class dl:
                                         cdm_only=cdm_only,
                                         vaults_only=vaults_only,
                                         export=export_path,
+                                        service_session=service.session,
                                     ),
                                     cdm=self.cdm,
                                     max_workers=workers,
@@ -2921,6 +2922,7 @@ class dl:
                                 cdm_only=cdm_only,
                                 vaults_only=vaults_only,
                                 export=export_path,
+                                service_session=service.session,
                             ),
                             cdm=self.cdm,
                             max_workers=workers,
@@ -3645,6 +3647,7 @@ class dl:
         cdm_only: bool = False,
         vaults_only: bool = False,
         export: Optional[Path] = None,
+        service_session: Optional[Any] = None,
     ) -> None:
         """
         Prepare the DRM by getting decryption data like KIDs, Keys, and such.
@@ -4175,7 +4178,7 @@ class dl:
                     from_vaults = drm.content_keys.copy()
 
                     try:
-                        drm.get_content_keys(licence=clearkey_licence or (lambda **_: None))
+                        drm.get_content_keys(licence=clearkey_licence or (lambda **_: None), session=service_session)
                     except Exception as e:
                         if drm.content_keys:
                             self.log.debug(f"License call failed but keys already in content_keys: {e}")

unshackle/core/drm/clearkey.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import base64
+import logging
 import shutil
 from pathlib import Path
 from typing import Optional, Union
@@ -14,6 +15,14 @@ from requests import Session
 from unshackle.core.session import RnetSession
 from unshackle.core.utilities import log_event
 
+log = logging.getLogger(__name__)
+
+# MPEG-TS: 188-byte packets, sync byte 0x47 at every packet boundary (shared with sample_aes).
+SYNC_BYTE = 0x47
+TS_PACKET_SIZE = 188
+# Consecutive packet boundaries that must carry the sync byte to call a segment clear MPEG-TS.
+TS_SYNC_CHECKS = 4
+
 
 class ClearKey:
     """AES Clear Key DRM System."""
@@ -39,12 +48,38 @@ class ClearKey:
 
         self.key: bytes = key
         self.iv: bytes = iv
+        self._warned_clear: bool = False
+
+    def _warn_clear(self, reason: str) -> None:
+        """Warn once per track that a segment was passed through undecrypted."""
+        if not self._warned_clear:
+            log.warning("ClearKey: manifest signaled AES-128 encryption but %s; passing segment through.", reason)
+            self._warned_clear = True
 
     def decrypt(self, path: Path) -> None:
         """Decrypt a Track with AES Clear Key DRM."""
         if not path or not path.exists():
             raise ValueError("Tried to decrypt a file that does not exist.")
 
+        data = path.read_bytes()
+
+        # Mislabeled-clear guard: plaintext MPEG-TS the manifest wrongly tags AES-128. Require the
+        # 0x47 sync byte at 4 consecutive packet boundaries plus an exact packet-count length, so
+        # real CBC ciphertext matching all of these is negligible (~1/256^4).
+        if (
+            len(data) % TS_PACKET_SIZE == 0
+            and len(data) >= TS_PACKET_SIZE * TS_SYNC_CHECKS
+            and all(data[i * TS_PACKET_SIZE] == SYNC_BYTE for i in range(TS_SYNC_CHECKS))
+        ):
+            self._warn_clear("the segment is clear MPEG-TS")
+            return
+
+        # Valid AES-128-CBC ciphertext is always a 16-byte multiple; anything else is not
+        # ciphertext (truncated/clear/SAMPLE-AES). Decrypting would raise, so pass it through.
+        if len(data) % AES.block_size != 0:
+            self._warn_clear(f"the segment length ({len(data)} bytes) is not a multiple of {AES.block_size}")
+            return
+
         log_event(
             "drm_decrypt",
             level="DEBUG",
@@ -53,15 +88,6 @@ class ClearKey:
             file=path.name,
         )
 
-        data = path.read_bytes()
-
-        # Guard mislabeled-clear: EXT-X-KEY says AES-128 but the segment is plaintext MPEG-TS.
-        # CBC ciphertext is 16-aligned, but so are some clear TS runs (752 bytes = 4 packets),
-        # so also sniff the 0x47 sync byte at packet 0 and 1.
-        is_ts = data[:1] == b"\x47" and data[188:189] == b"\x47"
-        if len(data) % AES.block_size != 0 or is_ts:
-            return
-
         decrypted = AES.new(self.key, AES.MODE_CBC, self.iv).decrypt(data)
 
         try:
@@ -105,11 +131,17 @@ class ClearKey:
             session.headers["User-Agent"] = "smartexoplayer/1.1.0 (Linux;Android 8.0.0) ExoPlayerLib/2.13.3"
 
         if m3u_key.uri.startswith("data:"):
-            media_types, data = m3u_key.uri[5:].split(",")
+            media_types, payload = m3u_key.uri[5:].split(",", 1)
             media_types = media_types.split(";")
+            # Materialise the key bytes here so the 16-byte check covers both encodings: base64
+            # decodes to bytes, a raw payload is latin-1 bytes. A raw str would otherwise reach
+            # __init__ and be reinterpreted through fromhex.
             if "base64" in media_types:
-                data = base64.b64decode(data)
-            key = data
+                key = base64.b64decode(payload)
+            else:
+                key = payload.encode("latin-1")
+            if len(key) != 16:
+                raise ValueError(f"Unexpected Length of Key ({len(key)} bytes) in M3U data: URI.")
         else:
             url = urljoin(m3u_key.base_uri, m3u_key.uri)
             log_event(

unshackle/core/drm/clearkey_cenc.py

@@ -132,14 +132,33 @@ class ClearKeyCENC:
             raise ValueError(f"Expected the ClearKey license to be bytes, str, or dict, not {response!r}")
 
         for jwk in document.get("keys") or []:
-            if jwk.get("kty") not in (None, "oct"):
+            kty = jwk.get("kty")
+            if kty not in (None, "oct"):
                 continue
+            if kty is None:
+                log_event(
+                    "drm_license_warning",
+                    level="WARNING",
+                    message="ClearKey JWK is missing 'kty'; assuming 'oct' per W3C EME",
+                    drm_type="ClearKeyCENC",
+                )
             kid_b64 = jwk.get("kid")
             key_b64 = jwk.get("k")
             if not kid_b64 or not key_b64:
                 continue
             kid = UUID(bytes=b64url_decode(kid_b64))
-            self.content_keys[kid] = b64url_decode(key_b64).hex()
+            key = b64url_decode(key_b64)
+            if len(key) != 16:
+                # A CENC content key is always 16 bytes; a wrong-length key would make the
+                # external tool fail opaquely. Skip it rather than pass it through.
+                log_event(
+                    "drm_license_warning",
+                    level="WARNING",
+                    message=f"ClearKey JWK for KID {kid.hex} has a {len(key)}-byte key (expected 16); skipping",
+                    drm_type="ClearKeyCENC",
+                )
+                continue
+            self.content_keys[kid] = key.hex()
 
         if not self.content_keys:
             raise ClearKeyCENC.Exceptions.EmptyLicense("No Content Keys were within the License")

unshackle/core/manifests/hls.py

@@ -709,6 +709,9 @@ class HLS:
         # First segment's Media Sequence Number - used as the AES-128 IV when EXT-X-KEY has
         # no explicit IV (RFC 8216 §5.2), where each segment's IV is its sequence number.
         media_sequence_start = getattr(master, "media_sequence", None) or 0
+        # Downloaded segment files are named by post-filter index; map that back to the wanted
+        # segment so the IV uses the absolute media sequence number, not the download index.
+        wanted_segments = [seg for seg in master.segments if seg not in unwanted_segments]
         map_data: Optional[tuple[m3u8.model.InitializationSection, bytes]] = None
         if session_drm:
             encryption_data: Optional[tuple[Optional[m3u8.Key], DRM_T]] = (initial_drm_key, session_drm)
@@ -802,7 +805,11 @@ class HLS:
                     seq_iv = isinstance(drm, ClearKey) and not (key_obj and key_obj.iv)
                     for file in files:
                         if seq_iv and file.stem.isdigit():
-                            drm.iv = (media_sequence_start + int(file.stem)).to_bytes(16, "big")
+                            download_i = int(file.stem)
+                            seq = getattr(wanted_segments[download_i], "media_sequence", None)
+                            if seq is None:
+                                seq = media_sequence_start + download_i
+                            drm.iv = int(seq).to_bytes(16, "big")
                         drm.decrypt(file)
                     merge(to=merged_path, via=files, delete=True, include_map_data=True)