From 191fbea66554df545a69790b1fbab27e6a9320fa Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Thu, 30 Sep 2021 10:10:20 +0300 Subject: [PATCH] Refactor password based encryptor into PasswordBasedStringEncryptor and PasswordBasedByteEncryptor This change allows to encrypt strings and bytes without any additional conversion done on the caller --- .../cc/resources/configuration_export.py | 4 +- .../cc/resources/configuration_import.py | 4 +- .../cc/server_utils/encryption/__init__.py | 7 ++-- ...n.py => password_based_byte_encryption.py} | 31 +++++---------- .../password_based_string_encryption.py | 38 +++++++++++++++++++ .../cc/resources/test_configuration_import.py | 4 +- .../test_password_based_encryption.py | 13 ++++--- 7 files changed, 64 insertions(+), 37 deletions(-) rename monkey/monkey_island/cc/server_utils/encryption/{password_based_encryption.py => password_based_byte_encryption.py} (64%) create mode 100644 monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryption.py diff --git a/monkey/monkey_island/cc/resources/configuration_export.py b/monkey/monkey_island/cc/resources/configuration_export.py index c550acc7d..111cfa177 100644 --- a/monkey/monkey_island/cc/resources/configuration_export.py +++ b/monkey/monkey_island/cc/resources/configuration_export.py @@ -4,7 +4,7 @@ import flask_restful from flask import request from monkey_island.cc.resources.auth.auth import jwt_required -from monkey_island.cc.server_utils.encryption import PasswordBasedEncryptor +from monkey_island.cc.server_utils.encryption import PasswordBasedStringEncryptor from monkey_island.cc.services.config import ConfigService @@ -21,7 +21,7 @@ class ConfigurationExport(flask_restful.Resource): password = data["password"] plaintext_config = json.dumps(plaintext_config) - pb_encryptor = PasswordBasedEncryptor(password) + pb_encryptor = PasswordBasedStringEncryptor(password) config_export = pb_encryptor.encrypt(plaintext_config) return {"config_export": config_export, "encrypted": should_encrypt} diff --git a/monkey/monkey_island/cc/resources/configuration_import.py b/monkey/monkey_island/cc/resources/configuration_import.py index 6c0575e94..3a66a2ed0 100644 --- a/monkey/monkey_island/cc/resources/configuration_import.py +++ b/monkey/monkey_island/cc/resources/configuration_import.py @@ -11,7 +11,7 @@ from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.server_utils.encryption import ( InvalidCiphertextError, InvalidCredentialsError, - PasswordBasedEncryptor, + PasswordBasedStringEncryptor, is_encrypted, ) from monkey_island.cc.services.config import ConfigService @@ -72,7 +72,7 @@ class ConfigurationImport(flask_restful.Resource): try: config = request_contents["config"] if ConfigurationImport.is_config_encrypted(request_contents["config"]): - pb_encryptor = PasswordBasedEncryptor(request_contents["password"]) + pb_encryptor = PasswordBasedStringEncryptor(request_contents["password"]) config = pb_encryptor.decrypt(config) return json.loads(config) except (JSONDecodeError, InvalidCiphertextError): diff --git a/monkey/monkey_island/cc/server_utils/encryption/__init__.py b/monkey/monkey_island/cc/server_utils/encryption/__init__.py index 7d806139c..84e6e6252 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/__init__.py +++ b/monkey/monkey_island/cc/server_utils/encryption/__init__.py @@ -1,11 +1,10 @@ from monkey_island.cc.server_utils.encryption.i_encryptor import IEncryptor from monkey_island.cc.server_utils.encryption.key_based_encryptor import KeyBasedEncryptor -from monkey_island.cc.server_utils.encryption.password_based_encryption import ( - InvalidCiphertextError, - InvalidCredentialsError, - PasswordBasedEncryptor, +from monkey_island.cc.server_utils.encryption.password_based_string_encryption import ( + PasswordBasedStringEncryptor, is_encrypted, ) +from .password_based_byte_encryption import InvalidCredentialsError, InvalidCiphertextError from monkey_island.cc.server_utils.encryption.data_store_encryptor import ( DataStoreEncryptor, get_datastore_encryptor, diff --git a/monkey/monkey_island/cc/server_utils/encryption/password_based_encryption.py b/monkey/monkey_island/cc/server_utils/encryption/password_based_byte_encryption.py similarity index 64% rename from monkey/monkey_island/cc/server_utils/encryption/password_based_encryption.py rename to monkey/monkey_island/cc/server_utils/encryption/password_based_byte_encryption.py index 20708ce31..569dc0d12 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/password_based_encryption.py +++ b/monkey/monkey_island/cc/server_utils/encryption/password_based_byte_encryption.py @@ -1,6 +1,6 @@ -import base64 import io import logging +from io import BytesIO import pyAesCrypt @@ -17,36 +17,28 @@ logger = logging.getLogger(__name__) # Note: password != key -class PasswordBasedEncryptor(IEncryptor): +class PasswordBasedByteEncryptor(IEncryptor): _BUFFER_SIZE = pyAesCrypt.crypto.bufferSizeDef def __init__(self, password: str): self.password = password - def encrypt(self, plaintext: str) -> str: - plaintext_stream = io.BytesIO(plaintext.encode()) + def encrypt(self, plaintext: BytesIO) -> BytesIO: ciphertext_stream = io.BytesIO() - pyAesCrypt.encryptStream( - plaintext_stream, ciphertext_stream, self.password, self._BUFFER_SIZE - ) + pyAesCrypt.encryptStream(plaintext, ciphertext_stream, self.password, self._BUFFER_SIZE) - ciphertext_b64 = base64.b64encode(ciphertext_stream.getvalue()) - logger.info("String encrypted.") + return ciphertext_stream - return ciphertext_b64.decode() - - def decrypt(self, ciphertext: str): - ciphertext = base64.b64decode(ciphertext) - ciphertext_stream = io.BytesIO(ciphertext) + def decrypt(self, ciphertext: BytesIO) -> BytesIO: plaintext_stream = io.BytesIO() - ciphertext_stream_len = len(ciphertext_stream.getvalue()) + ciphertext_stream_len = len(ciphertext.getvalue()) try: pyAesCrypt.decryptStream( - ciphertext_stream, + ciphertext, plaintext_stream, self.password, self._BUFFER_SIZE, @@ -59,7 +51,7 @@ class PasswordBasedEncryptor(IEncryptor): else: logger.info("The corrupt ciphertext provided.") raise InvalidCiphertextError - return plaintext_stream.getvalue().decode("utf-8") + return plaintext_stream class InvalidCredentialsError(Exception): @@ -68,8 +60,3 @@ class InvalidCredentialsError(Exception): class InvalidCiphertextError(Exception): """ Raised when ciphertext is corrupted """ - - -def is_encrypted(ciphertext: str) -> bool: - ciphertext = base64.b64decode(ciphertext) - return ciphertext.startswith(b"AES") diff --git a/monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryption.py b/monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryption.py new file mode 100644 index 000000000..ea796a441 --- /dev/null +++ b/monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryption.py @@ -0,0 +1,38 @@ +import base64 +import io +import logging + +import pyAesCrypt + +from monkey_island.cc.server_utils.encryption import IEncryptor +from monkey_island.cc.server_utils.encryption.password_based_byte_encryption import ( + PasswordBasedByteEncryptor, +) + +logger = logging.getLogger(__name__) + + +class PasswordBasedStringEncryptor(IEncryptor): + + _BUFFER_SIZE = pyAesCrypt.crypto.bufferSizeDef + + def __init__(self, password: str): + self.password = password + + def encrypt(self, plaintext: str) -> str: + plaintext_stream = io.BytesIO(plaintext.encode()) + ciphertext = PasswordBasedByteEncryptor(self.password).encrypt(plaintext_stream) + + return base64.b64encode(ciphertext.getvalue()).decode() + + def decrypt(self, ciphertext: str) -> str: + ciphertext = base64.b64decode(ciphertext) + ciphertext_stream = io.BytesIO(ciphertext) + + plaintext_stream = PasswordBasedByteEncryptor(self.password).decrypt(ciphertext_stream) + return plaintext_stream.getvalue().decode("utf-8") + + +def is_encrypted(ciphertext: str) -> bool: + ciphertext = base64.b64decode(ciphertext) + return ciphertext.startswith(b"AES") diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_configuration_import.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_configuration_import.py index fb397f234..bf7ccff80 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_configuration_import.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_configuration_import.py @@ -8,7 +8,7 @@ from tests.unit_tests.monkey_island.cc.services.utils.ciphertexts_for_encryption from common.utils.exceptions import InvalidConfigurationError from monkey_island.cc.resources.configuration_import import ConfigurationImport -from monkey_island.cc.server_utils.encryption import PasswordBasedEncryptor +from monkey_island.cc.server_utils.encryption import PasswordBasedStringEncryptor def test_is_config_encrypted__json(monkey_config_json): @@ -17,7 +17,7 @@ def test_is_config_encrypted__json(monkey_config_json): @pytest.mark.slow def test_is_config_encrypted__ciphertext(monkey_config_json): - pb_encryptor = PasswordBasedEncryptor(PASSWORD) + pb_encryptor = PasswordBasedStringEncryptor(PASSWORD) encrypted_config = pb_encryptor.encrypt(monkey_config_json) assert ConfigurationImport.is_config_encrypted(encrypted_config) diff --git a/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_password_based_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_password_based_encryption.py index d00609481..a231f3219 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_password_based_encryption.py +++ b/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_password_based_encryption.py @@ -4,7 +4,10 @@ from tests.unit_tests.monkey_island.cc.services.utils.ciphertexts_for_encryption VALID_CIPHER_TEXT, ) -from monkey_island.cc.server_utils.encryption import InvalidCredentialsError, PasswordBasedEncryptor +from monkey_island.cc.server_utils.encryption import ( + InvalidCredentialsError, + PasswordBasedStringEncryptor, +) MONKEY_CONFIGS_DIR_PATH = "monkey_configs" STANDARD_PLAINTEXT_MONKEY_CONFIG_FILENAME = "monkey_config_standard.json" @@ -14,27 +17,27 @@ INCORRECT_PASSWORD = "goodbye321" @pytest.mark.slow def test_encrypt_decrypt_string(monkey_config_json): - pb_encryptor = PasswordBasedEncryptor(PASSWORD) + pb_encryptor = PasswordBasedStringEncryptor(PASSWORD) encrypted_config = pb_encryptor.encrypt(monkey_config_json) assert pb_encryptor.decrypt(encrypted_config) == monkey_config_json @pytest.mark.slow def test_decrypt_string__wrong_password(monkey_config_json): - pb_encryptor = PasswordBasedEncryptor(INCORRECT_PASSWORD) + pb_encryptor = PasswordBasedStringEncryptor(INCORRECT_PASSWORD) with pytest.raises(InvalidCredentialsError): pb_encryptor.decrypt(VALID_CIPHER_TEXT) @pytest.mark.slow def test_decrypt_string__malformed_corrupted(): - pb_encryptor = PasswordBasedEncryptor(PASSWORD) + pb_encryptor = PasswordBasedStringEncryptor(PASSWORD) with pytest.raises(ValueError): pb_encryptor.decrypt(MALFORMED_CIPHER_TEXT_CORRUPTED) @pytest.mark.slow def test_decrypt_string__no_password(monkey_config_json): - pb_encryptor = PasswordBasedEncryptor("") + pb_encryptor = PasswordBasedStringEncryptor("") with pytest.raises(InvalidCredentialsError): pb_encryptor.decrypt(VALID_CIPHER_TEXT)