fix(hls): read the key id from the init segment when the PSSH has none

An EXT-X-KEY whose Widevine PSSH holds a content id and no key id raised KIDNotFound. Covers download, licence-only, --all-drm, session keys, import and the server CDM path. Fixes #166.
2026-11-12 03:55:05 +00:00 · 2026-10-13 12:14:34 -06:00
parent b4136e1cf5
commit af4e6d5d47
6 changed files with 256 additions and 32 deletions
@@ -0,0 +1,132 @@
+"""An HLS Widevine key whose PSSH holds a content id and no Key ID takes the Key ID of its init segment."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+from uuid import UUID
+
+import m3u8
+from pywidevine.license_protocol_pb2 import WidevinePsshData
+from pywidevine.pssh import PSSH
+
+from unshackle.core.constants import DOWNLOAD_LICENCE_ONLY
+from unshackle.core.drm import PlayReady, Widevine
+from unshackle.core.import_service import hls_playlist_kids
+from unshackle.core.manifests.hls import HLS
+from unshackle.core.tracks import Video
+from unshackle.core.tracks.track import DownloadContext, Track
+
+INIT_KID = UUID("11111111-1111-1111-1111-111111111111")
+PSSH_KID = UUID("22222222-2222-2222-2222-222222222222")
+TAG_KID = UUID("33333333-3333-3333-3333-333333333333")
+
+ZERO_KID = UUID(int=0)
+
+
+def no_kid_pssh(content_id: bytes = b"content-id-only") -> str:
+    return PSSH.new(
+        system_id=PSSH.SystemId.Widevine,
+        init_data=WidevinePsshData(content_id=content_id).SerializeToString(),
+    ).dumps()
+
+
+NO_KID_PSSH = no_kid_pssh()
+KEY_LINE = (
+    '#EXT-X-KEY:METHOD=SAMPLE-AES-CTR,URI="data:text/plain;base64,{pssh}",'
+    'KEYFORMAT="urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed",KEYFORMATVERSIONS="1"'
+)
+
+
+def keyed_segment(pssh: str, extra: str = "", init: str = "init.mp4") -> str:
+    return f'#EXT-X-MAP:URI="{init}"\n{KEY_LINE.format(pssh=pssh)}{extra}\n#EXTINF:6.0,\n{init}.m4s\n'
+
+
+def playlist(*segments: str) -> str:
+    body = "#EXT-X-DISCONTINUITY\n".join(segments)
+    return f"#EXTM3U\n#EXT-X-VERSION:7\n#EXT-X-TARGETDURATION:6\n{body}#EXT-X-ENDLIST\n"
+
+
+def key_of(pssh: str, extra: str = "") -> m3u8.Key:
+    return m3u8.loads(playlist(keyed_segment(pssh, extra))).keys[0]
+
+
+def licence_only(text: str, tmp_path: Path) -> list[tuple[list[UUID], Any]]:
+    """Run a licence-only download of the playlist and return the Key IDs of each licence call."""
+    media = tmp_path / "media.m3u8"
+    media.write_text(text, encoding="utf-8")
+    track = Video(
+        url="https://example.test/media.m3u8",
+        language="en",
+        descriptor=Track.Descriptor.HLS,
+        codec=Video.Codec.AVC,
+    )
+    track.from_file = media
+    licensed: list[tuple[list[UUID], Any]] = []
+    ctx = DownloadContext(
+        save_path=tmp_path / "video.mp4",
+        save_dir=tmp_path / "video_segments",
+        progress=lambda **_: None,
+        license_widevine=lambda drm, track_kid=None: licensed.append((drm.kids, track_kid)),
+        cdm=None,
+    )
+    DOWNLOAD_LICENCE_ONLY.set()
+    try:
+        HLS.download_track(track=track, ctx=ctx)
+    finally:
+        DOWNLOAD_LICENCE_ONLY.clear()
+    return licensed
+
+
+def test_download_track_licenses_with_the_init_segment_kid(tmp_path: Path, monkeypatch: Any) -> None:
+    # the probe itself needs the network and FFprobe; the order of the calls is under test
+    monkeypatch.setattr(HLS, "get_track_kid_from_init", staticmethod(lambda *args, **kwargs: INIT_KID))
+    assert licence_only(playlist(keyed_segment(NO_KID_PSSH)), tmp_path) == [([INIT_KID], INIT_KID)]
+
+
+def test_licence_only_reads_each_later_key_from_its_own_init_segment(tmp_path: Path, monkeypatch: Any) -> None:
+    def probe(master: Any, track: Any, session: Any, init_section: Any = None) -> UUID:
+        return PSSH_KID if init_section is not None and init_section.uri == "init2.mp4" else INIT_KID
+
+    monkeypatch.setattr(HLS, "get_track_kid_from_init", staticmethod(probe))
+    text = playlist(keyed_segment(no_kid_pssh(b"first")), keyed_segment(no_kid_pssh(b"later"), init="init2.mp4"))
+    assert licence_only(text, tmp_path) == [([INIT_KID], INIT_KID), ([PSSH_KID], PSSH_KID)]
+
+
+def test_get_drm_keeps_a_pssh_that_has_its_own_kid() -> None:
+    pssh = PSSH.new(key_ids=[PSSH_KID], system_id=PSSH.SystemId.Widevine).dumps()
+    drm = HLS.get_drm(key_of(pssh), kid=INIT_KID)
+    assert drm.kids == [PSSH_KID]
+    assert drm.pssh.dumps() == pssh
+
+
+def test_get_drm_prefers_the_tag_keyid_over_the_init_segment_kid() -> None:
+    drm = HLS.get_drm(key_of(NO_KID_PSSH, f",KEYID=0x{TAG_KID.hex}"), kid=INIT_KID)
+    assert drm.kids == [TAG_KID]
+
+
+def test_get_drm_replaces_a_placeholder_pssh_kid_and_ignores_a_placeholder_init_kid() -> None:
+    placeholder = PSSH.new(key_ids=[ZERO_KID], system_id=PSSH.SystemId.Widevine).dumps()
+    assert HLS.get_drm(key_of(placeholder), kid=INIT_KID).kids == [INIT_KID]
+    assert HLS.get_drm(key_of(placeholder), kid=ZERO_KID).kids == [ZERO_KID]
+
+
+def test_session_key_without_a_kid_does_not_stop_the_parse() -> None:
+    with_kid = PSSH.new(key_ids=[PSSH_KID], system_id=PSSH.SystemId.Widevine).dumps()
+    drm = HLS.get_all_drm([key_of(NO_KID_PSSH), key_of(with_kid)])
+    assert [d.kids for d in drm] == [[PSSH_KID]]
+
+
+def test_alternate_drm_builds_the_other_system_with_the_init_segment_kid() -> None:
+    chosen = object.__new__(PlayReady)
+    alternate = HLS.alternate_drm([key_of(NO_KID_PSSH)], chosen, kid=INIT_KID)
+    assert [type(d) for d in alternate] == [Widevine] and alternate[0].kids == [INIT_KID]
+
+
+def test_import_reads_the_kid_of_a_playlist_whose_key_has_none() -> None:
+    def get(url: str, headers: Any = None) -> Any:
+        return SimpleNamespace(text=playlist(keyed_segment(NO_KID_PSSH)), ok=True, content=b"init")
+
+    track = SimpleNamespace(id="vid", url="https://example.test/media.m3u8", get_key_id=lambda init: INIT_KID)
+    assert hls_playlist_kids(track, SimpleNamespace(get=get)) == {INIT_KID}
@@ -21,6 +21,7 @@ from Cryptodome.Cipher import AES
 from Cryptodome.Util.Padding import pad
 from pywidevine.pssh import PSSH

+from tests.core.test_hls_kid_from_init import KEY_LINE, no_kid_pssh
 from tests.core.test_ism_init import piff_fragment
 from unshackle.core.config import config
 from unshackle.core.constants import DOWNLOAD_CANCELLED
@@ -264,3 +265,34 @@ def test_hls_aes128_keeps_post_merge(server, tmp_path, monkeypatch):

    plain, rolling = download_both_ways(HLS, make, tmp_path, monkeypatch, False)
    assert rolling == plain == INIT + b"".join(PARTS)
+
+
+def test_hls_rekey_takes_the_kid_from_the_init_segment_of_the_new_key(server, tmp_path, monkeypatch):
+    """Each key carries a PSSH with a content id and no Key ID, so only the init segment names the KID."""
+    later_kid = UUID(hex="22222222222222222222222222222222")
+    later_init = bytes(range(100, 200))
+    server.routes["/init2.mp4"] = later_init
+    no_decrypt(monkeypatch, hls_module)
+    monkeypatch.setattr(config, "merge_segments", False)
+
+    def key_line(content_id: bytes) -> str:
+        return KEY_LINE.format(pssh=no_kid_pssh(content_id))
+
+    track = hls_track(server, False)
+    lines = ["#EXTM3U", "#EXT-X-VERSION:7", "#EXT-X-TARGETDURATION:4", '#EXT-X-MAP:URI="init.mp4"', key_line(b"a")]
+    for i in range(len(PARTS)):
+        if i == 2:
+            lines += ["#EXT-X-DISCONTINUITY", '#EXT-X-MAP:URI="init2.mp4"', key_line(b"b")]
+        lines += ["#EXTINF:4.0,", f"part{i}.m4s"]
+    server.routes["/media.m3u8"] = "\n".join(lines + ["#EXT-X-ENDLIST"]).encode()
+    # the real probe needs FFprobe and a tenc box; the KID follows the init bytes the parser hands over
+    monkeypatch.setattr(
+        type(track), "get_key_id", lambda self, init_data=None, **kw: {later_init: later_kid}.get(init_data, KID)
+    )
+
+    licensed: list[UUID] = []
+    ctx = make_ctx(tmp_path, "rekey.mp4")
+    ctx.license_widevine = lambda drm, track_kid=None: licensed.append(track_kid)
+    HLS.download_track(track, ctx)
+
+    assert licensed == [KID, later_kid]
@@ -13,6 +13,7 @@ from uuid import UUID
 import pytest
 from pywidevine.pssh import PSSH as WvPSSH

+from tests.core.test_hls_kid_from_init import NO_KID_PSSH
 from unshackle.core.api import handlers
 from unshackle.core.drm import Widevine
 from unshackle.core.tracks import Track
@@ -22,13 +23,14 @@ pytestmark = pytest.mark.unit
 KID = UUID("161219ec3df64e0eadbb3827c0ccc46d")


-def _media_playlist() -> str:
-    pssh = WvPSSH.new(key_ids=[KID], system_id=WvPSSH.SystemId.Widevine).dumps()
+def _media_playlist(pssh: str | None = None) -> str:
+    pssh = pssh or WvPSSH.new(key_ids=[KID], system_id=WvPSSH.SystemId.Widevine).dumps()
    return "\n".join(
        [
            "#EXTM3U",
            "#EXT-X-VERSION:6",
            "#EXT-X-TARGETDURATION:6",
+            '#EXT-X-MAP:URI="init.mp4"',
            f'#EXT-X-KEY:METHOD=SAMPLE-AES,URI="data:text/plain;base64,{pssh}",'
            'KEYFORMAT="urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed",KEYFORMATVERSIONS="1"',
            "#EXTINF:6.0,",
@@ -49,3 +51,21 @@ def test_hls_descriptor_alone_reaches_the_media_playlist_keys():
    assert track.drm and isinstance(track.drm[0], Widevine), "media playlist DRM not read"
    assert track.drm[0].kids == [KID]
    assert base64.b64decode(handlers.extract_pssh_from_track(track, "widevine"))
+
+
+def test_key_without_a_kid_takes_it_from_the_init_segment():
+    """A PSSH that holds a content id and no Key ID is DRM, and the server must not call the track clear."""
+    response = SimpleNamespace(text=_media_playlist(NO_KID_PSSH), ok=True, content=b"init")
+    session = SimpleNamespace(get=lambda url, headers=None: response)
+    track = SimpleNamespace(
+        id="vid",
+        url="https://cdn.example.invalid/media.m3u8",
+        descriptor=Track.Descriptor.HLS,
+        data={},
+        drm=None,
+        get_key_id=lambda init: KID,
+    )
+
+    handlers.ensure_track_drm(track, session)
+
+    assert track.drm and track.drm[0].kids == [KID]
@@ -4327,7 +4327,7 @@ def ensure_track_drm(track: Any, session: Any = None, init_data: Optional[bytes]
            keys = [k for k in (playlist.keys or []) + (playlist.session_keys or []) if k is not None]
            for key in keys:
                try:
-                    drm = HLS.get_drm(key)
+                    drm = HLS.get_drm_or_probe(key, playlist, track, probe_session=http)
                    if isinstance(drm, (Widevine, PlayReady)):
                        track.drm = [drm]
                        return
@@ -67,9 +67,10 @@ def hls_playlist_kids(track: Any, session: Optional[requests.Session] = None) ->
    key from its URI, so such a track takes no exported content key. An empty set means that
    the track is encrypted but names no KID that unshackle can read.
    """
