From 3499c0cf4dd19c84a5d172620309273d6fc8bca0 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Sat, 6 Aug 2022 10:03:35 +0100
Subject: [PATCH] RemoteCdm: Implement get_keys()

---
 pywidevine/remotecdm.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/pywidevine/remotecdm.py b/pywidevine/remotecdm.py
index 94cbcd2..37b0cc5 100644
--- a/pywidevine/remotecdm.py
+++ b/pywidevine/remotecdm.py
@@ -248,5 +248,38 @@ class RemoteCdm(Cdm):
             for key in r["keys"]
         ]
 
+    def get_keys(self, session_id: bytes, type_: Optional[Union[int, str]] = None) -> list[Key]:
+        try:
+            if isinstance(type_, str):
+                License.KeyContainer.KeyType.Value(type_)  # only test
+            elif isinstance(type_, int):
+                type_ = License.KeyContainer.KeyType.Name(type_)
+            elif type_ is None:
+                type_ = "ALL"
+            else:
+                raise TypeError(f"Expected type_ to be a {License.KeyContainer.KeyType} or int, not {type_!r}")
+        except ValueError as e:
+            raise ValueError(f"Could not parse type_ as a {License.KeyContainer.KeyType}, {e}")
+
+        r = self.__session.post(
+            url=f"{self.host}/{self.device_name}/get_keys/{type_}",
+            json={
+                "session_id": session_id.hex()
+            }
+        ).json()
+        if r["status"] != 200:
+            raise ValueError(f"Could not get {type_} Keys, {r['message']} [{r['status']}]")
+        r = r["data"]
+
+        return [
+            Key(
+                type_=key["type"],
+                kid=Key.kid_to_uuid(bytes.fromhex(key["key_id"])),
+                key=bytes.fromhex(key["key"]),
+                permissions=key["permissions"]
+            )
+            for key in r["keys"]
+        ]
+
 
 __ALL__ = (RemoteCdm,)