Integrate crypto_module with pubnub

pubnub/crypto.py

@@ -159,7 +159,7 @@ def encrypt_file(self, file_data, cryptor_id: str = None):
         cryptor_id = self._validate_cryptor_id(cryptor_id)
         crypto_payload = self.cryptor_map[cryptor_id].encrypt(file_data)
         header = self.encode_header(cryptor_id=cryptor_id, cryptor_data=crypto_payload['cryptor_data'])
-        return b64encode(header + crypto_payload['data']).decode()
+        return header + crypto_payload['data']
 
     def _get_cryptor(self, cryptor_id):
         if not cryptor_id or cryptor_id not in self.cryptor_map:
@@ -234,18 +234,20 @@ def decode_header(self, header: bytes) -> Union[None, CryptoHeader]:
 
 
 class LegacyCryptoModule(PubNubCryptoModule):
-    def __init__(self, cipher_key, use_random_iv) -> None:
+    def __init__(self, config) -> None:
         cryptor_map = {
-            PubNubLegacyCryptor.CRYPTOR_ID: PubNubLegacyCryptor(cipher_key, use_random_iv),
-            PubNubAesCbcCryptor.CRYPTOR_ID: PubNubAesCbcCryptor(),
+            PubNubLegacyCryptor.CRYPTOR_ID: PubNubLegacyCryptor(
+                config.cipher_key, config.use_random_initialization_vector),
+            PubNubAesCbcCryptor.CRYPTOR_ID: PubNubAesCbcCryptor(config.cipher_key),
         }
         super().__init__(cryptor_map, PubNubLegacyCryptor)
 
 
 class AesCbcCryptoModule(PubNubCryptoModule):
-    def __init__(self, cipher_key, use_random_iv) -> None:
+    def __init__(self, config) -> None:
         cryptor_map = {
-            PubNubLegacyCryptor.CRYPTOR_ID: PubNubLegacyCryptor(cipher_key, use_random_iv),
-            PubNubAesCbcCryptor.CRYPTOR_ID: PubNubAesCbcCryptor(cipher_key),
+            PubNubLegacyCryptor.CRYPTOR_ID: PubNubLegacyCryptor(
+                config.cipher_key, config.use_random_initialization_vector),
+            PubNubAesCbcCryptor.CRYPTOR_ID: PubNubAesCbcCryptor(config.cipher_key),
         }
         super().__init__(cryptor_map, PubNubAesCbcCryptor)

pubnub/crypto_core.py

@@ -6,6 +6,7 @@
 from abc import abstractmethod
 from Cryptodome.Cipher import AES
 from Cryptodome.Util.Padding import pad, unpad
+from pubnub.exceptions import PubNubException
 
 
 class PubNubCrypto:
@@ -51,6 +52,8 @@ class PubNubLegacyCryptor(PubNubCryptor):
     Initial16bytes = b'0123456789012345'
 
     def __init__(self, cipher_key, use_random_iv=False, cipher_mode=AES.MODE_CBC, fallback_cipher_mode=None):
+        if not cipher_key:
+            raise PubNubException('No cipher_key passed')
         self.cipher_key = cipher_key
         self.use_random_iv = use_random_iv
         self.mode = cipher_mode
@@ -62,7 +65,6 @@ def encrypt(self, msg, *, key=None, use_random_iv=None):
 
         secret = self.get_secret(key)
         initialization_vector = self.get_initialization_vector(use_random_iv)
-
         cipher = AES.new(bytes(secret[0:32], 'utf-8'), self.mode, initialization_vector)
         encrypted_message = cipher.encrypt(self.pad(msg))
         msg_with_iv = self.append_random_iv(encrypted_message, use_random_iv, initialization_vector)
@@ -98,6 +100,8 @@ def append_random_iv(self, message, use_random_iv, initialization_vector):
             return message
 
     def extract_random_iv(self, message, use_random_iv):
+        if not isinstance(message, bytes):
+            message = bytes(message, 'utf-8')
         if use_random_iv:
             return message[0:16], message[16:]
         else:

pubnub/endpoints/file_operations/download_file.py

@@ -4,6 +4,7 @@
 from pubnub.models.consumer.file import PNDownloadFileResult
 from pubnub.request_handlers.requests_handler import RequestsRequestHandler
 from pubnub.endpoints.file_operations.get_file_url import GetFileDownloadUrl
+from warnings import warn
 
 
 class DownloadFileNative(FileOperationEndpoint):
@@ -16,6 +17,7 @@ def __init__(self, pubnub):
         self._cipher_key = None
 
     def cipher_key(self, cipher_key):
+        warn('Deprecated: Usage of local cipher_keys is discouraged. Use pnconfiguration.cipher_key instead')
         self._cipher_key = cipher_key
         return self
 
@@ -40,10 +42,10 @@ def file_name(self, file_name):
         return self
 
     def decrypt_payload(self, data):
