diff --git a/monkey/monkey_island/cc/resources/auth/auth.py b/monkey/monkey_island/cc/resources/auth/auth.py index 06fcb9e07..124cd1f71 100644 --- a/monkey/monkey_island/cc/resources/auth/auth.py +++ b/monkey/monkey_island/cc/resources/auth/auth.py @@ -13,7 +13,10 @@ from monkey_island.cc.resources.auth.credential_utils import ( get_credentials_from_request, password_matches_hash, ) -from monkey_island.cc.server_utils.encryption.data_store_encryptor import setup_datastore_key +from monkey_island.cc.server_utils.encryption import ( + get_datastore_encryptor, + initialize_datastore_encryptor, +) logger = logging.getLogger(__name__) @@ -44,7 +47,8 @@ class Authenticate(flask_restful.Resource): username, password = get_credentials_from_request(request) if _credentials_match_registered_user(username, password): - setup_datastore_key(username, password) + if not get_datastore_encryptor(): + initialize_datastore_encryptor(username, password) access_token = _create_access_token(username) return make_response({"access_token": access_token, "error": ""}, 200) else: diff --git a/monkey/monkey_island/cc/resources/auth/registration.py b/monkey/monkey_island/cc/resources/auth/registration.py index f96a5ce82..064e80179 100644 --- a/monkey/monkey_island/cc/resources/auth/registration.py +++ b/monkey/monkey_island/cc/resources/auth/registration.py @@ -6,7 +6,10 @@ from flask import make_response, request import monkey_island.cc.environment.environment_singleton as env_singleton from common.utils.exceptions import InvalidRegistrationCredentialsError, RegistrationNotNeededError from monkey_island.cc.resources.auth.credential_utils import get_user_credentials_from_request -from monkey_island.cc.server_utils.encryption import remove_old_datastore_key, setup_datastore_key +from monkey_island.cc.server_utils.encryption import ( + initialize_datastore_encryptor, + remove_old_datastore_key, +) from monkey_island.cc.setup.mongo.database_initializer import reset_database logger = logging.getLogger(__name__) @@ -24,7 +27,7 @@ class Registration(flask_restful.Resource): env_singleton.env.try_add_user(credentials) remove_old_datastore_key() username, password = get_user_credentials_from_request(request) - setup_datastore_key(username, password) + initialize_datastore_encryptor(username, password) reset_database() return make_response({"error": ""}, 200) except (InvalidRegistrationCredentialsError, RegistrationNotNeededError) as e: diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 69ab0437a..82c1b3a58 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -11,6 +11,8 @@ from gevent.pywsgi import WSGIServer # Add the monkey_island directory to the path, to make sure imports that don't start with # "monkey_island." work. +from monkey_island.cc.server_utils.encryption import initialize_encryptor_factory + MONKEY_ISLAND_DIR_BASE_PATH = str(Path(__file__).parent.parent) if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path: sys.path.insert(0, MONKEY_ISLAND_DIR_BASE_PATH) @@ -27,7 +29,6 @@ from monkey_island.cc.server_utils.consts import ( # noqa: E402 GEVENT_EXCEPTION_LOG, MONGO_CONNECTION_TIMEOUT, ) -from monkey_island.cc.server_utils.encryption import initialize_datastore_encryptor # noqa: E402 from monkey_island.cc.server_utils.island_logger import reset_logger, setup_logging # noqa: E402 from monkey_island.cc.services.initialize import initialize_services # noqa: E402 from monkey_island.cc.services.reporting.exporter_init import populate_exporter_list # noqa: E402 @@ -87,7 +88,7 @@ def _configure_logging(config_options): def _initialize_globals(config_options: IslandConfigOptions, server_config_path: str): env_singleton.initialize_from_file(server_config_path) - initialize_datastore_encryptor(config_options.data_dir) + initialize_encryptor_factory(config_options.data_dir) initialize_services(config_options.data_dir) diff --git a/monkey/monkey_island/cc/server_utils/encryption/__init__.py b/monkey/monkey_island/cc/server_utils/encryption/__init__.py index e9c19d75e..fa9692ca3 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/__init__.py +++ b/monkey/monkey_island/cc/server_utils/encryption/__init__.py @@ -4,15 +4,15 @@ from monkey_island.cc.server_utils.encryption.password_based_string_encryptior i PasswordBasedStringEncryptor, is_encrypted, ) -from .password_based_bytes_encryption import InvalidCredentialsError, InvalidCiphertextError -from monkey_island.cc.server_utils.encryption.data_store_encryptor import ( - DataStoreEncryptor, - get_datastore_encryptor, - initialize_datastore_encryptor, +from .encryptor_factory import ( + FactoryNotInitializedError, remove_old_datastore_key, - setup_datastore_key, - EncryptorNotInitializedError, + get_encryptor_factory, + get_secret_from_credentials, + initialize_encryptor_factory, ) +from .data_store_encryptor import initialize_datastore_encryptor, get_datastore_encryptor +from .password_based_bytes_encryption import InvalidCredentialsError, InvalidCiphertextError from .dict_encryption.dict_encryptor import ( SensitiveField, encrypt_dict, diff --git a/monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py b/monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py index 624e663b4..f58401783 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py +++ b/monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py @@ -8,7 +8,11 @@ from typing import Union from Crypto import Random # noqa: DUO133 # nosec: B413 -from monkey_island.cc.server_utils.encryption import KeyBasedEncryptor +from monkey_island.cc.server_utils.encryption import FactoryNotInitializedError, KeyBasedEncryptor +from monkey_island.cc.server_utils.encryption.encryptor_factory import ( + get_encryptor_factory, + get_secret_from_credentials, +) from monkey_island.cc.server_utils.encryption.password_based_bytes_encryption import ( PasswordBasedBytesEncryptor, ) @@ -19,33 +23,27 @@ _encryptor: Union[None, DataStoreEncryptor] = None class DataStoreEncryptor: _BLOCK_SIZE = 32 - _KEY_FILENAME = "mongo_key.bin" - def __init__(self, key_file_dir: str): - self.key_file_path = os.path.join(key_file_dir, self._KEY_FILENAME) - self._key_based_encryptor = None - - def init_key(self, secret: str): - if os.path.exists(self.key_file_path): - self._load_existing_key(secret) + def __init__(self, key_file_path: str, secret: str): + if os.path.exists(key_file_path): + self._key_based_encryptor = DataStoreEncryptor._load_existing_key(key_file_path, secret) else: - self._create_new_key(secret) + self._key_based_encryptor = DataStoreEncryptor._create_new_key(key_file_path, secret) - def _load_existing_key(self, secret: str): - with open(self.key_file_path, "rb") as f: + @staticmethod + def _load_existing_key(key_file_path: str, secret: str): + with open(key_file_path, "rb") as f: encrypted_key = f.read() cipher_key = PasswordBasedBytesEncryptor(secret).decrypt(encrypted_key) - self._key_based_encryptor = KeyBasedEncryptor(cipher_key) + return KeyBasedEncryptor(cipher_key) - def _create_new_key(self, secret: str): - cipher_key = Random.new().read(self._BLOCK_SIZE) + @staticmethod + def _create_new_key(key_file_path: str, secret: str): + cipher_key = Random.new().read(DataStoreEncryptor._BLOCK_SIZE) encrypted_key = PasswordBasedBytesEncryptor(secret).encrypt(cipher_key) - with open_new_securely_permissioned_file(self.key_file_path, "wb") as f: + with open_new_securely_permissioned_file(key_file_path, "wb") as f: f.write(encrypted_key) - self._key_based_encryptor = KeyBasedEncryptor(cipher_key) - - def is_key_setup(self) -> bool: - return self._key_based_encryptor is not None + return KeyBasedEncryptor(cipher_key) def enc(self, message: str): return self._key_based_encryptor.encrypt(message) @@ -54,43 +52,14 @@ class DataStoreEncryptor: return self._key_based_encryptor.decrypt(enc_message) -def initialize_datastore_encryptor(key_file_dir: str): +def initialize_datastore_encryptor(username: str, password: str): global _encryptor - _encryptor = DataStoreEncryptor(key_file_dir) - - -class EncryptorNotInitializedError(Exception): - pass - - -def _get_secret_from_credentials(username: str, password: str) -> str: - return f"{username}:{password}" - - -def encryptor_initialized_key_not_set(f): - def inner_function(*args, **kwargs): - if _encryptor is None: - raise EncryptorNotInitializedError - else: - if not _encryptor.is_key_setup(): - return f(*args, **kwargs) - else: - pass - - return inner_function - - -@encryptor_initialized_key_not_set -def remove_old_datastore_key(): - if os.path.isfile(_encryptor.key_file_path): - os.remove(_encryptor.key_file_path) - - -@encryptor_initialized_key_not_set -def setup_datastore_key(username: str, password: str): - secret = _get_secret_from_credentials(username, password) - _encryptor.init_key(secret) + factory = get_encryptor_factory() + if not factory: + raise FactoryNotInitializedError + secret = get_secret_from_credentials(username, password) + _encryptor = DataStoreEncryptor(factory.key_file_path, secret) def get_datastore_encryptor(): diff --git a/monkey/monkey_island/cc/server_utils/encryption/encryptor_factory.py b/monkey/monkey_island/cc/server_utils/encryption/encryptor_factory.py new file mode 100644 index 000000000..67d3e6ee4 --- /dev/null +++ b/monkey/monkey_island/cc/server_utils/encryption/encryptor_factory.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import os +from ctypes import Union + +_factory: Union[None, EncryptorFactory] = None + + +class EncryptorFactory: + + _KEY_FILENAME = "mongo_key.bin" + + def __init__(self, key_file_dir: str): + self.key_file_path = os.path.join(key_file_dir, self._KEY_FILENAME) + + +class FactoryNotInitializedError(Exception): + pass + + +def get_secret_from_credentials(username: str, password: str) -> str: + return f"{username}:{password}" + + +def remove_old_datastore_key(): + if _factory is None: + raise FactoryNotInitializedError + if os.path.isfile(_factory.key_file_path): + os.remove(_factory.key_file_path) + + +def initialize_encryptor_factory(key_file_dir: str): + global _factory + _factory = EncryptorFactory(key_file_dir) + + +def get_encryptor_factory(): + return _factory diff --git a/monkey/tests/unit_tests/monkey_island/cc/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/conftest.py index ee3e8aafa..cbd141ac7 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/conftest.py @@ -10,8 +10,10 @@ from tests.unit_tests.monkey_island.cc.server_utils.encryption.test_password_bas STANDARD_PLAINTEXT_MONKEY_CONFIG_FILENAME, ) -from monkey_island.cc.server_utils.encryption import initialize_datastore_encryptor -from monkey_island.cc.server_utils.encryption.data_store_encryptor import setup_datastore_key +from monkey_island.cc.server_utils.encryption import ( + initialize_datastore_encryptor, + initialize_encryptor_factory, +) @pytest.fixture @@ -34,5 +36,5 @@ MOCK_PASSWORD = "3cr3t_p455w0rd" @pytest.fixture def uses_encryptor(data_for_tests_dir): - initialize_datastore_encryptor(data_for_tests_dir) - setup_datastore_key(MOCK_USERNAME, MOCK_PASSWORD) + initialize_encryptor_factory(data_for_tests_dir) + initialize_datastore_encryptor(MOCK_USERNAME, MOCK_PASSWORD) diff --git a/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_data_store_encryptor.py b/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_data_store_encryptor.py index 1a394c025..e8901862b 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_data_store_encryptor.py +++ b/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_data_store_encryptor.py @@ -2,14 +2,15 @@ import pytest from tests.unit_tests.monkey_island.cc.conftest import MOCK_PASSWORD, MOCK_USERNAME from monkey_island.cc.server_utils.encryption import ( - DataStoreEncryptor, - EncryptorNotInitializedError, + FactoryNotInitializedError, data_store_encryptor, + encryptor_factory, get_datastore_encryptor, initialize_datastore_encryptor, + initialize_encryptor_factory, remove_old_datastore_key, - setup_datastore_key, ) +from monkey_island.cc.server_utils.encryption.encryptor_factory import EncryptorFactory PLAINTEXT = "Hello, Monkey!" @@ -25,45 +26,38 @@ def test_encryption(data_for_tests_dir): @pytest.fixture def initialized_key_dir(tmpdir): - initialize_datastore_encryptor(tmpdir) - setup_datastore_key(MOCK_USERNAME, MOCK_PASSWORD) + initialize_encryptor_factory(tmpdir) + initialize_datastore_encryptor(MOCK_USERNAME, MOCK_PASSWORD) yield tmpdir data_store_encryptor._encryptor = None + encryptor_factory._factory = None def test_key_creation(initialized_key_dir): - assert (initialized_key_dir / DataStoreEncryptor._KEY_FILENAME).isfile() - - -def test_key_removal_fails_if_key_initialized(initialized_key_dir): - remove_old_datastore_key() - assert (initialized_key_dir / DataStoreEncryptor._KEY_FILENAME).isfile() + assert (initialized_key_dir / EncryptorFactory._KEY_FILENAME).isfile() def test_key_removal(initialized_key_dir, monkeypatch): - monkeypatch.setattr(DataStoreEncryptor, "is_key_setup", lambda _: False) remove_old_datastore_key() - assert not (initialized_key_dir / DataStoreEncryptor._KEY_FILENAME).isfile() + assert not (initialized_key_dir / EncryptorFactory._KEY_FILENAME).isfile() def test_key_removal__no_key(tmpdir): - initialize_datastore_encryptor(tmpdir) - assert not (tmpdir / DataStoreEncryptor._KEY_FILENAME).isfile() + initialize_encryptor_factory(tmpdir) + assert not (tmpdir / EncryptorFactory._KEY_FILENAME).isfile() # Make sure no error thrown when we try to remove an non-existing key remove_old_datastore_key() - - data_store_encryptor._encryptor = None + encryptor_factory._factory = None def test_encryptor_not_initialized(): - with pytest.raises(EncryptorNotInitializedError): + with pytest.raises(FactoryNotInitializedError): remove_old_datastore_key() - setup_datastore_key() + initialize_datastore_encryptor(MOCK_USERNAME, MOCK_PASSWORD) -def test_setup_datastore_key(tmpdir): - initialize_datastore_encryptor(tmpdir) - assert not (tmpdir / DataStoreEncryptor._KEY_FILENAME).isfile() - setup_datastore_key(MOCK_USERNAME, MOCK_PASSWORD) - assert (tmpdir / DataStoreEncryptor._KEY_FILENAME).isfile() - assert get_datastore_encryptor().is_key_setup() +def test_initialize_encryptor(tmpdir): + initialize_encryptor_factory(tmpdir) + assert not (tmpdir / EncryptorFactory._KEY_FILENAME).isfile() + initialize_datastore_encryptor(MOCK_USERNAME, MOCK_PASSWORD) + assert (tmpdir / EncryptorFactory._KEY_FILENAME).isfile()