fix(drm): load a bare PlayReady WRMHEADER from DASH, HLS and ISM manifests

PlayReady.from_pssh_b64 wraps a bare UTF-16LE WRMHEADER in a PlayReady Object and keeps that as pssh_b64. DASH and ISM log and skip a PlayReady entry they cannot parse.

tests/core/test_playready_bare_header.py

@@ -0,0 +1,112 @@
+"""A bare UTF-16LE WRMHEADER in place of a PSSH box or PlayReady Object loads as PlayReady DRM."""
+
+from __future__ import annotations
+
+import base64
+import logging
+from uuid import UUID
+
+import m3u8
+import pytest
+from lxml import etree
+from pyplayready.misc.exceptions import InvalidPssh
+from pyplayready.system.pssh import PSSH
+from pywidevine.pssh import PSSH as WV_PSSH
+
+from unshackle.core.drm import PlayReady, Widevine, drm_from_dict
+from unshackle.core.manifests.dash import DASH
+from unshackle.core.manifests.hls import HLS
+from unshackle.core.manifests.ism import ISM
+
+KID = UUID("00112233-4455-6677-8899-aabbccddeeff")
+KID_B64 = base64.b64encode(KID.bytes_le).decode()
+PR_URN = "urn:uuid:9a04f079-9840-4286-ab92-e65be0885f95"
+WV_URN = "urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed"
+
+
+def bare_header(version: str) -> str:
+    """Base64 of a bare WRMHEADER, with the KID where the given header version puts it."""
+    if version == "4.0.0.0":
+        data = f"<PROTECTINFO><KEYLEN>16</KEYLEN><ALGID>AESCTR</ALGID></PROTECTINFO><KID>{KID_B64}</KID>"
+    else:
+        data = f'<PROTECTINFO><KIDS><KID ALGID="AESCTR" VALUE="{KID_B64}"></KID></KIDS></PROTECTINFO>'
+    xml = (
+        f'<WRMHEADER xmlns="http://schemas.microsoft.com/DRM/2007/03/PlayReadyHeader" version="{version}">'
+        f"<DATA>{data}<LA_URL>https://example.invalid/pr</LA_URL></DATA></WRMHEADER>"
+    )
+    return base64.b64encode(xml.encode("utf-16-le")).decode()
+
+
+def pr_protection(child: str, value: str) -> etree._Element:
+    return etree.fromstring(
+        f'<ContentProtection xmlns:cenc="urn:mpeg:cenc:2013" xmlns:mspr="urn:microsoft:playready" '
+        f'schemeIdUri="{PR_URN}"><{child}>{value}</{child}></ContentProtection>'
+    )
+
+
+@pytest.mark.parametrize("version", ["4.0.0.0", "4.1.0.0", "4.2.0.0", "4.3.0.0"])
+@pytest.mark.parametrize("child", ["mspr:pro", "cenc:pssh"])
+def test_dash_loads_a_bare_header(child: str, version: str) -> None:
+    drm = DASH.get_drm([pr_protection(child, bare_header(version))])
+
+    assert len(drm) == 1
+    assert isinstance(drm[0], PlayReady)
+    assert KID in drm[0].kids
+    assert drm[0].pssh_b64
+    PSSH(drm[0].pssh_b64)
+    assert KID in drm_from_dict(drm[0].to_dict()).kids
+
+
+def test_wrapped_header_bytes_are_unchanged() -> None:
+    header = bare_header("4.0.0.0")
+    drm = PlayReady.from_pssh_b64(header)
+
+    assert drm.pssh.wrm_headers[0].dumps().encode("utf-16-le") == base64.b64decode(header)
+
+
+def test_dash_skips_an_unparseable_playready_entry_and_keeps_widevine(caplog: pytest.LogCaptureFixture) -> None:
+    widevine = etree.fromstring(
+        f'<ContentProtection xmlns:cenc="urn:mpeg:cenc:2013" schemeIdUri="{WV_URN}">'
+        f"<cenc:pssh>{WV_PSSH.new(system_id=WV_PSSH.SystemId.Widevine, key_ids=[KID]).dumps()}</cenc:pssh>"
+        "</ContentProtection>"
+    )
+    garbage = pr_protection("mspr:pro", base64.b64encode(b"\x00\x00not a header").decode())
+
+    with caplog.at_level(logging.WARNING, logger="DASH"):
+        drm = DASH.get_drm([widevine, garbage])
+
+    assert [type(d) for d in drm] == [Widevine]
+    assert "Skipping unparseable PlayReady ContentProtection" in caplog.text
+
+
+@pytest.mark.parametrize(
+    "data",
+    [b"\x00\x00not a header", "<WRMHEADER/>".encode("utf-8")],
+    ids=["garbage", "utf-8-header"],
+)
+def test_other_bad_data_keeps_the_original_error(data: bytes) -> None:
+    with pytest.raises(InvalidPssh):
+        PlayReady.from_pssh_b64(base64.b64encode(data).decode())
+
+
+def test_hls_loads_a_bare_header() -> None:
+    key = m3u8.loads(
+        "#EXTM3U\n"
+        f'#EXT-X-KEY:METHOD=SAMPLE-AES-CTR,URI="data:text/plain;charset=UTF-16;base64,{bare_header("4.0.0.0")}",'
+        'KEYFORMAT="com.microsoft.playready",KEYFORMATVERSIONS="1"\n'
+        "#EXTINF:4.0,\nseg0.mp4\n"
+    ).keys[0]
+
+    drm = HLS.get_drm(key)
+
+    assert isinstance(drm, PlayReady)
+    assert KID in drm.kids
+
+
+def test_ism_loads_a_bare_header() -> None:
+    header = etree.fromstring(f'<ProtectionHeader SystemID="{PR_URN[9:]}">{bare_header("4.0.0.0")}</ProtectionHeader>')
+
+    drm = ISM.get_drm([header])
+
+    assert len(drm) == 1
+    assert KID in drm[0].kids

