diff --git a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py index df5d0e417..971a2bf5d 100644 --- a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py @@ -13,14 +13,23 @@ class MongoCredentialsRepository(ICredentialsRepository): def get_configured_credentials(self) -> Sequence[Credentials]: try: - configured_credentials = list(mongo.db.configured_credentials.find({})) + configured_credentials = [] + list_configured_credentials = list(mongo.db.configured_credentials.find({})) + for c in list_configured_credentials: + del c["_id"] + configured_credentials.append(Credentials.from_mapping(c)) + return configured_credentials except Exception as err: raise RetrievalError(err) def get_stolen_credentials(self) -> Sequence[Credentials]: try: - stolen_credentials = list(mongo.db.stolen_credentials.find({})) + stolen_credentials = [] + list_stolen_credentials = list(mongo.db.stolen_credentials.find({})) + for c in list_stolen_credentials: + del c["_id"] + stolen_credentials.append(Credentials.from_mapping(c)) return stolen_credentials except Exception as err: raise RetrievalError(err) @@ -36,25 +45,27 @@ class MongoCredentialsRepository(ICredentialsRepository): def save_configured_credentials(self, credentials: Sequence[Credentials]): try: - mongo.db.configured_credentials.insert_many(credentials) + for c in credentials: + mongo.db.configured_credentials.insert_one(Credentials.to_mapping(c)) except Exception as err: raise StorageError(err) def save_stolen_credentials(self, credentials: Sequence[Credentials]): try: - mongo.db.stolen_credentials.insert_many(credentials) + for c in credentials: + mongo.db.stolen_credentials.insert_one(Credentials.to_mapping(c)) except Exception as err: raise StorageError(err) def remove_configured_credentials(self): try: - mongo.db.configured_credentials.remove({}) + mongo.db.configured_credentials.delete_many({}) except Exception as err: raise RemovalError(err) def remove_stolen_credentials(self): try: - mongo.db.stolen_credentials.remove({}) + mongo.db.stolen_credentials.delete_many({}) except Exception as err: raise RemovalError(err) 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 28b1c2eab..8d32b961b 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 @@ -1,8 +1,47 @@ import mongoengine import pytest +from common.credentials import Credentials from monkey_island.cc.repository import MongoCredentialsRepository +USER1 = "test_user_1" +USER2 = "test_user_2" +USER3 = "test_user_3" +PASSWORD = "12435" +PASSWORD2 = "password" +PASSWORD3 = "lozinka" +LM_HASH = "AEBD4DE384C7EC43AAD3B435B51404EE" +NT_HASH = "7A21990FCD3D759941E45C490F143D5F" +PUBLIC_KEY = "MY_PUBLIC_KEY" +PRIVATE_KEY = "MY_PRIVATE_KEY" + +CREDENTIALS_DICT_1 = { + "identities": [ + {"credential_type": "USERNAME", "username": USER1}, + {"credential_type": "USERNAME", "username": USER2}, + ], + "secrets": [ + {"credential_type": "PASSWORD", "password": PASSWORD}, + {"credential_type": "LM_HASH", "lm_hash": LM_HASH}, + {"credential_type": "NT_HASH", "nt_hash": NT_HASH}, + { + "credential_type": "SSH_KEYPAIR", + "public_key": PUBLIC_KEY, + "private_key": PRIVATE_KEY, + }, + ], +} + +CREDENTIALS_DICT_2 = { + "identities": [ + {"credential_type": "USERNAME", "username": USER3}, + ], + "secrets": [ + {"credential_type": "PASSWORD", "password": PASSWORD2}, + {"credential_type": "PASSWORD", "password": PASSWORD3}, + ], +} + @pytest.fixture def fake_mongo(monkeypatch): @@ -29,3 +68,79 @@ def test_mongo_repository_get_all(fake_mongo): actual_credentials = MongoCredentialsRepository().get_all_credentials() assert actual_credentials == [] + + +def test_mongo_repository_configured(fake_mongo): + + credentials = [ + Credentials.from_mapping(CREDENTIALS_DICT_1), + Credentials.from_mapping(CREDENTIALS_DICT_2), + ] + + mongo_repository = MongoCredentialsRepository() + + mongo_repository.save_configured_credentials(credentials) + + actual_configured_credentials = mongo_repository.get_configured_credentials() + + assert actual_configured_credentials == credentials + + mongo_repository.remove_configured_credentials() + + actual_configured_credentials = mongo_repository.get_configured_credentials() + + assert actual_configured_credentials == [] + + +def test_mongo_repository_stolen(fake_mongo): + + stolen_credentials = [Credentials.from_mapping(CREDENTIALS_DICT_1)] + + configured_credentials = [Credentials.from_mapping(CREDENTIALS_DICT_2)] + + mongo_repository = MongoCredentialsRepository() + + mongo_repository.save_configured_credentials(configured_credentials) + mongo_repository.save_stolen_credentials(stolen_credentials) + + actual_stolen_credentials = mongo_repository.get_stolen_credentials() + + assert actual_stolen_credentials == stolen_credentials + + mongo_repository.remove_stolen_credentials() + + actual_stolen_credentials = mongo_repository.get_stolen_credentials() + + assert actual_stolen_credentials == [] + + # Must remove configured also for the next tests + mongo_repository.remove_configured_credentials() + + +def test_mongo_repository_all(fake_mongo): + + configured_credentials = [Credentials.from_mapping(CREDENTIALS_DICT_1)] + stolen_credentials = [Credentials.from_mapping(CREDENTIALS_DICT_2)] + all_credentials = [ + Credentials.from_mapping(CREDENTIALS_DICT_1), + Credentials.from_mapping(CREDENTIALS_DICT_2), + ] + + mongo_repository = MongoCredentialsRepository() + + mongo_repository.save_configured_credentials(configured_credentials) + mongo_repository.save_stolen_credentials(stolen_credentials) + + actual_credentials = mongo_repository.get_all_credentials() + + assert actual_credentials == all_credentials + + mongo_repository.remove_all_credentials() + + actual_credentials = mongo_repository.get_all_credentials() + actual_stolen_credentials = mongo_repository.get_stolen_credentials() + actual_configured_credentials = mongo_repository.get_configured_credentials() + + assert actual_credentials == [] + assert actual_stolen_credentials == [] + assert actual_configured_credentials == []