Island: Modify MongoCredentialsRepository

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)
 

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 == []