Agent, UT: Update credentials store using `setdefault().update`

* get_credentials use PropgationCredentials type
* private stored credentials in Aggregating Credentials Store
* initial values in credentials store constructor
* build_puppet accepts ICredentialsStore
* private telemetry_messenger in monkey

monkey/infection_monkey/credential_store/aggregating_credentials_store.py

@@ -1,9 +1,10 @@
 import logging
-from typing import Iterable, Mapping
+from typing import Any, Iterable, Mapping
 
 from common.common_consts.credential_component_type import CredentialComponentType
 from infection_monkey.i_control_channel import IControlChannel
 from infection_monkey.i_puppet import Credentials
+from infection_monkey.typing import PropagationCredentials
 
 from .i_credentials_store import ICredentialsStore
 
@@ -12,63 +13,73 @@ logger = logging.getLogger(__name__)
 
 class AggregatingCredentialsStore(ICredentialsStore):
     def __init__(self, control_channel: IControlChannel):
-        self.stored_credentials = {}
+        self._stored_credentials = {
+            "exploit_user_list": set(),
+            "exploit_password_list": set(),
+            "exploit_lm_hash_list": set(),
+            "exploit_ntlm_hash_list": set(),
+            "exploit_ssh_keys": [],
+        }
         self._control_channel = control_channel
 
-    def add_credentials(self, credentials_to_add: Iterable[Credentials]) -> None:
+    def add_credentials(self, credentials_to_add: Iterable[Credentials]):
         for credentials in credentials_to_add:
-            usernames = [
+            usernames = {
                 identity.username
                 for identity in credentials.identities
                 if identity.credential_type is CredentialComponentType.USERNAME
-            ]
+            }
-            self._set_attribute("exploit_user_list", usernames)
+            self._stored_credentials.setdefault("exploit_user_list", set()).update(usernames)
 
             for secret in credentials.secrets:
                 if secret.credential_type is CredentialComponentType.PASSWORD:
-                    self._set_attribute("exploit_password_list", [secret.password])
+                    self._stored_credentials.setdefault("exploit_password_list", set()).update(
+                        [secret.password]
+                    )
                 elif secret.credential_type is CredentialComponentType.LM_HASH:
-                    self._set_attribute("exploit_lm_hash_list", [secret.lm_hash])
+                    self._stored_credentials.setdefault("exploit_lm_hash_list", set()).update(
+                        [secret.lm_hash]
+                    )
                 elif secret.credential_type is CredentialComponentType.NT_HASH:
