Island: Unlock RepositoryEncryptor in AuthenticationService

monkey/monkey_island/cc/services/authentication_service.py

@@ -10,6 +10,7 @@ from common.utils.exceptions import (
 from monkey_island.cc.models import UserCredentials
 from monkey_island.cc.repository import IUserRepository
 from monkey_island.cc.server_utils.encryption import (
+    ILockableEncryptor,
     reset_datastore_encryptor,
     unlock_datastore_encryptor,
 )
@@ -17,9 +18,15 @@ from monkey_island.cc.setup.mongo.database_initializer import reset_database
 
 
 class AuthenticationService:
-    def __init__(self, data_dir: Path, user_datastore: IUserRepository):
+    def __init__(
+        self,
+        data_dir: Path,
+        user_datastore: IUserRepository,
+        repository_encryptor: ILockableEncryptor,
+    ):
         self._data_dir = data_dir
         self._user_datastore = user_datastore
+        self._repository_encryptor = repository_encryptor
 
     def needs_registration(self) -> bool:
         return not self._user_datastore.has_registered_users()
@@ -30,7 +37,7 @@ class AuthenticationService:
 
         credentials = UserCredentials(username, _hash_password(password))
         self._user_datastore.add_user(credentials)
-        self._reset_datastore_encryptor(username, password)
+        self._reset_repository_encryptor(username, password)
         reset_database()
 
     def authenticate(self, username: str, password: str):
@@ -42,14 +49,21 @@ class AuthenticationService:
         if not _credentials_match_registered_user(username, password, registered_user):
             raise IncorrectCredentialsError()
 
-        self._unlock_datastore_encryptor(username, password)
+        self._unlock_repository_encryptor(username, password)
 
-    def _unlock_datastore_encryptor(self, username: str, password: str):
+    def _unlock_repository_encryptor(self, username: str, password: str):
         secret = _get_secret_from_credentials(username, password)
+        self._repository_encryptor.unlock(secret.encode())
+
+        # Legacy datastore encryptor will be removed soon
         unlock_datastore_encryptor(self._data_dir, secret)
 
-    def _reset_datastore_encryptor(self, username: str, password: str):
+    def _reset_repository_encryptor(self, username: str, password: str):
         secret = _get_secret_from_credentials(username, password)
+        self._repository_encryptor.reset_key()
+        self._repository_encryptor.unlock(secret.encode())
+
+        # Legacy datastore encryptor will be removed soon
         reset_datastore_encryptor(self._data_dir, secret)
 
 

monkey/tests/unit_tests/monkey_island/cc/services/test_authentication_service.py

@@ -10,6 +10,7 @@ from common.utils.exceptions import (
 )
 from monkey_island.cc.models import UserCredentials
 from monkey_island.cc.repository import IUserRepository
+from monkey_island.cc.server_utils.encryption import ILockableEncryptor
 from monkey_island.cc.services import AuthenticationService, authentication_service
 
 USERNAME = "user1"
@@ -48,6 +49,11 @@ def mock_unlock_datastore_encryptor():
     return MagicMock()
 
 
+@pytest.fixture
+def mock_repository_encryptor():
+    return MagicMock(spec=ILockableEncryptor)
+
+
 @pytest.fixture(autouse=True)
 def patch_datastore_utils(
     monkeypatch,
@@ -64,20 +70,20 @@ def patch_datastore_utils(
     )
 
 
-def test_needs_registration__true(tmp_path):
+def test_needs_registration__true(tmp_path, mock_repository_encryptor):
     has_registered_users = False
     mock_user_datastore = MockUserDatastore(lambda: has_registered_users, None, None)
 
-    a_s = AuthenticationService(tmp_path, mock_user_datastore)
+    a_s = AuthenticationService(tmp_path, mock_user_datastore, mock_repository_encryptor)
 
     assert a_s.needs_registration()
 
 
-def test_needs_registration__false(tmp_path):
+def test_needs_registration__false(tmp_path, mock_repository_encryptor):
     has_registered_users = True
     mock_user_datastore = MockUserDatastore(lambda: has_registered_users, None, None)
 
-    a_s = AuthenticationService(tmp_path, mock_user_datastore)
+    a_s = AuthenticationService(tmp_path, mock_user_datastore, mock_repository_encryptor)
 
     assert not a_s.needs_registration()
 
@@ -85,39 +91,45 @@ def test_needs_registration__false(tmp_path):
 @pytest.mark.slow
 @pytest.mark.parametrize("error", [InvalidRegistrationCredentialsError, AlreadyRegisteredError])
 def test_register_new_user__fails(
-    tmp_path, mock_reset_datastore_encryptor, mock_reset_database, error
+    tmp_path, mock_reset_datastore_encryptor, mock_reset_database, mock_repository_encryptor, error
 ):
     mock_user_datastore = MockUserDatastore(lambda: True, MagicMock(side_effect=error), None)
 
-    a_s = AuthenticationService(tmp_path, mock_user_datastore)
+    a_s = AuthenticationService(tmp_path, mock_user_datastore, mock_repository_encryptor)
 
     with pytest.raises(error):
         a_s.register_new_user(USERNAME, PASSWORD)
 
     mock_reset_datastore_encryptor.assert_not_called()
+    mock_repository_encryptor.reset_key().assert_not_called()
+    mock_repository_encryptor.unlock.assert_not_called()
     mock_reset_database.assert_not_called()
 
 
 def test_register_new_user__empty_password_fails(
-    tmp_path, mock_reset_datastore_encryptor, mock_reset_database
+    tmp_path, mock_reset_datastore_encryptor, mock_reset_database, mock_repository_encryptor
 ):
     mock_user_datastore = MockUserDatastore(lambda: False, None, None)
 
-    a_s = AuthenticationService(tmp_path, mock_user_datastore)
+    a_s = AuthenticationService(tmp_path, mock_user_datastore, mock_repository_encryptor)
 
     with pytest.raises(InvalidRegistrationCredentialsError):
         a_s.register_new_user(USERNAME, "")
 
     mock_reset_datastore_encryptor.assert_not_called()
+    mock_repository_encryptor.reset_key().assert_not_called()
+    mock_repository_encryptor.unlock.assert_not_called()
     mock_reset_database.assert_not_called()
 
 
 @pytest.mark.slow
-def test_register_new_user(tmp_path, mock_reset_datastore_encryptor, mock_reset_database):
+def test_register_new_user(
+    tmp_path, mock_reset_datastore_encryptor, mock_reset_database, mock_repository_encryptor
+):
     mock_add_user = MagicMock()
     mock_user_datastore = MockUserDatastore(lambda: False, mock_add_user, None)
 
-    a_s = AuthenticationService(tmp_path, mock_user_datastore)
+    a_s = AuthenticationService(tmp_path, mock_user_datastore, mock_repository_encryptor)
 
     a_s.register_new_user(USERNAME, PASSWORD)
 
@@ -127,30 +139,16 @@ def test_register_new_user(tmp_path, mock_reset_datastore_encryptor, mock_reset_
     mock_reset_datastore_encryptor.assert_called_once()
     assert mock_reset_datastore_encryptor.call_args[0][1] != USERNAME
 
+    mock_repository_encryptor.reset_key.assert_called_once()
+    mock_repository_encryptor.unlock.assert_called_once()
+    assert mock_repository_encryptor.unlock.call_args[0][0] != USERNAME
+
     mock_reset_database.assert_called_once()
 
 
 @pytest.mark.slow
-def test_authenticate__success(tmp_path, mock_unlock_datastore_encryptor):
-    mock_user_datastore = MockUserDatastore(
-        lambda: True,
-        None,
-        lambda _: UserCredentials(USERNAME, PASSWORD_HASH),
-    )
-
-    a_s = AuthenticationService(tmp_path, mock_user_datastore)
-
-    # If authentication fails, this function will raise an exception and the test will fail.
-    a_s.authenticate(USERNAME, PASSWORD)
-    mock_unlock_datastore_encryptor.assert_called_once()
-
-
-@pytest.mark.slow
-@pytest.mark.parametrize(
-    ("username", "password"), [("wrong_username", PASSWORD), (USERNAME, "wrong_password")]
-)
-def test_authenticate__failed_wrong_credentials(
-    tmp_path, mock_unlock_datastore_encryptor, username, password
+def test_authenticate__success(
+    tmp_path, mock_unlock_datastore_encryptor, mock_repository_encryptor
 ):
     mock_user_datastore = MockUserDatastore(
         lambda: True,
@@ -158,22 +156,47 @@ def test_authenticate__failed_wrong_credentials(
         lambda _: UserCredentials(USERNAME, PASSWORD_HASH),
     )
 
-    a_s = AuthenticationService(tmp_path, mock_user_datastore)
+    a_s = AuthenticationService(tmp_path, mock_user_datastore, mock_repository_encryptor)
+
+    # If authentication fails, this function will raise an exception and the test will fail.
+    a_s.authenticate(USERNAME, PASSWORD)
+    mock_unlock_datastore_encryptor.assert_called_once()
+    mock_repository_encryptor.unlock.assert_called_once()
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize(
+    ("username", "password"), [("wrong_username", PASSWORD), (USERNAME, "wrong_password")]
+)
+def test_authenticate__failed_wrong_credentials(
+    tmp_path, mock_unlock_datastore_encryptor, username, password, mock_repository_encryptor
+):
+    mock_user_datastore = MockUserDatastore(
+        lambda: True,
+        None,
+        lambda _: UserCredentials(USERNAME, PASSWORD_HASH),
+    )
+
+    a_s = AuthenticationService(tmp_path, mock_user_datastore, mock_repository_encryptor)
 
     with pytest.raises(IncorrectCredentialsError):
         a_s.authenticate(username, password)
 
     mock_unlock_datastore_encryptor.assert_not_called()
+    mock_repository_encryptor.unlock.assert_not_called()
 
 
-def test_authenticate__failed_no_registered_user(tmp_path, mock_unlock_datastore_encryptor):
+def test_authenticate__failed_no_registered_user(
+    tmp_path, mock_unlock_datastore_encryptor, mock_repository_encryptor
+):
     mock_user_datastore = MockUserDatastore(
         lambda: True, None, MagicMock(side_effect=UnknownUserError)
     )
 
-    a_s = AuthenticationService(tmp_path, mock_user_datastore)
+    a_s = AuthenticationService(tmp_path, mock_user_datastore, mock_repository_encryptor)
 
     with pytest.raises(IncorrectCredentialsError):
         a_s.authenticate(USERNAME, PASSWORD)
 
     mock_unlock_datastore_encryptor.assert_not_called()
+    mock_repository_encryptor.unlock.assert_not_called()