-        return PubNubFileCrypto(self._pubnub.config).decrypt(
-            self._cipher_key or self._pubnub.config.cipher_key,
-            data
-        )
+        if self._cipher_key:
+            return PubNubFileCrypto(self._pubnub.config).decrypt(self._cipher_key, data)
+        else:
+            return self._pubnub.crypto.decrypt_file(data)
 
     def validate_params(self):
         self.validate_subscribe_key()

pubnub/endpoints/file_operations/publish_file_message.py

@@ -3,6 +3,8 @@
 from pubnub import utils
 from pubnub.models.consumer.file import PNPublishFileMessageResult
 from pubnub.endpoints.mixins import TimeTokenOverrideMixin
+from pubnub.crypto import PubNubCryptodome
+from warnings import warn
 
 
 class PublishFileMessage(FileOperationEndpoint, TimeTokenOverrideMixin):
@@ -30,6 +32,7 @@ def should_store(self, should_store):
         return self
 
     def cipher_key(self, cipher_key):
+        warn('Deprecated: Usage of local cipher_keys is discouraged. Use pnconfiguration.cipher_key instead')
         self._cipher_key = cipher_key
         return self
 
@@ -50,11 +53,10 @@ def file_name(self, file_name):
         return self
 
     def _encrypt_message(self, message):
-        if self._cipher_key or self._pubnub.config.cipher_key:
-            return self._pubnub.config.crypto.encrypt(
-                self._cipher_key or self._pubnub.config.cipher_key,
-                utils.write_value_as_string(message)
-            )
+        if self._cipher_key:
+            return PubNubCryptodome(self._pubnub.config).encrypt(self._cipher_key, utils.write_value_as_string(message))
+        elif self._pubnub.config.cipher_key:
+            return self._pubnub.crypto.encrypt(utils.write_value_as_string(message))
         else:
             return message
 

pubnub/endpoints/file_operations/send_file.py

@@ -7,6 +7,7 @@
 from pubnub.endpoints.file_operations.fetch_upload_details import FetchFileUploadS3Data
 from pubnub.request_handlers.requests_handler import RequestsRequestHandler
 from pubnub.endpoints.mixins import TimeTokenOverrideMixin
+from warnings import warn
 
 
 class SendFileNative(FileOperationEndpoint, TimeTokenOverrideMixin):
@@ -107,6 +108,7 @@ def file_name(self, file_name):
         return self
 
     def cipher_key(self, cipher_key):
+        warn('Deprecated: Usage of local cipher_keys is discouraged. Use pnconfiguration.cipher_key instead')
         self._cipher_key = cipher_key
         return self
 

pubnub/pnconfiguration.py

@@ -1,7 +1,7 @@
 from Cryptodome.Cipher import AES
 from pubnub.enums import PNHeartbeatNotificationOptions, PNReconnectionPolicy
 from pubnub.exceptions import PubNubException
-from pubnub.crypto import PubNubCrypto
+from pubnub.crypto import PubNubCrypto, LegacyCryptoModule
 
 
 class PNConfiguration(object):
@@ -11,6 +11,7 @@ class PNConfiguration(object):
     RECONNECTION_INTERVAL = 3
     RECONNECTION_MIN_EXPONENTIAL_BACKOFF = 1
     RECONNECTION_MAX_EXPONENTIAL_BACKOFF = 32
+    DEFAULT_CRYPTO_MODULE = LegacyCryptoModule
 
     def __init__(self):
         # TODO: add validation
@@ -46,6 +47,7 @@ def __init__(self):
         self._heartbeat_interval = PNConfiguration.DEFAULT_HEARTBEAT_INTERVAL
         self.cryptor = None
         self.file_cryptor = None
+        self._crypto_module = None
 
     def validate(self):
         PNConfiguration.validate_not_empty_string(self.uuid)
@@ -124,6 +126,16 @@ def file_crypto(self) -> PubNubCrypto:
 
         return self.file_crypto_instance
 
+    @property
+    def crypto_module(self):
+        if not self._crypto_module:
+            self._crypto_module = self.DEFAULT_CRYPTO_MODULE(self)
+        return self._crypto_module
+
+    @crypto_module.setter
+    def crypto_module(self, crypto_module: PubNubCrypto):
+        self._crypto_module = crypto_module
+
     @property
     def port(self):
         return 443 if self.ssl == "https" else 80

pubnub/pubnub_core.py

@@ -1,5 +1,6 @@
 import logging
 import time
+from warnings import warn
 from pubnub.endpoints.entities.membership.add_memberships import AddSpaceMembers, AddUserSpaces
 from pubnub.endpoints.entities.membership.update_memberships import UpdateSpaceMembers, UpdateUserSpaces
 from pubnub.endpoints.entities.membership.fetch_memberships import FetchSpaceMemberships, FetchUserMemberships
@@ -121,6 +122,10 @@ def sdk_platform(self):
     def uuid(self):
         return self.config.uuid
 
+    @property
+    def crypto(self):
+        return self.config.crypto_module
+
     def add_listener(self, listener):
         self._validate_subscribe_manager_enabled()
 
