Decouple and add configurability of crypto instance (#168)

* Decouple and add configurability of crypto instance

* Fix naming

Author: Sebastian Molenda

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
 

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

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()
 

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)