Refactor DataStoreEncryptor by splitting up initialization related methods into EncryptorFactory

This makes encryptor initialization workflow more straight-forward and the files become smaller, easier to read

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:

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:

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)
 
 

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,

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

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

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)

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