+    http = session or requests.Session()
    try:
-        text = (session or requests.Session()).get(str(track.url)).text
-        keys = [k for k in m3u8.loads(text, uri=str(track.url)).keys if k is not None and k.method != "NONE"]
+        playlist = m3u8.loads(http.get(str(track.url)).text, uri=str(track.url))
+        keys = [k for k in playlist.keys if k is not None and k.method != "NONE"]
    except Exception as e:
        log.warning(f"Cannot read the media playlist of {track.id} for its KIDs: {e!r}")
        return None
@@ -78,7 +79,7 @@ def hls_playlist_kids(track: Any, session: Optional[requests.Session] = None) ->
        if key.method == "AES-128":
            continue
        try:
-            kids.update(real_kids(HLS.get_drm(key)))
+            kids.update(real_kids(HLS.get_drm_or_probe(key, playlist, track, probe_session=http)))
        except Exception as e:
            log.debug(f"Skipping HLS key {key.keyformat or key.method} of {track.id}: {e!r}")
    if kids or any(k.method != "AES-128" for k in keys):
@@ -706,9 +706,14 @@ class HLS:
            media_playlist_key = HLS.get_supported_key(cdm_media_keys) if cdm_media_keys else None

            if media_playlist_key:
-                media_drm = HLS.get_drm(media_playlist_key, session)
+                init_kid = (
+                    None
+                    if media_playlist_key.method == "AES-128"
+                    else HLS.get_track_kid_from_init(master, track, session)
+                )
+                media_drm = HLS.get_drm(media_playlist_key, session, **({"kid": init_kid} if init_kid else {}))
                if isinstance(media_drm, (Widevine, PlayReady)):
-                    track_kid = HLS.get_track_kid_from_init(master, track, session) or media_drm.kid
+                    track_kid = init_kid or media_drm.kid
                    # Preserve pre-existing keys (e.g. from server_cdm)
                    if track.drm:
                        for existing_drm in track.drm:
@@ -716,7 +721,7 @@ class HLS:
                                media_drm.content_keys.update(existing_drm.content_keys)
                    track.drm = [media_drm]
                    if DOWNLOAD_ALL_DRM.is_set():
-                        track.drm += HLS.alternate_drm(media_keys, media_drm, session)
+                        track.drm += HLS.alternate_drm(media_keys, media_drm, session, init_kid)
                    try:
                        if not license_widevine:
                            raise ValueError("license_widevine func must be supplied to use DRM")
@@ -772,7 +777,7 @@ class HLS:
                if key is None or key.method == "AES-128" or key in licensed_keys:
                    continue
                licensed_keys.append(key)
-                later_drm = HLS.get_drm(key, session)
+                later_drm = HLS.get_drm_or_probe(key, master, track, session, segment.init_section)
                if isinstance(later_drm, (Widevine, PlayReady)) and license_widevine:
                    license_widevine(later_drm, track_kid=later_drm.kid)
            progress(downloaded="[yellow]SKIPPED")
@@ -1152,22 +1157,18 @@ class HLS:
                if key is None:
                    encryption_data = None
                elif not encryption_data or encryption_data[0] != key:
-                    drm = HLS.get_drm(key, session)
-                    if isinstance(drm, (Widevine, PlayReady)):
-                        try:
-                            if map_data:
-                                track_kid = track.get_key_id(map_data[1])
-                            else:
-                                track_kid = None
-                            if not track_kid:
-                                track_kid = drm.kid
+                    try:
+                        init_kid = track.get_key_id(map_data[1]) if map_data and key.method != "AES-128" else None
+                        drm = HLS.get_drm(key, session, **({"kid": init_kid} if init_kid else {}))
+                        if isinstance(drm, (Widevine, PlayReady)):
+                            track_kid = init_kid or drm.kid
                            progress(downloaded="LICENSING")
                            license_widevine(drm, track_kid=track_kid)
                            progress(downloaded="[yellow]LICENSED")
-                        except Exception:  # noqa
-                            DOWNLOAD_CANCELLED.set()  # skip pending track downloads
-                            progress(downloaded="[red]FAILED")
-                            raise
+                    except Exception:  # noqa
+                        DOWNLOAD_CANCELLED.set()  # skip pending track downloads
+                        progress(downloaded="[red]FAILED")
+                        raise
                    if (
                        encryption_data
                        and isinstance(drm, (Widevine, PlayReady))
@@ -1462,6 +1463,7 @@ class HLS:
        keys: list[Union[m3u8.model.SessionKey, m3u8.model.Key]],
        chosen: DRM_T,
        session: Optional[Union[Session, RnetSession]] = None,
+        kid: Optional[UUID] = None,
    ) -> list[DRM_T]:
        """Build the Widevine or PlayReady DRM that the EXT-X-KEY entries offer next to ``chosen``.

@@ -1472,7 +1474,7 @@ class HLS:
        if not other_keys:
            return []
        try:
-            return [HLS.get_drm(other_keys[0], session)]
+            return [HLS.get_drm(other_keys[0], session, **({"kid": kid} if kid else {}))]
        except Exception as e:
            logging.getLogger("HLS").warning(f"Could not read the {system} EXT-X-KEY entry: {e}")
            return []
@@ -1482,14 +1484,15 @@ class HLS:
        master: M3U8,
        track: AnyTrack,
        session: Union[Session, RnetSession],
+        init_section: Optional[m3u8.model.InitializationSection] = None,
    ) -> Optional[UUID]:
        """
