Initial Commit

pyplayready/device.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import base64
+from enum import IntEnum
+from pathlib import Path
+from typing import Union, Any
+
+from construct import Struct, Const, Int8ub, Bytes, this, Int32ub
+
+from pyplayready.bcert import CertificateChain
+from pyplayready.ecc_key import ECCKey
+
+
+class SecurityLevel(IntEnum):
+    SL150 = 150
+    SL2000 = 2000
+    SL3000 = 3000
+
+
+class _DeviceStructs:
+    magic = Const(b"PRD")
+
+    v1 = Struct(
+        "signature" / magic,
+        "version" / Int8ub,
+        "group_key_length" / Int32ub,
+        "group_key" / Bytes(this.group_key_length),
+        "group_certificate_length" / Int32ub,
+        "group_certificate" / Bytes(this.group_certificate_length)
+    )
+
+    v2 = Struct(
+        "signature" / magic,
+        "version" / Int8ub,
+        "group_certificate_length" / Int32ub,
+        "group_certificate" / Bytes(this.group_certificate_length),
+        "encryption_key" / Bytes(96),
+        "signing_key" / Bytes(96),
+    )
+
+
+class Device:
+    CURRENT_STRUCT = _DeviceStructs.v2
+
+    def __init__(
+            self,
+            *_: Any,
+            group_certificate: Union[str, bytes],
+            encryption_key: Union[str, bytes],
+            signing_key: Union[str, bytes],
+            **__: Any
+    ):
+        if isinstance(group_certificate, str):
+            group_certificate = base64.b64decode(group_certificate)
+        if not isinstance(group_certificate, bytes):
+            raise ValueError(f"Expecting Bytes or Base64 input, got {group_certificate!r}")
+
+        if isinstance(encryption_key, str):
+            encryption_key = base64.b64decode(encryption_key)
+        if not isinstance(encryption_key, bytes):
+            raise ValueError(f"Expecting Bytes or Base64 input, got {encryption_key!r}")
+        if isinstance(signing_key, str):
+            signing_key = base64.b64decode(signing_key)
+        if not isinstance(signing_key, bytes):
+            raise ValueError(f"Expecting Bytes or Base64 input, got {signing_key!r}")
+
+        self.group_certificate = CertificateChain.loads(group_certificate)
+        self.encryption_key = ECCKey.loads(encryption_key)
+        self.signing_key = ECCKey.loads(signing_key)
+        self.security_level = self.group_certificate.get_security_level()
+
+    @classmethod
+    def loads(cls, data: Union[str, bytes]) -> Device:
+        if isinstance(data, str):
+            data = base64.b64decode(data)
+        if not isinstance(data, bytes):
+            raise ValueError(f"Expecting Bytes or Base64 input, got {data!r}")
+        return cls(**cls.CURRENT_STRUCT.parse(data))
+
+    @classmethod
+    def load(cls, path: Union[Path, str]) -> Device:
+        if not isinstance(path, (Path, str)):
+            raise ValueError(f"Expecting Path object or path string, got {path!r}")
+        with Path(path).open(mode="rb") as f:
+            return cls.loads(f.read())
+
+    def dumps(self) -> bytes:
+        return self.CURRENT_STRUCT.build(dict(
+            version=2,
+            group_certificate_length=len(self.group_certificate.dumps()),
+            group_certificate=self.group_certificate.dumps(),
+            encryption_key=self.encryption_key.dumps(),
+            signing_key=self.signing_key.dumps()
+        ))
+
+    def dump(self, path: Union[Path, str]) -> None:
+        if not isinstance(path, (Path, str)):
+            raise ValueError(f"Expecting Path object or path string, got {path!r}")
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_bytes(self.dumps())
+
+    def get_name(self):
+        name = f"{self.group_certificate.get_name()}_sl{self.group_certificate.get_security_level()}"
+        return ''.join(char for char in name if char.isascii()).strip().lower().replace(" ", "_")