Initial Commit

pyplayready/pssh.py

@@ -0,0 +1,87 @@
+import base64
+from typing import Union
+from uuid import UUID
+
+from construct import Struct, Int32ul, Int16ul, Array, this, Bytes, PaddedString, Switch, Int32ub, Const, Container
+
+
+class _PlayreadyPSSHStructs:
+    PSSHBox = Struct(
+        "length" / Int32ub,
+        "pssh" / Const(b"pssh"),
+        "fullbox" / Int32ub,
+        "system_id" / Bytes(16),
+        "data_length" / Int32ub,
+        "data" / Bytes(this.data_length)
+    )
+
+    PlayreadyObject = Struct(
+        "type" / Int16ul,
+        "length" / Int16ul,
+        "data" / Switch(
+            this.type,
+            {
+                1: PaddedString(this.length, "utf16")
+            },
+            default=Bytes(this.length)
+        )
+    )
+
+    PlayreadyHeader = Struct(
+        "length" / Int32ul,
+        "record_count" / Int16ul,
+        "records" / Array(this.record_count, PlayreadyObject)
+    )
+
+
+class PSSH:
+    SYSTEM_ID = UUID(hex="9a04f07998404286ab92e65be0885f95")
+
+    def __init__(
+            self,
+            data: Union[str, bytes]
+    ):
+        """Represents a PlayReady PSSH"""
+        if not data:
+            raise ValueError("Data must not be empty")
+
+        if isinstance(data, str):
+            try:
+                data = base64.b64decode(data)
+            except Exception as e:
+                raise Exception(f"Could not decode data as Base64, {e}")
+
+        try:
+            if self._is_playready_pssh_box(data):
+                pssh_box = _PlayreadyPSSHStructs.PSSHBox.parse(data)
+                if bool(self._is_utf_16(pssh_box.data)):
+                    self.wrm_headers = [pssh_box.data.decode("utf-16-le")]
+                elif bool(self._is_utf_16(pssh_box.data[6:])):
+                    self.wrm_headers = [pssh_box.data[6:].decode("utf-16-le")]
+                elif bool(self._is_utf_16(pssh_box.data[10:])):
+                    self.wrm_headers = [pssh_box.data[10:].decode("utf-16-le")]
+                else:
+                    self.wrm_headers = list(self._get_wrm_headers(_PlayreadyPSSHStructs.PlayreadyHeader.parse(pssh_box.data)))
+            elif bool(self._is_utf_16(data)):
+                self.wrm_headers = [data.decode("utf-16-le")]
+            elif bool(self._is_utf_16(data[6:])):
+                self.wrm_headers = [data[6:].decode("utf-16-le")]
+            elif bool(self._is_utf_16(data[10:])):
+                self.wrm_headers = [data[10:].decode("utf-16-le")]
+            else:
+                self.wrm_headers = list(self._get_wrm_headers(_PlayreadyPSSHStructs.PlayreadyHeader.parse(data)))
+        except Exception:
+            raise Exception("Could not parse data as a PSSH Box nor a PlayReadyHeader")
+
+    def _is_playready_pssh_box(self, data: bytes) -> bool:
+        return data[12:28] == self.SYSTEM_ID.bytes
+
+    @staticmethod
+    def _is_utf_16(data: bytes) -> bool:
+        return all(map(lambda i: data[i] == 0, range(1, len(data), 2)))
+
+    @staticmethod
+    def _get_wrm_headers(wrm_header: Container):
+        for record in wrm_header.records:
+            if record.type == 1:
+                yield record.data