-        Extract the track's Key ID from its init segment (EXT-X-MAP).
+        Extract the track's Key ID from an init segment (EXT-X-MAP), by default the first in the playlist.

        Returns None if no init segment exists or KID extraction fails.
        The caller should fall back to drm.kid from the PSSH if this returns None.
        """
-        map_section = next((seg.init_section for seg in master.segments if seg.init_section), None)
+        map_section = init_section or next((seg.init_section for seg in master.segments if seg.init_section), None)
        if not map_section:
            return None

@@ -1503,6 +1506,10 @@ class HLS:
            map_res = session.get(url=map_uri, headers=headers)
            if map_res.ok:
                return track.get_key_id(map_res.content)
+            logging.getLogger("HLS").warning(
+                f"Could not read the Key ID from the init segment {safe_display_url(map_uri)}: "
+                f"HTTP {map_res.status_code}"
+            )
        # best-effort key-id probe over network + init parse; caller falls back without it
        except Exception as e:
            logging.getLogger("HLS").debug(f"Init map key-id probe failed for {map_uri}: {e!r}")
@@ -1550,6 +1557,7 @@ class HLS:
    def get_drm(
        key: Union[m3u8.model.SessionKey, m3u8.model.Key],
        session: Optional[Union[Session, RnetSession]] = None,
+        kid: Optional[UUID] = None,
    ) -> DRM_T:
        """
        Convert HLS EXT-X-KEY data to an initialized DRM object.
