From 01416f6513246e5d4966b847fcecb337a0bb6c5a Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Sat, 6 Aug 2022 08:44:25 +0100
Subject: [PATCH] Cdm: Add a method to get keys from loaded license

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

diff --git a/pywidevine/cdm.py b/pywidevine/cdm.py
index 6bd345a..670440a 100644
--- a/pywidevine/cdm.py
+++ b/pywidevine/cdm.py
@@ -390,6 +390,39 @@ class Cdm:
 
         del session.context[licence.id.request_id]
 
+    def get_keys(self, session_id: bytes, type_: Optional[Union[int, str]] = None) -> list[Key]:
+        """
+        Get Keys from the loaded License message.
+
+        Parameters:
+            session_id: Session identifier.
+            type_: (optional) Key Type to filter by and return.
+
+        Raises:
+            InvalidSession: If the Session identifier is invalid.
+            TypeError: If the provided type_ is an unexpected value type.
+            ValueError: If the provided type_ is not a valid Key Type.
+        """
+        session = self._sessions.get(session_id)
+        if not session:
+            raise InvalidSession(f"Session identifier {session_id!r} is invalid.")
+
+        try:
+            if isinstance(type_, str):
+                type_ = License.KeyContainer.KeyType.Value(type_)
+            elif isinstance(type_, int):
+                License.KeyContainer.KeyType.Name(type_)  # only test
+            elif type_ is not None:
+                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}")
+
+        return [
+            key
+            for key in session.keys
+            if not type_ or key.type == License.KeyContainer.KeyType.Name(type_)
+        ]
+
     def decrypt(
         self,
         session_id: bytes,