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"), ) ]