@@ -1558,6 +1566,8 @@ class HLS:
            key: The m3u8 `EXT-X-KEY` object.
            session: Optional HTTP session used to request AES-128 URIs.
                Useful to set headers, proxies, cookies, and so forth.
+            kid: The Key ID from the init segment. Widevine uses it only when the tag
+                and the PSSH both carry no real Key ID.

        Raises a NotImplementedError when unshackle cannot use the DRM system.
        """
@@ -1573,13 +1583,16 @@ class HLS:
            drm = Widevine(pssh=WV_PSSH.new(key_ids=[key.uri.split(",")[-1]], system_id=WV_PSSH.SystemId.Widevine))
        elif key.keyformat and key.keyformat.lower() == WidevineCdm.urn:
            extra_params = dict(key._extra_params)  # noqa
-            # a v0 PSSH carries no KID, leaving the tag's own KEYID as the only source
-            kid = extra_params.pop("keyid", None) or extra_params.pop("kid", None)
-            if isinstance(kid, str):
-                kid = kid.removeprefix("0x").removeprefix("0X")
+            pssh = WV_PSSH(key.uri.split(",")[-1])
+            if kid in Widevine.PLACEHOLDER_KIDS or any(k not in Widevine.PLACEHOLDER_KIDS for k in pssh.key_ids):
+                kid = None
+            # a v0 PSSH can carry no KID: the tag's own KEYID is the first source, the init segment the last
+            tag_kid = extra_params.pop("keyid", None) or extra_params.pop("kid", None)
+            if isinstance(tag_kid, str):
+                tag_kid = tag_kid.removeprefix("0x").removeprefix("0X")
            drm = Widevine(
-                pssh=WV_PSSH(key.uri.split(",")[-1]),
-                kid=kid,
+                pssh=pssh,
+                kid=tag_kid or kid,
                **extra_params,
            )
        elif key.keyformat and key.keyformat.lower() in {f"urn:uuid:{PR_PSSH.SYSTEM_ID}", "com.microsoft.playready"}:
@@ -1592,6 +1605,28 @@ class HLS:

        return drm

+    @staticmethod
+    def get_drm_or_probe(
+        key: Union[m3u8.model.SessionKey, m3u8.model.Key],
+        master: M3U8,
+        track: AnyTrack,
+        session: Optional[Union[Session, RnetSession]] = None,
+        init_section: Optional[m3u8.model.InitializationSection] = None,
+        probe_session: Any = None,
+    ) -> DRM_T:
+        """
+        Convert HLS EXT-X-KEY data to a DRM object, with the Key ID of the init segment when the key carries none.
+
+        The init segment request uses `probe_session` when the caller has no HTTP session that `get_drm` accepts.
+        """
+        try:
+            return HLS.get_drm(key, session)
+        except Widevine.Exceptions.KIDNotFound:
+            kid = HLS.get_track_kid_from_init(master, track, probe_session or session or Session(), init_section)
+            if not kid:
+                raise
+            return HLS.get_drm(key, session, kid=kid)
+
    @staticmethod
    def get_all_drm(
        keys: list[Union[m3u8.model.SessionKey, m3u8.model.Key]], proxy: Optional[str] = None
@@ -1617,6 +1652,10 @@ class HLS:
                drm_objects.append(drm)
            except NotImplementedError:
                unsupported_keys.append(key)
+            except Widevine.Exceptions.KIDNotFound:
+                logging.getLogger("HLS").warning(
+                    "Skipped an EXT-X-SESSION-KEY entry that carries no Key ID. The media playlist must supply the key."
+                )

        if not drm_objects and unsupported_keys:
            logging.debug(