-                    self._set_attribute("exploit_ntlm_hash_list", [secret.nt_hash])
+                    self._stored_credentials.setdefault("exploit_ntlm_hash_list", set()).update(
+                        [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):
+    def get_credentials(self) -> PropagationCredentials:
         try:
             propagation_credentials = self._control_channel.get_credentials_for_propagation()
+
+            # Needs to be reworked when exploiters accepts sequence of Credentials
             self._aggregate_credentials(propagation_credentials)
-            return self.stored_credentials
+
+            return self._stored_credentials
         except Exception as ex:
-            self.stored_credentials = {}
+            self._stored_credentials = {}
             logger.error(f"Error while attempting to retrieve credentials for propagation: {ex}")
 
     def _aggregate_credentials(self, credentials_to_aggr: Mapping):
         for cred_attr, credentials_values in credentials_to_aggr.items():
             self._set_attribute(cred_attr, credentials_values)
 
-    def _set_attribute(self, attribute_to_be_set, credentials_values):
+    def _set_attribute(self, attribute_to_be_set: str, credentials_values: Iterable[Any]):
-        if attribute_to_be_set not in self.stored_credentials:
+        if not credentials_values:
-            self.stored_credentials[attribute_to_be_set] = []
+            return
 
-        if credentials_values:
+        if isinstance(credentials_values[0], dict):
-            if isinstance(credentials_values[0], dict):
+            self._stored_credentials[attribute_to_be_set] = []
-                self.stored_credentials.setdefault(attribute_to_be_set, []).extend(
+            self._stored_credentials.setdefault(attribute_to_be_set, []).extend(credentials_values)
-                    credentials_values
+            self._stored_credentials[attribute_to_be_set] = [
-                )
+                dict(s_c)
-                self.stored_credentials[attribute_to_be_set] = [
+                for s_c in set(
-                    dict(s_c)
+                    frozenset(d_c.items()) for d_c in self._stored_credentials[attribute_to_be_set]
-                    for s_c in set(
-                        frozenset(d_c.items())
-                        for d_c in self.stored_credentials[attribute_to_be_set]
-                    )
-                ]
-            else:
-                self.stored_credentials[attribute_to_be_set] = sorted(
-                    list(
-                        set(self.stored_credentials[attribute_to_be_set]).union(credentials_values)
-                    )
                 )
+            ]
+        else:
+            self._stored_credentials.setdefault(attribute_to_be_set, set()).update(
+                credentials_values
+            )

monkey/infection_monkey/credential_store/i_credentials_store.py

@@ -1,20 +1,22 @@
 import abc
-from typing import Iterable, Mapping
+from typing import Iterable
 
 from infection_monkey.i_puppet import Credentials
+from infection_monkey.typing import PropagationCredentials
 
 
 class ICredentialsStore(metaclass=abc.ABCMeta):
     @abc.abstractmethod
-    def add_credentials(self, credentials_to_add: Iterable[Credentials]) -> None:
+    def add_credentials(self, credentials_to_add: Iterable[Credentials]):
-        """a
+        """
-        Method that adds credentials to the CredentialStore
+        Adds credentials to the CredentialStore
-        :param Credentials credentials: The credentials that will be added
+        :param Iterable[Credentials] credentials: The credentials that will be added
         """
 
     @abc.abstractmethod
-    def get_credentials(self) -> Mapping:
+    def get_credentials(self) -> PropagationCredentials:
         """
-        Method that retrieves credentials from the store
+        Retrieves credentials from the store
         :return: Credentials that can be used for propagation
+        :type: PropagationCredentials
         """

monkey/infection_monkey/monkey.py

@@ -15,7 +15,7 @@ from infection_monkey.credential_collectors import (
     MimikatzCredentialCollector,
     SSHCredentialCollector,
 )
-from infection_monkey.credential_store import AggregatingCredentialsStore
+from infection_monkey.credential_store import AggregatingCredentialsStore, ICredentialsStore
 from infection_monkey.exploit import CachingAgentRepository, ExploiterWrapper
 from infection_monkey.exploit.hadoop import HadoopExploiter
 from infection_monkey.exploit.log4shell import Log4ShellExploiter
@@ -89,8 +89,7 @@ class InfectionMonkey:
         self._default_server = self._opts.server
         # TODO used in propogation phase
         self._monkey_inbound_tunnel = None
-        self._credentials_store = None
+        self._telemetry_messenger = LegacyTelemetryMessengerAdapter()
-        self.telemetry_messenger = LegacyTelemetryMessengerAdapter()
         self._current_depth = self._opts.depth
         self._master = None
 
@@ -125,7 +124,7 @@ class InfectionMonkey:
         if is_windows_os():
             T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
 
-        run_aws_environment_check(self.telemetry_messenger)
+        run_aws_environment_check(self._telemetry_messenger)
 
         should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop()
         if should_stop:
@@ -183,21 +182,21 @@ class InfectionMonkey:
         local_network_interfaces = InfectionMonkey._get_local_network_interfaces()
 
         control_channel = ControlChannel(self._default_server, GUID)
-        self._credentials_store = AggregatingCredentialsStore(control_channel)
+        credentials_store = AggregatingCredentialsStore(control_channel)
 
-        puppet = self._build_puppet()
+        puppet = self._build_puppet(credentials_store)
 
         victim_host_factory = self._build_victim_host_factory(local_network_interfaces)
 
         telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-            self.telemetry_messenger, self._monkey_inbound_tunnel
+            self._telemetry_messenger, self._monkey_inbound_tunnel
         )
 
         telemetry_messenger = CredentialsInterceptingTelemetryMessenger(
             ExploitInterceptingTelemetryMessenger(
-                self.telemetry_messenger, self._monkey_inbound_tunnel
+                self._telemetry_messenger, self._monkey_inbound_tunnel
             ),
-            self._credentials_store,
+            credentials_store,
         )
 
         self._master = AutomatedMaster(
@@ -207,7 +206,7 @@ class InfectionMonkey:
             victim_host_factory,
             control_channel,
             local_network_interfaces,
-            self._credentials_store,
+            credentials_store,
         )
 
     @staticmethod
@@ -218,7 +217,7 @@ class InfectionMonkey:
 
         return local_network_interfaces
 
-    def _build_puppet(self) -> IPuppet:
+    def _build_puppet(self, credentials_store: ICredentialsStore) -> IPuppet:
         puppet = Puppet()
 
         puppet.load_plugin(
@@ -228,7 +227,7 @@ class InfectionMonkey:
         )
         puppet.load_plugin(
             "SSHCollector",
-            SSHCredentialCollector(self.telemetry_messenger),
+            SSHCredentialCollector(self._telemetry_messenger),
             PluginType.CREDENTIAL_COLLECTOR,
         )
 
@@ -241,7 +240,7 @@ class InfectionMonkey:
         agent_repository = CachingAgentRepository(
             f"https://{self._default_server}", ControlClient.proxies
         )
-        exploit_wrapper = ExploiterWrapper(self.telemetry_messenger, agent_repository)
+        exploit_wrapper = ExploiterWrapper(self._telemetry_messenger, agent_repository)
 
         puppet.load_plugin(
             "HadoopExploiter", exploit_wrapper.wrap(HadoopExploiter), PluginType.EXPLOITER
@@ -260,7 +259,7 @@ class InfectionMonkey:
         )
 
         zerologon_telemetry_messenger = CredentialsInterceptingTelemetryMessenger(
-            self.telemetry_messenger, self._credentials_store
+            self._telemetry_messenger, credentials_store
         )
         zerologon_wrapper = ExploiterWrapper(zerologon_telemetry_messenger, agent_repository)
         puppet.load_plugin(
@@ -271,54 +270,54 @@ class InfectionMonkey:
 
         puppet.load_plugin(
             "CommunicateAsBackdoorUser",
-            CommunicateAsBackdoorUser(self.telemetry_messenger),
+            CommunicateAsBackdoorUser(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
             "ModifyShellStartupFiles",
-            ModifyShellStartupFiles(self.telemetry_messenger),
+            ModifyShellStartupFiles(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
-            "HiddenFiles", HiddenFiles(self.telemetry_messenger), PluginType.POST_BREACH_ACTION
+            "HiddenFiles", HiddenFiles(self._telemetry_messenger), PluginType.POST_BREACH_ACTION
         )
         puppet.load_plugin(
             "TrapCommand",
-            CommunicateAsBackdoorUser(self.telemetry_messenger),
+            CommunicateAsBackdoorUser(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
             "ChangeSetuidSetgid",
-            ChangeSetuidSetgid(self.telemetry_messenger),
+            ChangeSetuidSetgid(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
-            "ScheduleJobs", ScheduleJobs(self.telemetry_messenger), PluginType.POST_BREACH_ACTION
+            "ScheduleJobs", ScheduleJobs(self._telemetry_messenger), PluginType.POST_BREACH_ACTION
         )
         puppet.load_plugin(
-            "Timestomping", Timestomping(self.telemetry_messenger), PluginType.POST_BREACH_ACTION
+            "Timestomping", Timestomping(self._telemetry_messenger), PluginType.POST_BREACH_ACTION
         )
         puppet.load_plugin(
             "AccountDiscovery",
-            AccountDiscovery(self.telemetry_messenger),
+            AccountDiscovery(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
             "ProcessListCollection",
-            ProcessListCollection(self.telemetry_messenger),
+            ProcessListCollection(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
-            "TrapCommand", TrapCommand(self.telemetry_messenger), PluginType.POST_BREACH_ACTION
+            "TrapCommand", TrapCommand(self._telemetry_messenger), PluginType.POST_BREACH_ACTION
         )
         puppet.load_plugin(
             "SignedScriptProxyExecution",
-            SignedScriptProxyExecution(self.telemetry_messenger),
+            SignedScriptProxyExecution(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
         puppet.load_plugin(
             "ClearCommandHistory",
-            ClearCommandHistory(self.telemetry_messenger),
+            ClearCommandHistory(self._telemetry_messenger),
             PluginType.POST_BREACH_ACTION,
         )
 

monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_credentials_store.py

@@ -6,7 +6,7 @@ from infection_monkey.credential_collectors import Password, SSHKeypair, Usernam
 from infection_monkey.credential_store import AggregatingCredentialsStore
 from infection_monkey.i_puppet import Credentials
 
-DEFAULT_CREDENTIALS = {
+CONTROL_CHANNEL_CREDENTIALS = {
     "exploit_user_list": ["Administrator", "root", "user1"],
     "exploit_password_list": ["123456", "123456789", "password", "root"],
     "exploit_lm_hash_list": ["aasdf23asd1fdaasadasdfas"],
@@ -28,7 +28,7 @@ PROPAGATION_CREDENTIALS = {
     "exploit_ssh_keys": [{"public_key": "some_public_key", "private_key": "some_private_key"}],
 }
 
-TELEM_CREDENTIALS = [
+CREDENTIALS_COLLECTION = [
     Credentials(
         [Username("user1"), Username("user3")],
         [
@@ -43,58 +43,51 @@ TELEM_CREDENTIALS = [
 @pytest.fixture
 def aggregating_credentials_store() -> AggregatingCredentialsStore:
     control_channel = MagicMock()
-    control_channel.get_credentials_for_propagation.return_value = DEFAULT_CREDENTIALS
+    control_channel.get_credentials_for_propagation.return_value = CONTROL_CHANNEL_CREDENTIALS
     return AggregatingCredentialsStore(control_channel)
 
 
 def test_get_credentials_from_store(aggregating_credentials_store):
-    aggregating_credentials_store.get_credentials()
+    actual_stored_credentials = aggregating_credentials_store.get_credentials()
 
-    actual_stored_credentials = aggregating_credentials_store.stored_credentials
+    print(actual_stored_credentials)
 
-    assert (
+    assert actual_stored_credentials["exploit_user_list"] == set(
-        actual_stored_credentials["exploit_user_list"] == DEFAULT_CREDENTIALS["exploit_user_list"]
+        CONTROL_CHANNEL_CREDENTIALS["exploit_user_list"]
     )
-    assert (
+    assert actual_stored_credentials["exploit_password_list"] == set(
-        actual_stored_credentials["exploit_password_list"]
+        CONTROL_CHANNEL_CREDENTIALS["exploit_password_list"]
-        == DEFAULT_CREDENTIALS["exploit_password_list"]
     )
-    assert (
+    assert actual_stored_credentials["exploit_ntlm_hash_list"] == set(
-        actual_stored_credentials["exploit_ntlm_hash_list"]
+        CONTROL_CHANNEL_CREDENTIALS["exploit_ntlm_hash_list"]
-        == DEFAULT_CREDENTIALS["exploit_ntlm_hash_list"]
     )
 
     for ssh_keypair in actual_stored_credentials["exploit_ssh_keys"]:
-        assert ssh_keypair in DEFAULT_CREDENTIALS["exploit_ssh_keys"]
+        assert ssh_keypair in CONTROL_CHANNEL_CREDENTIALS["exploit_ssh_keys"]
 
 
-def test_add_credentials_to_empty_store(aggregating_credentials_store):
+def test_add_credentials_to_store(aggregating_credentials_store):
-    aggregating_credentials_store.add_credentials(TELEM_CREDENTIALS)
+    aggregating_credentials_store.add_credentials(CREDENTIALS_COLLECTION)
 
-    assert aggregating_credentials_store.stored_credentials == PROPAGATION_CREDENTIALS
+    actual_stored_credentials = aggregating_credentials_store.get_credentials()
 
-
+    assert actual_stored_credentials["exploit_user_list"] == set(
-def test_add_credentials_to_full_store(aggregating_credentials_store):
+        [
-
+            "Administrator",
-    aggregating_credentials_store.get_credentials()
+            "root",
-
+            "user1",
-    aggregating_credentials_store.add_credentials(TELEM_CREDENTIALS)
+            "user3",
-
+        ]
-    actual_stored_credentials = aggregating_credentials_store.stored_credentials
+    )
-
+    assert actual_stored_credentials["exploit_password_list"] == set(
-    assert actual_stored_credentials["exploit_user_list"] == [
+        [
-        "Administrator",
+            "123456",
-        "root",
+            "123456789",
-        "user1",
+            "abcdefg",
-        "user3",
+            "password",
-    ]
+            "root",
-    assert actual_stored_credentials["exploit_password_list"] == [
+        ]
-        "123456",
+    )
-        "123456789",
-        "abcdefg",
-        "password",
-        "root",
-    ]
 
     for ssh_keypair in actual_stored_credentials["exploit_ssh_keys"]:
-        assert ssh_keypair in DEFAULT_CREDENTIALS["exploit_ssh_keys"]
+        assert ssh_keypair in CONTROL_CHANNEL_CREDENTIALS["exploit_ssh_keys"]

monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_credentials_intercepting_telemetry_messenger.py

@@ -1,7 +1,6 @@
 from unittest.mock import MagicMock
 
 from infection_monkey.credential_collectors import Password, SSHKeypair, Username
-from infection_monkey.credential_store import AggregatingCredentialsStore
 from infection_monkey.i_puppet import Credentials
 from infection_monkey.telemetry.base_telem import BaseTelem
 from infection_monkey.telemetry.credentials_telem import CredentialsTelem
@@ -56,19 +55,14 @@ def test_credentials_generic_telemetry():
 
 def test_successful_intercepting_credentials_telemetry():
     mock_telemetry_messenger = MagicMock()
-    aggregating_credentials_store = AggregatingCredentialsStore(MagicMock())
+    mock_credentials_store = MagicMock()
-    mock_empty_credentials_telem = MockCredentialsTelem([])
+    mock_empty_credentials_telem = MockCredentialsTelem(TELEM_CREDENTIALS)
 
     telemetry_messenger = CredentialsInterceptingTelemetryMessenger(
-        mock_telemetry_messenger, aggregating_credentials_store
+        mock_telemetry_messenger, mock_credentials_store
     )
 
     telemetry_messenger.send_telemetry(mock_empty_credentials_telem)
 
     assert mock_telemetry_messenger.send_telemetry.called
-    assert not aggregating_credentials_store.stored_credentials
+    assert mock_credentials_store.add_credentials.called
-
-    mock_credentials_telem = MockCredentialsTelem(TELEM_CREDENTIALS)
-    telemetry_messenger.send_telemetry(mock_credentials_telem)
-
-    assert aggregating_credentials_store.stored_credentials