diff --git a/monkey/monkey_island/cc/server_utils/encryption/__init__.py b/monkey/monkey_island/cc/server_utils/encryption/__init__.py index 1b302a6fc..7fbc882cb 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/__init__.py +++ b/monkey/monkey_island/cc/server_utils/encryption/__init__.py @@ -12,9 +12,9 @@ from .password_based_bytes_encryptor import ( InvalidCiphertextError, ) from .data_store_encryptor import ( - initialize_datastore_encryptor, get_datastore_encryptor, - remove_old_datastore_key, + initialize_datastore_encryptor, + reinitialize_datastore_encryptor, ) from .dict_encryptor import ( SensitiveField, 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 6a44b5cff..7e9a279d4 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 @@ -53,11 +53,24 @@ class DataStoreEncryptor(IEncryptor): if self._key_file_path.is_file(): self._key_file_path.unlink() + self._key_based_encryptor = None + + +def reinitialize_datastore_encryptor( + key_file_dir: str, secret: str, key_file_name: str = "mongo_key.bin" +): + _delete_encryptor() + initialize_datastore_encryptor(key_file_dir, secret, key_file_name) + + +def _delete_encryptor(): + global _encryptor -def remove_old_datastore_key(): if _encryptor: _encryptor.erase_key() + _encryptor = None + def initialize_datastore_encryptor( key_file_dir: str, secret: str, key_file_name: str = "mongo_key.bin" diff --git a/monkey/monkey_island/cc/services/authentication.py b/monkey/monkey_island/cc/services/authentication.py index 2d7940055..9cb0121c5 100644 --- a/monkey/monkey_island/cc/services/authentication.py +++ b/monkey/monkey_island/cc/services/authentication.py @@ -1,7 +1,7 @@ from monkey_island.cc.server_utils.encryption import ( get_datastore_encryptor, initialize_datastore_encryptor, - remove_old_datastore_key, + reinitialize_datastore_encryptor, ) @@ -22,8 +22,8 @@ class AuthenticationService: @staticmethod def reset_datastore_encryptor(username: str, password: str): - remove_old_datastore_key() - AuthenticationService._init_encryptor_from_credentials(username, password) + secret = AuthenticationService._get_secret_from_credentials(username, password) + reinitialize_datastore_encryptor(AuthenticationService.KEY_FILE_DIRECTORY, secret) @staticmethod def _init_encryptor_from_credentials(username: str, password: str): 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 7d5616185..4ae6b15e1 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 @@ -1,10 +1,11 @@ import pytest +from common.utils.file_utils import get_file_sha256_hash from monkey_island.cc.server_utils.encryption import ( data_store_encryptor, get_datastore_encryptor, initialize_datastore_encryptor, - remove_old_datastore_key, + reinitialize_datastore_encryptor, ) PLAINTEXT = "Hello, Monkey!" @@ -42,28 +43,47 @@ def test_key_creation(key_file, tmp_path): assert key_file.is_file() +@pytest.mark.slow +def test_existing_key_reused(key_file, tmp_path): + assert not key_file.is_file() + + initialize_datastore_encryptor(tmp_path, MOCK_SECRET, KEY_FILENAME) + key_file_hash_1 = get_file_sha256_hash(key_file) + + initialize_datastore_encryptor(tmp_path, MOCK_SECRET, KEY_FILENAME) + key_file_hash_2 = get_file_sha256_hash(key_file) + + assert key_file_hash_1 == key_file_hash_2 + + @pytest.mark.slow def test_key_removal(key_file, tmp_path): initialize_datastore_encryptor(tmp_path, MOCK_SECRET, KEY_FILENAME) assert key_file.is_file() - remove_old_datastore_key() + get_datastore_encryptor().erase_key() assert not key_file.is_file() -def test_key_removal__no_key(key_file): - assert not key_file.is_file() - # Make sure no error thrown when we try to remove an non-existing key - remove_old_datastore_key() - - -def test_key_removal__no_key_2(key_file, tmp_path): +@pytest.mark.slow +def test_key_removal__no_key(key_file, tmp_path): assert not key_file.is_file() initialize_datastore_encryptor(tmp_path, MOCK_SECRET, KEY_FILENAME) assert key_file.is_file() - key_file.unlink() + get_datastore_encryptor().erase_key() assert not key_file.is_file() # Make sure no error thrown when we try to remove an non-existing key get_datastore_encryptor().erase_key() + + +@pytest.mark.slow +def test_reinitialize_datastore_encryptor(key_file, tmp_path): + initialize_datastore_encryptor(tmp_path, MOCK_SECRET, KEY_FILENAME) + key_file_hash_1 = get_file_sha256_hash(key_file) + + reinitialize_datastore_encryptor(tmp_path, MOCK_SECRET, KEY_FILENAME) + key_file_hash_2 = get_file_sha256_hash(key_file) + + assert key_file_hash_1 != key_file_hash_2