From 0f2fc0902fcd0cc12c1a69d188982fb278be8db4 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Mon, 18 Jul 2022 11:29:52 +0200
Subject: [PATCH] Agent: Simplify credentials object in aggregating credentials
 store

---
 .../aggregating_credentials_store.py          | 46 +++++++++----------
 .../test_aggregating_credentials_store.py     | 21 ++++-----
 2 files changed, 32 insertions(+), 35 deletions(-)

diff --git a/monkey/infection_monkey/credential_store/aggregating_credentials_store.py b/monkey/infection_monkey/credential_store/aggregating_credentials_store.py
index 6ff6bb288..dacfbdcee 100644
--- a/monkey/infection_monkey/credential_store/aggregating_credentials_store.py
+++ b/monkey/infection_monkey/credential_store/aggregating_credentials_store.py
@@ -26,31 +26,29 @@ class AggregatingCredentialsStore(ICredentialsStore):
 
     def add_credentials(self, credentials_to_add: Iterable[Credentials]):
         for credentials in credentials_to_add:
-            usernames = {
-                identity.username
-                for identity in credentials.identities
-                if identity.credential_type is CredentialComponentType.USERNAME
-            }
-            self._stored_credentials.setdefault("exploit_user_list", set()).update(usernames)
+            self._stored_credentials.setdefault("exploit_user_list", set()).add(
+                credentials.identity.username
+            )
 
-            for secret in credentials.secrets:
-                if secret.credential_type is CredentialComponentType.PASSWORD:
-                    self._stored_credentials.setdefault("exploit_password_list", set()).add(
-                        secret.password
-                    )
-                elif secret.credential_type is CredentialComponentType.LM_HASH:
-                    self._stored_credentials.setdefault("exploit_lm_hash_list", set()).add(
-                        secret.lm_hash
-                    )
-                elif secret.credential_type is CredentialComponentType.NT_HASH:
-                    self._stored_credentials.setdefault("exploit_ntlm_hash_list", set()).add(
-                        secret.nt_hash
-                    )
-                elif secret.credential_type is CredentialComponentType.SSH_KEYPAIR:
-                    self._set_attribute(
-                        "exploit_ssh_keys",
-                        [{"public_key": secret.public_key, "private_key": secret.private_key}],
-                    )
+            secret = credentials.secret
+
+            if secret.credential_type is CredentialComponentType.PASSWORD:
+                self._stored_credentials.setdefault("exploit_password_list", set()).add(
+                    secret.password
+                )
+            elif secret.credential_type is CredentialComponentType.LM_HASH:
+                self._stored_credentials.setdefault("exploit_lm_hash_list", set()).add(
+                    secret.lm_hash
+                )
+            elif secret.credential_type is CredentialComponentType.NT_HASH:
+                self._stored_credentials.setdefault("exploit_ntlm_hash_list", set()).add(
+                    secret.nt_hash
+                )
+            elif secret.credential_type is CredentialComponentType.SSH_KEYPAIR:
+                self._set_attribute(
+                    "exploit_ssh_keys",
+                    [{"public_key": secret.public_key, "private_key": secret.private_key}],
+                )
 
     def get_credentials(self) -> PropagationCredentials:
         try:
diff --git a/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_credentials_store.py b/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_credentials_store.py
index baff7df65..54d43f59c 100644
--- a/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_credentials_store.py
+++ b/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_credentials_store.py
@@ -30,21 +30,20 @@ EMPTY_CHANNEL_CREDENTIALS = {
 
 TEST_CREDENTIALS = [
     Credentials(
-        [Username("user1"), Username("user3")],
-        [
-            Password("abcdefg"),
-            Password("root"),
-            SSHKeypair(public_key="some_public_key_1", private_key="some_private_key_1"),
-        ],
-    )
+        identity=Username("user1"),
+        secret=Password("root"),
+    ),
+    Credentials(identity=Username("user1"), secret=Password("abcdefg")),
+    Credentials(
+        identity=Username("user3"),
+        secret=SSHKeypair(public_key="some_public_key_1", private_key="some_private_key_1"),
+    ),
 ]
 
 SSH_KEYS_CREDENTIALS = [
     Credentials(
-        [Username("root")],
-        [
-            SSHKeypair(public_key="some_public_key", private_key="some_private_key"),
-        ],
+        Username("root"),
+        SSHKeypair(public_key="some_public_key", private_key="some_private_key"),
     )
 ]