From 5c7bab7a0d0b406c0176584a1595f475647634f5 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Fri, 11 Jun 2021 16:05:32 +0300 Subject: [PATCH] Refactored json parsing out of encryption/decryption functionality. --- .../cc/resources/configuration_export.py | 5 +- .../cc/resources/configuration_import.py | 19 +++--- .../cc/services/utils/config_encryption.py | 62 ------------------- .../cc/services/utils/encryption.py | 59 ++++++++++++++++++ .../unit_tests/monkey_island/cc/conftest.py | 5 ++ .../cc/resources/test_configuration_import.py | 13 ++-- .../services/utils/test_config_encryption.py | 28 ++++----- 7 files changed, 96 insertions(+), 95 deletions(-) delete mode 100644 monkey/monkey_island/cc/services/utils/config_encryption.py create mode 100644 monkey/monkey_island/cc/services/utils/encryption.py diff --git a/monkey/monkey_island/cc/resources/configuration_export.py b/monkey/monkey_island/cc/resources/configuration_export.py index 40dd90044..4a7aeec24 100644 --- a/monkey/monkey_island/cc/resources/configuration_export.py +++ b/monkey/monkey_island/cc/resources/configuration_export.py @@ -5,7 +5,7 @@ from flask import request from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.services.config import ConfigService -from monkey_island.cc.services.utils.config_encryption import encrypt_config +from monkey_island.cc.services.utils.encryption import encrypt_string class ConfigurationExport(flask_restful.Resource): @@ -19,6 +19,7 @@ class ConfigurationExport(flask_restful.Resource): config_export = plaintext_config if should_encrypt: password = data["password"] - config_export = encrypt_config(plaintext_config, password) + plaintext_config = json.dumps(plaintext_config) + config_export = encrypt_string(plaintext_config, password) 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 50a4cf955..bcb3d8b1f 100644 --- a/monkey/monkey_island/cc/resources/configuration_import.py +++ b/monkey/monkey_island/cc/resources/configuration_import.py @@ -9,9 +9,10 @@ from flask import request from common.utils.exceptions import InvalidConfigurationError from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.services.config import ConfigService -from monkey_island.cc.services.utils.config_encryption import ( +from monkey_island.cc.services.utils.encryption import ( + InvalidCiphertextError, InvalidCredentialsError, - decrypt_config, + decrypt_ciphertext, is_encrypted, ) @@ -68,13 +69,13 @@ class ConfigurationImport(flask_restful.Resource): @staticmethod def _get_plaintext_config_from_request(request_contents: dict) -> dict: - if ConfigurationImport.is_config_encrypted(request_contents["config"]): - return decrypt_config(request_contents["config"], request_contents["password"]) - else: - try: - return json.loads(request_contents["config"]) - except JSONDecodeError: - raise InvalidConfigurationError + try: + config = request_contents["config"] + if ConfigurationImport.is_config_encrypted(request_contents["config"]): + config = decrypt_ciphertext(config, request_contents["password"]) + return json.loads(config) + except (JSONDecodeError, InvalidCiphertextError): + raise InvalidConfigurationError @staticmethod def import_config(config_json): diff --git a/monkey/monkey_island/cc/services/utils/config_encryption.py b/monkey/monkey_island/cc/services/utils/config_encryption.py deleted file mode 100644 index 70c1e12ee..000000000 --- a/monkey/monkey_island/cc/services/utils/config_encryption.py +++ /dev/null @@ -1,62 +0,0 @@ -import base64 -import io -import json -import logging -from typing import Dict - -import pyAesCrypt - -from common.utils.exceptions import InvalidConfigurationError - -BUFFER_SIZE = pyAesCrypt.crypto.bufferSizeDef - -logger = logging.getLogger(__name__) - - -def encrypt_config(config: Dict, password: str) -> str: - plaintext_config_stream = io.BytesIO(json.dumps(config).encode()) - ciphertext_config_stream = io.BytesIO() - - pyAesCrypt.encryptStream( - plaintext_config_stream, ciphertext_config_stream, password, BUFFER_SIZE - ) - - ciphertext_b64 = base64.b64encode(ciphertext_config_stream.getvalue()) - logger.info("Configuration encrypted.") - - return ciphertext_b64.decode() - - -def decrypt_config(ciphertext: str, password: str) -> Dict: - ciphertext = base64.b64decode(ciphertext) - ciphertext_config_stream = io.BytesIO(ciphertext) - dec_plaintext_config_stream = io.BytesIO() - - len_ciphertext_config_stream = len(ciphertext_config_stream.getvalue()) - - try: - pyAesCrypt.decryptStream( - ciphertext_config_stream, - dec_plaintext_config_stream, - password, - BUFFER_SIZE, - len_ciphertext_config_stream, - ) - except ValueError as ex: - if str(ex).startswith("Wrong password"): - logger.info("Wrong password for configuration provided.") - raise InvalidCredentialsError - else: - logger.info("The provided configuration file is corrupt.") - raise InvalidConfigurationError - plaintext_config = json.loads(dec_plaintext_config_stream.getvalue().decode("utf-8")) - return plaintext_config - - -def is_encrypted(ciphertext: str) -> bool: - ciphertext = base64.b64decode(ciphertext) - return ciphertext.startswith(b"AES") - - -class InvalidCredentialsError(Exception): - """ Raise when credentials supplied are invalid """ diff --git a/monkey/monkey_island/cc/services/utils/encryption.py b/monkey/monkey_island/cc/services/utils/encryption.py new file mode 100644 index 000000000..ae4af2257 --- /dev/null +++ b/monkey/monkey_island/cc/services/utils/encryption.py @@ -0,0 +1,59 @@ +import base64 +import io +import logging + +import pyAesCrypt + +BUFFER_SIZE = pyAesCrypt.crypto.bufferSizeDef + +logger = logging.getLogger(__name__) + + +def encrypt_string(plaintext: str, password: str) -> str: + plaintext_stream = io.BytesIO(plaintext.encode()) + ciphertext_stream = io.BytesIO() + + pyAesCrypt.encryptStream(plaintext_stream, ciphertext_stream, password, BUFFER_SIZE) + + ciphertext_b64 = base64.b64encode(ciphertext_stream.getvalue()) + logger.info("String encrypted.") + + return ciphertext_b64.decode() + + +def decrypt_ciphertext(ciphertext: str, password: str) -> str: + ciphertext = base64.b64decode(ciphertext) + ciphertext_stream = io.BytesIO(ciphertext) + plaintext_stream = io.BytesIO() + + ciphertext_stream_len = len(ciphertext_stream.getvalue()) + + try: + pyAesCrypt.decryptStream( + ciphertext_stream, + plaintext_stream, + password, + BUFFER_SIZE, + ciphertext_stream_len, + ) + except ValueError as ex: + if str(ex).startswith("Wrong password"): + logger.info("Wrong password provided for decryption.") + raise InvalidCredentialsError + else: + logger.info("The corrupt ciphertext provided.") + raise InvalidCiphertextError + return plaintext_stream.getvalue().decode("utf-8") + + +def is_encrypted(ciphertext: str) -> bool: + ciphertext = base64.b64decode(ciphertext) + return ciphertext.startswith(b"AES") + + +class InvalidCredentialsError(Exception): + """ Raised when password for decryption is invalid """ + + +class InvalidCiphertextError(Exception): + """ Raised when ciphertext is corrupted """ diff --git a/monkey/tests/unit_tests/monkey_island/cc/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/conftest.py index 370e0cc4d..c89c4c294 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/conftest.py @@ -18,3 +18,8 @@ def monkey_config(data_for_tests_dir): ) plaintext_config = json.loads(open(plaintext_monkey_config_standard_path, "r").read()) return plaintext_config + + +@pytest.fixture +def monkey_config_json(monkey_config): + return json.dumps(monkey_config) 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 1cfd6b914..ed1d908cf 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 @@ -1,5 +1,3 @@ -import json - import pytest from tests.unit_tests.monkey_island.cc.services.utils.ciphertexts_for_encryption_test import ( MALFORMED_CIPHER_TEXT_CORRUPTED, @@ -8,16 +6,15 @@ from tests.unit_tests.monkey_island.cc.services.utils.test_config_encryption imp from common.utils.exceptions import InvalidConfigurationError from monkey_island.cc.resources.configuration_import import ConfigurationImport -from monkey_island.cc.services.utils.config_encryption import encrypt_config +from monkey_island.cc.services.utils.encryption import encrypt_string -def test_is_config_encrypted__json(monkey_config): - monkey_config = json.dumps(monkey_config) - assert not ConfigurationImport.is_config_encrypted(monkey_config) +def test_is_config_encrypted__json(monkey_config_json): + assert not ConfigurationImport.is_config_encrypted(monkey_config_json) -def test_is_config_encrypted__ciphertext(monkey_config): - encrypted_config = encrypt_config(monkey_config, PASSWORD) +def test_is_config_encrypted__ciphertext(monkey_config_json): + encrypted_config = encrypt_string(monkey_config_json, PASSWORD) assert ConfigurationImport.is_config_encrypted(encrypted_config) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/utils/test_config_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/services/utils/test_config_encryption.py index fb7dcb929..23b399cc5 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/utils/test_config_encryption.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/utils/test_config_encryption.py @@ -3,10 +3,10 @@ from tests.unit_tests.monkey_island.cc.services.utils.ciphertexts_for_encryption MALFORMED_CIPHER_TEXT_CORRUPTED, ) -from monkey_island.cc.services.utils.config_encryption import ( +from monkey_island.cc.services.utils.encryption import ( InvalidCredentialsError, - decrypt_config, - encrypt_config, + decrypt_ciphertext, + encrypt_string, ) MONKEY_CONFIGS_DIR_PATH = "monkey_configs" @@ -15,23 +15,23 @@ PASSWORD = "hello123" INCORRECT_PASSWORD = "goodbye321" -def test_encrypt_decrypt_config(monkey_config): - encrypted_config = encrypt_config(monkey_config, PASSWORD) - assert decrypt_config(encrypted_config, PASSWORD) == monkey_config +def test_encrypt_decrypt_string(monkey_config_json): + encrypted_config = encrypt_string(monkey_config_json, PASSWORD) + assert decrypt_ciphertext(encrypted_config, PASSWORD) == monkey_config_json -def test_encrypt_decrypt_config__wrong_password(monkey_config): - encrypted_config = encrypt_config(monkey_config, PASSWORD) +def test_encrypt_decrypt_string__wrong_password(monkey_config_json): + encrypted_config = encrypt_string(monkey_config_json, PASSWORD) with pytest.raises(InvalidCredentialsError): - decrypt_config(encrypted_config, INCORRECT_PASSWORD) + decrypt_ciphertext(encrypted_config, INCORRECT_PASSWORD) -def test_encrypt_decrypt_config__malformed_corrupted(): +def test_encrypt_decrypt_string__malformed_corrupted(): with pytest.raises(ValueError): - decrypt_config(MALFORMED_CIPHER_TEXT_CORRUPTED, PASSWORD) + decrypt_ciphertext(MALFORMED_CIPHER_TEXT_CORRUPTED, PASSWORD) -def test_encrypt_decrypt_config__decrypt_no_password(monkey_config): - encrypted_config = encrypt_config(monkey_config, PASSWORD) +def test_encrypt_decrypt_string__decrypt_no_password(monkey_config_json): + encrypted_config = encrypt_string(monkey_config_json, PASSWORD) with pytest.raises(InvalidCredentialsError): - decrypt_config(encrypted_config, "") + decrypt_ciphertext(encrypted_config, "")