@@ -326,9 +331,11 @@ def publish_file_message(self):
         return PublishFileMessage(self)
 
     def decrypt(self, cipher_key, file):
+        warn('Deprecated: Usage of decrypt with cipher key will be removed. Use PubNub.crypto.decrypt instead')
         return self.config.file_crypto.decrypt(cipher_key, file)
 
     def encrypt(self, cipher_key, file):
+        warn('Deprecated: Usage of encrypt with cipher key will be removed. Use PubNub.crypto.encrypt instead')
         return self.config.file_crypto.encrypt(cipher_key, file)
 
     @staticmethod

tests/functional/event_engine/test_subscribe.py

@@ -126,15 +126,15 @@ async def test_handshake_failed_reconnect():
     config.subscribe_key = 'totally-fake-key'
     config.enable_subscribe = True
     config.reconnect_policy = PNReconnectionPolicy.EXPONENTIAL
-    config.maximum_reconnection_retries = 5
-    config.subscribe_request_timeout = 2
+    config.maximum_reconnection_retries = 2
+    config.subscribe_request_timeout = 1
 
     callback = TestCallback()
 
     pubnub = PubNubAsyncio(config, subscription_manager=EventEngineSubscriptionManager)
     pubnub.add_listener(callback)
     pubnub.subscribe().channels('foo').execute()
-    await asyncio.sleep(16)
+    await asyncio.sleep(7)
     assert pubnub._subscription_manager.event_engine.get_state_name() == states.HandshakeReconnectingState.__name__
     await asyncio.sleep(1)
 

tests/integrational/asyncio/test_file_upload.py

@@ -3,7 +3,7 @@
 from unittest.mock import patch
 from pubnub.pubnub_asyncio import PubNubAsyncio
 from tests.integrational.vcr_helper import pn_vcr
-from tests.helper import pnconf_file_copy
+from tests.helper import pnconf_file_copy, pnconf_enc_env_copy
 from pubnub.endpoints.file_operations.publish_file_message import PublishFileMessage
 from pubnub.models.consumer.file import (
     PNSendFileResult, PNGetFilesResult, PNDownloadFileResult,
@@ -86,12 +86,12 @@ async def test_send_and_download_file(event_loop, file_for_upload):
 
 
 @pn_vcr.use_cassette(
-    "tests/integrational/fixtures/asyncio/file_upload/send_and_download_encrypted_file.yaml",
-    filter_query_parameters=['uuid', 'l_file', 'pnsdk']
+    "tests/integrational/fixtures/asyncio/file_upload/send_and_download_encrypted_file_cipher_key.json",
+    filter_query_parameters=['uuid', 'l_file', 'pnsdk'], serializer='pn_json'
 )
 @pytest.mark.asyncio
-async def test_send_and_download_file_encrypted(event_loop, file_for_upload, file_upload_test_data):
-    pubnub = PubNubAsyncio(pnconf_file_copy(), custom_event_loop=event_loop)
+async def test_send_and_download_file_encrypted_cipher_key(event_loop, file_for_upload, file_upload_test_data):
+    pubnub = PubNubAsyncio(pnconf_enc_env_copy(), custom_event_loop=event_loop)
 
     with patch("pubnub.crypto.PubNubCryptodome.get_initialization_vector", return_value="knightsofni12345"):
         envelope = await send_file(pubnub, file_for_upload, cipher_key="test")
@@ -107,6 +107,27 @@ async def test_send_and_download_file_encrypted(event_loop, file_for_upload, fil
         await pubnub.stop()
 
 
+@pn_vcr.use_cassette(
+    "tests/integrational/fixtures/asyncio/file_upload/send_and_download_encrypted_file_crypto_module.json",
+    filter_query_parameters=['uuid', 'l_file', 'pnsdk'], serializer='pn_json'
+)
+@pytest.mark.asyncio
+async def test_send_and_download_encrypted_file_crypto_module(event_loop, file_for_upload, file_upload_test_data):
+    pubnub = PubNubAsyncio(pnconf_enc_env_copy(), custom_event_loop=event_loop)
+
+    with patch("pubnub.crypto_core.PubNubLegacyCryptor.get_initialization_vector", return_value=b"knightsofni12345"):
+        envelope = await send_file(pubnub, file_for_upload)
+        download_envelope = await pubnub.download_file().\
+            channel(CHANNEL).\
+            file_id(envelope.result.file_id).\
+            file_name(envelope.result.name).\
+            future()
+
+        assert isinstance(download_envelope.result, PNDownloadFileResult)
+        assert download_envelope.result.data == bytes(file_upload_test_data["FILE_CONTENT"], "utf-8")
+        await pubnub.stop()
+
+
 @pn_vcr.use_cassette(
     "tests/integrational/fixtures/asyncio/file_upload/get_file_url.yaml",
     filter_query_parameters=['uuid', 'l_file', 'pnsdk']