From 3fd27c650388ee7f52135337ea173d5fcd4bf0e4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 15 Sep 2022 12:00:04 -0400 Subject: [PATCH] Island: Add ICredentialsRepository.reset() --- .../cc/repository/i_credentials_repository.py | 7 +++++++ .../cc/repository/mongo_credentials_repository.py | 3 +++ .../monkey_island/in_memory_credentials_repository.py | 3 +++ .../cc/repository/test_mongo_credentials_repository.py | 5 +++++ 4 files changed, 18 insertions(+) diff --git a/monkey/monkey_island/cc/repository/i_credentials_repository.py b/monkey/monkey_island/cc/repository/i_credentials_repository.py index 381782533..ad891ff39 100644 --- a/monkey/monkey_island/cc/repository/i_credentials_repository.py +++ b/monkey/monkey_island/cc/repository/i_credentials_repository.py @@ -84,3 +84,10 @@ class ICredentialsRepository(ABC): :raises RemovalError: If an error is encountered while attempting to remove the credentials """ pass + + def reset(self): + """ + An alias for remove_all_credentials() + + :raises RemovalError: If an error is encountered while attempting to remove the credentials + """ diff --git a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py index ab7c9f18b..438ce48b6 100644 --- a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py @@ -52,6 +52,9 @@ class MongoCredentialsRepository(ICredentialsRepository): self.remove_configured_credentials() self.remove_stolen_credentials() + def reset(self): + self.remove_all_credentials() + def _get_credentials_from_collection(self, collection) -> Sequence[Credentials]: try: collection_result = [] diff --git a/monkey/tests/monkey_island/in_memory_credentials_repository.py b/monkey/tests/monkey_island/in_memory_credentials_repository.py index 6eb8155e8..389a31e08 100644 --- a/monkey/tests/monkey_island/in_memory_credentials_repository.py +++ b/monkey/tests/monkey_island/in_memory_credentials_repository.py @@ -33,3 +33,6 @@ class InMemoryCredentialsRepository(ICredentialsRepository): def remove_all_credentials(self): self.remove_configured_credentials() self.remove_stolen_credentials() + + def reset(self): + self.remove_all_credentials() diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py index 88c584534..f921c9df5 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py @@ -142,6 +142,11 @@ def test_mongo_repository_remove_credentials__removal_error(error_raising_creden error_raising_credentials_repository.remove_stolen_credentials() +def test_mongo_repository_reset__removal_error(error_raising_credentials_repository): + with pytest.raises(RemovalError): + error_raising_credentials_repository.reset() + + @pytest.mark.parametrize("credentials", CREDENTIALS) def test_configured_secrets_encrypted( mongo_repository: MongoCredentialsRepository,