Island: Add reinitialize_datastore_encryptor()

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,

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"

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()
+        secret = AuthenticationService._get_secret_from_credentials(username, password)
-        AuthenticationService._init_encryptor_from_credentials(username, password)
+        reinitialize_datastore_encryptor(AuthenticationService.KEY_FILE_DIRECTORY, secret)
 
     @staticmethod
     def _init_encryptor_from_credentials(username: str, password: str):

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):
+@pytest.mark.slow
-    assert not key_file.is_file()
+def test_key_removal__no_key(key_file, tmp_path):
-    # 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):
     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