unshackle/core/drm/playready.py

@@ -15,6 +15,7 @@ import m3u8
 from construct import Container
 from pymp4.parser import Box
 from pyplayready.cdm import Cdm as PlayReadyCdm
+from pyplayready.misc.exceptions import InvalidPssh
 from pyplayready.system.pssh import PSSH
 from pyplayready.system.wrmheader import WRMHeader
 from requests import Session
@@ -170,6 +171,26 @@ class PlayReady:
 
         return []
 
+    @classmethod
+    def from_pssh_b64(cls, pssh_b64: str, kid: Optional[UUID] = None) -> PlayReady:
+        """Load a base64 PSSH box or PlayReady Object, or a bare UTF-16LE WRMHEADER that some packagers emit.
+
+        It wraps a bare WRMHEADER in a PlayReady Object and keeps that as `pssh_b64`, so remote CDMs and
+        the REST API can parse it again.
+        """
+        try:
+            pssh = PSSH(pssh_b64)
+        except InvalidPssh:
+            try:
+                raw = base64.b64decode(pssh_b64)
+            except ValueError:
+                raw = b""
+            if not raw.startswith(b"<\x00"):
+                raise
+            pssh_b64 = cls.pro_b64_for(WRMHeader(raw))
+            pssh = PSSH(pssh_b64)
+        return cls(pssh=pssh, kid=kid, pssh_b64=pssh_b64)
+
     @classmethod
     def from_track(cls, track: AnyTrack, session: Optional[Session] = None) -> PlayReady:
         if not session:

unshackle/core/manifests/dash.py

@@ -20,7 +20,6 @@ from zlib import crc32
 import requests
 from langcodes import Language, tag_is_valid
 from lxml.etree import Element, ElementTree, XMLSyntaxError
-from pyplayready.system.pssh import PSSH as PR_PSSH
 from pywidevine.cdm import Cdm as WidevineCdm
 from pywidevine.pssh import PSSH
 from requests import Session
@@ -1355,7 +1354,6 @@ class DASH:
                 )
                 if not pr_pssh_b64:
                     continue
-                pr_pssh = PR_PSSH(pr_pssh_b64)
                 kid_b64 = protection.findtext("kid")
                 kid = None
                 if kid_b64:
@@ -1364,7 +1362,10 @@ class DASH:
                     except ValueError:
                         kid = None
 
-                drm.append(PlayReady(pssh=pr_pssh, kid=kid, pssh_b64=pr_pssh_b64))
+                try:
+                    drm.append(PlayReady.from_pssh_b64(pr_pssh_b64, kid=kid))
+                except Exception as e:
+                    logging.getLogger("DASH").warning(f"Skipping unparseable PlayReady ContentProtection: {e!r}")
 
             elif urn == ClearKeyCENC.urn:
                 # W3C EME ClearKey (org.w3.clearkey): match the scheme UUID alone,

unshackle/core/manifests/hls.py

@@ -1578,10 +1578,7 @@ class HLS:
                 **extra_params,
             )
         elif key.keyformat and key.keyformat.lower() in {f"urn:uuid:{PR_PSSH.SYSTEM_ID}", "com.microsoft.playready"}:
-            drm = PlayReady(
-                pssh=PR_PSSH(key.uri.split(",")[-1]),
-                pssh_b64=key.uri.split(",")[-1],
-            )
+            drm = PlayReady.from_pssh_b64(key.uri.split(",")[-1])
         else:
             raise NotImplementedError(f"The key system is not supported: {key}")
 

unshackle/core/manifests/ism.py

@@ -15,7 +15,6 @@ from typing import Any, BinaryIO, Optional, Union
 import requests
 from langcodes import Language, tag_is_valid
 from lxml.etree import Element
-from pyplayready.system.pssh import PSSH as PR_PSSH
 from pywidevine.pssh import PSSH
 from requests import Session
 
@@ -104,11 +103,9 @@ class ISM:
                 drm.append(Widevine(pssh=pssh, kid=kid))
             elif system_id == "9a04f079-9840-4286-ab92-e65be0885f95":
                 try:
-                    pr_pssh = PR_PSSH(data)
+                    drm.append(PlayReady.from_pssh_b64(data))
                 except Exception as e:
                     logging.getLogger("ISM").warning(f"Skipping unparseable PlayReady ProtectionHeader: {e!r}")
-                    continue
-                drm.append(PlayReady(pssh=pr_pssh, pssh_b64=data))
         return drm
 
     @staticmethod

unshackle/core/tracks/track.py

@@ -1174,7 +1174,7 @@ class Track:
                     drm_list.append(drm)
                 elif fmt in {f"urn:uuid:{PR_PSSH.SYSTEM_ID}", "com.microsoft.playready"}:
                     pssh_b64 = key.uri.split(",")[-1]
-                    drm = PlayReady(pssh=PR_PSSH(pssh_b64), pssh_b64=pssh_b64)
+                    drm = PlayReady.from_pssh_b64(pssh_b64)
                     drm_list.append(drm)
 
             if drm_list: