Decouple and add configurability of crypto instance (#168)

* Decouple and add configurability of crypto instance * Fix naming
pubnub · Sep 5, 2023 · 701ece7 · 701ece7
1 parent 1b44c1f
commit 701ece7
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 7 deletions.
diff --git a/pubnub/crypto.py b/pubnub/crypto.py
@@ -16,7 +16,7 @@ class PubNubCryptodome(PubNubCrypto):
     fallback_mode = None
 
     def __init__(self, pubnub_config):
-        self.pubnub_configuration = pubnub_config
+        super().__init__(pubnub_config)
         self.mode = pubnub_config.cipher_mode
         self.fallback_mode = pubnub_config.fallback_cipher_mode
 

diff --git a/pubnub/crypto_core.py b/pubnub/crypto_core.py
@@ -2,6 +2,9 @@
 
 
 class PubNubCrypto:
+    def __init__(self, pubnub_config):
+        self.pubnub_configuration = pubnub_config
+
     @abstractmethod
     def encrypt(self, key, msg):
         pass

diff --git a/pubnub/pnconfiguration.py b/pubnub/pnconfiguration.py
@@ -1,6 +1,7 @@
 from Cryptodome.Cipher import AES
 from pubnub.enums import PNHeartbeatNotificationOptions, PNReconnectionPolicy
 from pubnub.exceptions import PubNubException
+from pubnub.crypto import PubNubCrypto
 
 
 class PNConfiguration(object):
@@ -43,6 +44,8 @@ def __init__(self):
         self.heartbeat_default_values = True
         self._presence_timeout = PNConfiguration.DEFAULT_PRESENCE_TIMEOUT
         self._heartbeat_interval = PNConfiguration.DEFAULT_HEARTBEAT_INTERVAL
+        self.cryptor = None
+        self.file_cryptor = None
 
     def validate(self):
         PNConfiguration.validate_not_empty_string(self.uuid)
@@ -102,15 +105,20 @@ def crypto(self):
         return self.crypto_instance
 
     def _init_cryptodome(self):
-        from .crypto import PubNubCryptodome
-        self.crypto_instance = PubNubCryptodome(self)
+        if not self.cryptor:
+            from pubnub.crypto import PubNubCryptodome
+            self.cryptor = PubNubCryptodome
+        self.crypto_instance = self.cryptor(self)
 
     def _init_file_crypto(self):
         from .crypto import PubNubFileCrypto
-        self.file_crypto_instance = PubNubFileCrypto(self)
+        if not self.file_cryptor:
+            from pubnub.crypto import PubNubFileCrypto
+            self.file_cryptor = PubNubFileCrypto
+        self.file_crypto_instance = self.file_cryptor(self)
 
     @property
-    def file_crypto(self):
+    def file_crypto(self) -> PubNubCrypto:
         if not self.file_crypto_instance:
             self._init_file_crypto()
 

diff --git a/tests/unit/test_crypto.py b/tests/unit/test_crypto.py
@@ -1,6 +1,7 @@
 from pubnub.pubnub import PubNub
-from pubnub.crypto import PubNubCryptodome
-from tests.helper import pnconf_file_copy, hardcoded_iv_config_copy
+from pubnub.crypto import PubNubCryptodome, PubNubCrypto
+from tests.helper import pnconf_file_copy, hardcoded_iv_config_copy, pnconf_env_copy
+
 
 crypto = PubNubCryptodome(pnconf_file_copy())
 crypto_hardcoded_iv = PubNubCryptodome(hardcoded_iv_config_copy())
@@ -62,3 +63,19 @@ def test_encrypt_and_decrypt_file(self, file_for_upload, file_upload_test_data):
 
         decrypted_file = pubnub.decrypt(KEY, encrypted_file)
         assert file_upload_test_data["FILE_CONTENT"] == decrypted_file.decode("utf-8")
+
+
+class TestPubNubCryptoInterface:
+    def test_get_default_crypto(self):
+        config = pnconf_env_copy()
+        assert isinstance(config.crypto, PubNubCrypto)
+        assert isinstance(config.crypto, PubNubCryptodome)
+
+    def test_get_custom_crypto(self):
+        class CustomCryptor(PubNubCrypto):
+            pass
+
+        config = pnconf_env_copy()
+        config.cryptor = CustomCryptor
+        assert isinstance(config.crypto, PubNubCrypto)
+        assert isinstance(config.crypto, CustomCryptor)