From f8de381e0165e89c8b1b98b791358ad6c923db85 Mon Sep 17 00:00:00 2001
From: hyugogirubato <65763543+hyugogirubato@users.noreply.github.com>
Date: Sun, 20 Oct 2024 16:06:43 +0200
Subject: [PATCH] Fix unknown func display

---
 keydive/cdm.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/keydive/cdm.py b/keydive/cdm.py
index ccca9f1..e8883d9 100644
--- a/keydive/cdm.py
+++ b/keydive/cdm.py
@@ -14,6 +14,8 @@ from pywidevine.device import Device, DeviceTypes
 from pywidevine.license_protocol_pb2 import (SignedMessage, LicenseRequest, ClientIdentification, SignedDrmCertificate,
                                              DrmCertificate, EncryptedClientIdentification)
 
+from keydive.constants import OEM_CRYPTO_API
+
 
 class Cdm:
     """
@@ -22,6 +24,13 @@ class Cdm:
     """
 
     def __init__(self):
+        """
+        Initializes the Cdm object, setting up a logger and containers for client IDs and private keys.
+
+        Attributes:
+            client_id (dict[int, ClientIdentification]): Stores client identification info mapped by key modulus.
+            private_key (dict[int, RsaKey]): Stores private keys mapped by key modulus.
+        """
         self.logger = logging.getLogger(self.__class__.__name__)
         # https://github.com/devine-dl/pywidevine
         self.client_id: dict[int, ClientIdentification] = {}
@@ -96,17 +105,22 @@ class Cdm:
         except Exception as e:
             self.logger.debug('Failed to set challenge data: %s', e)
 
-    def set_private_key(self, data: bytes) -> None:
+    def set_private_key(self, data: bytes, name: str) -> None:
         """
         Sets the private key from the provided data.
 
         Args:
             data (bytes): The private key data.
+            name (str): The name of the function.
         """
         try:
             key = RSA.import_key(data)
             if key.n not in self.private_key:
                 self.logger.debug('Receive private key: \n\n%s\n', key.exportKey('PEM').decode('utf-8'))
+
+                if name not in OEM_CRYPTO_API:
+                    self.logger.warning(f'The function "{name}" does not belong to the referenced functions. Communicate it to the developer to improve the tool.')
+
             self.private_key[key.n] = key
         except Exception as e:
             self.logger.debug('Failed to set private key: %s', e)