diff --git a/monkey/infection_monkey/credential_store/__init__.py b/monkey/infection_monkey/credential_store/__init__.py new file mode 100644 index 000000000..e05ce3160 --- /dev/null +++ b/monkey/infection_monkey/credential_store/__init__.py @@ -0,0 +1,2 @@ +from .i_credentials_store import ICredentialsStore +from .aggregating_credentials_store import AggregatingCredentialsStore diff --git a/monkey/infection_monkey/credential_store/aggregating_credentials_store.py b/monkey/infection_monkey/credential_store/aggregating_credentials_store.py new file mode 100644 index 000000000..eef6f72fc --- /dev/null +++ b/monkey/infection_monkey/credential_store/aggregating_credentials_store.py @@ -0,0 +1,91 @@ +import logging +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 infection_monkey.utils.decorators import request_cache + +from .i_credentials_store import ICredentialsStore + +logger = logging.getLogger(__name__) + +CREDENTIALS_POLL_PERIOD_SEC = 30 + + +class AggregatingCredentialsStore(ICredentialsStore): + def __init__(self, control_channel: IControlChannel): + 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]): + 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) + + 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}], + ) + + def get_credentials(self) -> PropagationCredentials: + try: + propagation_credentials = self._get_credentials_from_control_channel() + + # Needs to be reworked when exploiters accepts sequence of Credentials + self._aggregate_credentials(propagation_credentials) + + return self._stored_credentials + except Exception as ex: + self._stored_credentials = {} + logger.error(f"Error while attempting to retrieve credentials for propagation: {ex}") + + @request_cache(CREDENTIALS_POLL_PERIOD_SEC) + def _get_credentials_from_control_channel(self) -> PropagationCredentials: + return self._control_channel.get_credentials_for_propagation() + + 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: str, credentials_values: Iterable[Any]): + if not credentials_values: + return + + if isinstance(credentials_values[0], dict): + self._stored_credentials.setdefault(attribute_to_be_set, []).extend(credentials_values) + self._stored_credentials[attribute_to_be_set] = [ + dict(s_c) + for s_c in set( + frozenset(d_c.items()) for d_c in self._stored_credentials[attribute_to_be_set] + ) + ] + else: + self._stored_credentials.setdefault(attribute_to_be_set, set()).update( + credentials_values + ) diff --git a/monkey/infection_monkey/credential_store/i_credentials_store.py b/monkey/infection_monkey/credential_store/i_credentials_store.py new file mode 100644 index 000000000..72ef12d44 --- /dev/null +++ b/monkey/infection_monkey/credential_store/i_credentials_store.py @@ -0,0 +1,22 @@ +import abc +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]): + """ + Adds credentials to the CredentialStore + :param Iterable[Credentials] credentials: The credentials that will be added + """ + + @abc.abstractmethod + def get_credentials(self) -> PropagationCredentials: + """ + Retrieves credentials from the store + :return: Credentials that can be used for propagation + :type: PropagationCredentials + """ diff --git a/monkey/infection_monkey/exploit/sshexec.py b/monkey/infection_monkey/exploit/sshexec.py index 0e6a9c038..dab29ae03 100644 --- a/monkey/infection_monkey/exploit/sshexec.py +++ b/monkey/infection_monkey/exploit/sshexec.py @@ -60,7 +60,7 @@ class SSHExploiter(HostExploiter): for user, ssh_key_pair in ssh_key_pairs_iterator: # Creating file-like private key for paramiko pkey = io.StringIO(ssh_key_pair["private_key"]) - ssh_string = "%s@%s" % (ssh_key_pair["user"], ssh_key_pair["ip"]) + ssh_string = "%s@%s" % (user, self.host.ip_addr) ssh = paramiko.SSHClient() ssh.set_missing_host_key_policy(paramiko.WarningPolicy()) diff --git a/monkey/infection_monkey/exploit/zerologon.py b/monkey/infection_monkey/exploit/zerologon.py index 7a6a41dca..9be006d8e 100644 --- a/monkey/infection_monkey/exploit/zerologon.py +++ b/monkey/infection_monkey/exploit/zerologon.py @@ -91,8 +91,9 @@ class ZerologonExploiter(HostExploiter): @staticmethod def connect_to_dc(dc_ip) -> object: binding = epm.hept_map(dc_ip, nrpc.MSRPC_UUID_NRPC, protocol="ncacn_ip_tcp") - rpc_con = transport.DCERPCTransportFactory(binding).get_dce_rpc() - rpc_con.set_connect_timeout(LONG_REQUEST_TIMEOUT) + rpc_transport = transport.DCERPCTransportFactory(binding) + rpc_transport.set_connect_timeout(LONG_REQUEST_TIMEOUT) + rpc_con = rpc_transport.get_dce_rpc() rpc_con.connect() rpc_con.bind(nrpc.MSRPC_UUID_NRPC) return rpc_con diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 1e837bb74..b499b578d 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -3,6 +3,7 @@ import threading import time from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple +from infection_monkey.credential_store import ICredentialsStore from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError from infection_monkey.i_master import IMaster from infection_monkey.i_puppet import IPuppet @@ -36,6 +37,7 @@ class AutomatedMaster(IMaster): victim_host_factory: VictimHostFactory, control_channel: IControlChannel, local_network_interfaces: List[NetworkInterface], + credentials_store: ICredentialsStore, ): self._current_depth = current_depth self._puppet = puppet @@ -43,9 +45,8 @@ class AutomatedMaster(IMaster): self._control_channel = control_channel ip_scanner = IPScanner(self._puppet, NUM_SCAN_THREADS) - exploiter = Exploiter( - self._puppet, NUM_EXPLOIT_THREADS, self._control_channel.get_credentials_for_propagation - ) + + exploiter = Exploiter(self._puppet, NUM_EXPLOIT_THREADS, credentials_store.get_credentials) self._propagator = Propagator( self._telemetry_messenger, ip_scanner, diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index d6781af7f..eb6dd12d5 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -7,14 +7,12 @@ from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from infection_monkey.config import WormConfiguration from infection_monkey.control import ControlClient from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError -from infection_monkey.utils.decorators import request_cache +from infection_monkey.typing import PropagationCredentials requests.packages.urllib3.disable_warnings() logger = logging.getLogger(__name__) -CREDENTIALS_POLL_PERIOD_SEC = 30 - class ControlChannel(IControlChannel): def __init__(self, server: str, agent_id: str): @@ -69,8 +67,7 @@ class ControlChannel(IControlChannel): ) as e: raise IslandCommunicationError(e) - @request_cache(CREDENTIALS_POLL_PERIOD_SEC) - def get_credentials_for_propagation(self) -> dict: + def get_credentials_for_propagation(self) -> PropagationCredentials: propagation_credentials_url = ( f"https://{self._control_channel_server}/api/propagation-credentials/{self._agent_id}" ) diff --git a/monkey/infection_monkey/master/exploiter.py b/monkey/infection_monkey/master/exploiter.py index 3f0087af8..5f8e25b4d 100644 --- a/monkey/infection_monkey/master/exploiter.py +++ b/monkey/infection_monkey/master/exploiter.py @@ -9,6 +9,7 @@ from typing import Callable, Dict, List, Mapping from infection_monkey.i_puppet import ExploiterResultData, IPuppet from infection_monkey.model import VictimHost +from infection_monkey.typing import PropagationCredentials from infection_monkey.utils.threading import interruptible_iter, run_worker_threads QUEUE_TIMEOUT = 2 @@ -24,7 +25,7 @@ class Exploiter: self, puppet: IPuppet, num_workers: int, - get_updated_credentials_for_propagation: Callable[[], Mapping], + get_updated_credentials_for_propagation: Callable[[], PropagationCredentials], ): self._puppet = puppet self._num_workers = num_workers @@ -160,7 +161,7 @@ class Exploiter: exploitation_success=False, propagation_success=False, error_message=msg ) - def _get_credentials_for_propagation(self) -> Mapping: + def _get_credentials_for_propagation(self) -> PropagationCredentials: try: return self._get_updated_credentials_for_propagation() except Exception as ex: diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 1f0a53a63..abcd96772 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -15,6 +15,7 @@ from infection_monkey.credential_collectors import ( MimikatzCredentialCollector, SSHCredentialCollector, ) +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 @@ -54,6 +55,9 @@ from infection_monkey.puppet.puppet import Puppet from infection_monkey.system_singleton import SystemSingleton from infection_monkey.telemetry.attack.t1106_telem import T1106Telem from infection_monkey.telemetry.attack.t1107_telem import T1107Telem +from infection_monkey.telemetry.messengers.credentials_intercepting_telemetry_messenger import ( + CredentialsInterceptingTelemetryMessenger, +) from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import ( ExploitInterceptingTelemetryMessenger, ) @@ -84,7 +88,7 @@ class InfectionMonkey: self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server) self._default_server = self._opts.server self._monkey_inbound_tunnel = None - self.telemetry_messenger = LegacyTelemetryMessengerAdapter() + self._telemetry_messenger = LegacyTelemetryMessengerAdapter() self._current_depth = self._opts.depth self._master = None @@ -119,7 +123,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: @@ -175,12 +179,23 @@ class InfectionMonkey: def _build_master(self): local_network_interfaces = InfectionMonkey._get_local_network_interfaces() - puppet = self._build_puppet() + + control_channel = ControlChannel(self._default_server, GUID) + credentials_store = AggregatingCredentialsStore(control_channel) + + 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 + ), + credentials_store, ) self._master = AutomatedMaster( @@ -188,8 +203,9 @@ class InfectionMonkey: puppet, telemetry_messenger, victim_host_factory, - ControlChannel(self._default_server, GUID), + control_channel, local_network_interfaces, + credentials_store, ) @staticmethod @@ -200,7 +216,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( @@ -210,7 +226,7 @@ class InfectionMonkey: ) puppet.load_plugin( "SSHCollector", - SSHCredentialCollector(self.telemetry_messenger), + SSHCredentialCollector(self._telemetry_messenger), PluginType.CREDENTIAL_COLLECTOR, ) @@ -223,7 +239,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 @@ -240,62 +256,67 @@ class InfectionMonkey: puppet.load_plugin( "MSSQLExploiter", exploit_wrapper.wrap(MSSQLExploiter), PluginType.EXPLOITER ) + + zerologon_telemetry_messenger = CredentialsInterceptingTelemetryMessenger( + self._telemetry_messenger, credentials_store + ) + zerologon_wrapper = ExploiterWrapper(zerologon_telemetry_messenger, agent_repository) puppet.load_plugin( "ZerologonExploiter", - exploit_wrapper.wrap(ZerologonExploiter), + zerologon_wrapper.wrap(ZerologonExploiter), PluginType.EXPLOITER, ) 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, ) diff --git a/monkey/infection_monkey/telemetry/credentials_telem.py b/monkey/infection_monkey/telemetry/credentials_telem.py index c0573d942..4f5c43aa4 100644 --- a/monkey/infection_monkey/telemetry/credentials_telem.py +++ b/monkey/infection_monkey/telemetry/credentials_telem.py @@ -17,6 +17,10 @@ class CredentialsTelem(BaseTelem): """ self._credentials = credentials + @property + def credentials(self) -> Iterable[Credentials]: + return iter(self._credentials) + def send(self, log_data=True): super().send(log_data=False) diff --git a/monkey/infection_monkey/telemetry/messengers/credentials_intercepting_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/credentials_intercepting_telemetry_messenger.py new file mode 100644 index 000000000..541800577 --- /dev/null +++ b/monkey/infection_monkey/telemetry/messengers/credentials_intercepting_telemetry_messenger.py @@ -0,0 +1,38 @@ +from functools import singledispatch + +from infection_monkey.credential_store import ICredentialsStore +from infection_monkey.telemetry.credentials_telem import CredentialsTelem +from infection_monkey.telemetry.i_telem import ITelem +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger + + +class CredentialsInterceptingTelemetryMessenger(ITelemetryMessenger): + def __init__( + self, telemetry_messenger: ITelemetryMessenger, credentials_store: ICredentialsStore + ): + self._telemetry_messenger = telemetry_messenger + self._credentials_store = credentials_store + + def send_telemetry(self, telemetry: ITelem): + _send_telemetry(telemetry, self._telemetry_messenger, self._credentials_store) + + +# Note: We can use @singledispatchmethod instead of @singledispatch if we migrate to Python 3.8 or +# later. +@singledispatch +def _send_telemetry( + telemetry: ITelem, + telemetry_messenger: ITelemetryMessenger, + credentials_store: ICredentialsStore, +): + telemetry_messenger.send_telemetry(telemetry) + + +@_send_telemetry.register +def _( + telemetry: CredentialsTelem, + telemetry_messenger: ITelemetryMessenger, + credentials_store: ICredentialsStore, +): + credentials_store.add_credentials(telemetry.credentials) + telemetry_messenger.send_telemetry(telemetry) diff --git a/monkey/infection_monkey/typing.py b/monkey/infection_monkey/typing.py new file mode 100644 index 000000000..6e1eab9aa --- /dev/null +++ b/monkey/infection_monkey/typing.py @@ -0,0 +1,3 @@ +from typing import Iterable, Mapping + +PropagationCredentials = Mapping[str, Iterable[str]] 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 new file mode 100644 index 000000000..0b6fd8545 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_credentials_store.py @@ -0,0 +1,95 @@ +from unittest.mock import MagicMock + +import pytest + +from infection_monkey.credential_collectors import Password, SSHKeypair, Username +from infection_monkey.credential_store import AggregatingCredentialsStore +from infection_monkey.i_puppet import Credentials + +CONTROL_CHANNEL_CREDENTIALS = { + "exploit_user_list": ["Administrator", "root", "user1"], + "exploit_password_list": ["123456", "123456789", "password", "root"], + "exploit_lm_hash_list": ["aasdf23asd1fdaasadasdfas"], + "exploit_ntlm_hash_list": ["asdfadvxvsdftw3e3421234123412", "qw4trklxklvznksbhasd1231"], + "exploit_ssh_keys": [ + {"public_key": "some_public_key", "private_key": "some_private_key"}, + { + "public_key": "ssh-ed25519 AAAAC3NzEIFaJ7xH+Yoxd\n", + "private_key": "-----BEGIN OPENSSH PRIVATE KEY-----\nb3BdHIAAAAGYXjl0j66VAKruPEKjS3A=\n" + "-----END OPENSSH PRIVATE KEY-----\n", + }, + ], +} + +TEST_CREDENTIALS = [ + Credentials( + [Username("user1"), Username("user3")], + [ + Password("abcdefg"), + Password("root"), + 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"), + ], + ) +] + + +@pytest.fixture +def aggregating_credentials_store() -> AggregatingCredentialsStore: + control_channel = MagicMock() + control_channel.get_credentials_for_propagation.return_value = CONTROL_CHANNEL_CREDENTIALS + return AggregatingCredentialsStore(control_channel) + + +def test_get_credentials_from_store(aggregating_credentials_store): + actual_stored_credentials = aggregating_credentials_store.get_credentials() + + print(actual_stored_credentials) + + assert actual_stored_credentials["exploit_user_list"] == set( + CONTROL_CHANNEL_CREDENTIALS["exploit_user_list"] + ) + assert actual_stored_credentials["exploit_password_list"] == set( + CONTROL_CHANNEL_CREDENTIALS["exploit_password_list"] + ) + assert actual_stored_credentials["exploit_ntlm_hash_list"] == set( + CONTROL_CHANNEL_CREDENTIALS["exploit_ntlm_hash_list"] + ) + + for ssh_keypair in actual_stored_credentials["exploit_ssh_keys"]: + assert ssh_keypair in CONTROL_CHANNEL_CREDENTIALS["exploit_ssh_keys"] + + +def test_add_credentials_to_store(aggregating_credentials_store): + aggregating_credentials_store.add_credentials(TEST_CREDENTIALS) + aggregating_credentials_store.add_credentials(SSH_KEYS_CREDENTIALS) + + actual_stored_credentials = aggregating_credentials_store.get_credentials() + + assert actual_stored_credentials["exploit_user_list"] == set( + [ + "Administrator", + "root", + "user1", + "user3", + ] + ) + assert actual_stored_credentials["exploit_password_list"] == set( + [ + "123456", + "123456789", + "abcdefg", + "password", + "root", + ] + ) + + assert len(actual_stored_credentials["exploit_ssh_keys"]) == 3 diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py b/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py index 4bb7b4294..cf0112d59 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py @@ -14,7 +14,7 @@ INTERVAL = 0.001 def test_terminate_without_start(): - m = AutomatedMaster(None, None, None, None, MagicMock(), []) + m = AutomatedMaster(None, None, None, None, MagicMock(), [], MagicMock()) # Test that call to terminate does not raise exception m.terminate() @@ -34,7 +34,7 @@ def test_stop_if_cant_get_config_from_island(monkeypatch): monkeypatch.setattr( "infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL ) - m = AutomatedMaster(None, None, None, None, cc, []) + m = AutomatedMaster(None, None, None, None, cc, [], MagicMock()) m.start() assert cc.get_config.call_count == CHECK_FOR_CONFIG_COUNT @@ -73,7 +73,7 @@ def test_stop_if_cant_get_stop_signal_from_island(monkeypatch, sleep_and_return_ "infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL ) - m = AutomatedMaster(None, None, None, None, cc, []) + m = AutomatedMaster(None, None, None, None, cc, [], MagicMock()) m.start() assert cc.should_agent_stop.call_count == CHECK_FOR_STOP_AGENT_COUNT diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/conftest.py b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/conftest.py new file mode 100644 index 000000000..c29555262 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/conftest.py @@ -0,0 +1,18 @@ +import pytest + +from infection_monkey.telemetry.base_telem import BaseTelem + + +@pytest.fixture(scope="package") +def TestTelem(): + class InnerTestTelem(BaseTelem): + telem_category = None + __test__ = False + + def __init__(self): + pass + + def get_data(self): + return {} + + return InnerTestTelem diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_credentials_intercepting_telemetry_messenger.py b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_credentials_intercepting_telemetry_messenger.py new file mode 100644 index 000000000..7481bff34 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_credentials_intercepting_telemetry_messenger.py @@ -0,0 +1,56 @@ +from unittest.mock import MagicMock + +from infection_monkey.credential_collectors import Password, SSHKeypair, Username +from infection_monkey.i_puppet import Credentials +from infection_monkey.telemetry.credentials_telem import CredentialsTelem +from infection_monkey.telemetry.messengers.credentials_intercepting_telemetry_messenger import ( + CredentialsInterceptingTelemetryMessenger, +) + +TELEM_CREDENTIALS = [ + Credentials( + [Username("user1"), Username("user3")], + [ + Password("abcdefg"), + Password("root"), + SSHKeypair(public_key="some_public_key", private_key="some_private_key"), + ], + ) +] + + +class MockCredentialsTelem(CredentialsTelem): + def __init(self, credentials): + super().__init__(credentials) + + def get_data(self): + return {} + + +def test_credentials_generic_telemetry(TestTelem): + mock_telemetry_messenger = MagicMock() + mock_credentials_store = MagicMock() + + telemetry_messenger = CredentialsInterceptingTelemetryMessenger( + mock_telemetry_messenger, mock_credentials_store + ) + + telemetry_messenger.send_telemetry(TestTelem()) + + assert mock_telemetry_messenger.send_telemetry.called + assert not mock_credentials_store.add_credentials.called + + +def test_successful_intercepting_credentials_telemetry(): + mock_telemetry_messenger = MagicMock() + mock_credentials_store = MagicMock() + mock_empty_credentials_telem = MockCredentialsTelem(TELEM_CREDENTIALS) + + telemetry_messenger = CredentialsInterceptingTelemetryMessenger( + mock_telemetry_messenger, mock_credentials_store + ) + + telemetry_messenger.send_telemetry(mock_empty_credentials_telem) + + assert mock_telemetry_messenger.send_telemetry.called + assert mock_credentials_store.add_credentials.called diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py index f949738f6..969489107 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py @@ -2,23 +2,12 @@ from unittest.mock import MagicMock from infection_monkey.i_puppet.i_puppet import ExploiterResultData from infection_monkey.model.host import VictimHost -from infection_monkey.telemetry.base_telem import BaseTelem from infection_monkey.telemetry.exploit_telem import ExploitTelem from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import ( ExploitInterceptingTelemetryMessenger, ) -class TestTelem(BaseTelem): - telem_category = None - - def __init__(self): - pass - - def get_data(self): - return {} - - class MockExpliotTelem(ExploitTelem): def __init__(self, propagation_success): erd = ExploiterResultData() @@ -29,7 +18,7 @@ class MockExpliotTelem(ExploitTelem): return {} -def test_generic_telemetry(): +def test_generic_telemetry(TestTelem): mock_telemetry_messenger = MagicMock() mock_tunnel = MagicMock() diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py b/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py index a3d1e3f6f..13c93f60f 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py @@ -1,37 +1,51 @@ import json +import pytest + from infection_monkey.credential_collectors import Password, SSHKeypair, Username from infection_monkey.i_puppet import Credentials from infection_monkey.telemetry.credentials_telem import CredentialsTelem +USERNAME = "m0nkey" +PASSWORD = "mmm" +PUBLIC_KEY = "pub_key" +PRIVATE_KEY = "priv_key" -def test_credential_telem_send(spy_send_telemetry): - username = "m0nkey" - password = "mmm" - public_key = "pub_key" - private_key = "priv_key" + +@pytest.fixture +def credentials_for_test(): + + return Credentials( + [Username(USERNAME)], [Password(PASSWORD), SSHKeypair(PRIVATE_KEY, PUBLIC_KEY)] + ) + + +def test_credential_telem_send(spy_send_telemetry, credentials_for_test): expected_data = [ { - "identities": [{"username": username, "credential_type": "USERNAME"}], + "identities": [{"username": USERNAME, "credential_type": "USERNAME"}], "secrets": [ - {"password": password, "credential_type": "PASSWORD"}, + {"password": PASSWORD, "credential_type": "PASSWORD"}, { - "private_key": "pub_key", - "public_key": "priv_key", + "private_key": PRIVATE_KEY, + "public_key": PUBLIC_KEY, "credential_type": "SSH_KEYPAIR", }, ], } ] - credentials = Credentials( - [Username(username)], [Password(password), SSHKeypair(public_key, private_key)] - ) - - telem = CredentialsTelem([credentials]) + telem = CredentialsTelem([credentials_for_test]) telem.send() expected_data = json.dumps(expected_data, cls=telem.json_encoder) assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "credentials" + + +def test_credentials_property(credentials_for_test): + telem = CredentialsTelem([credentials_for_test]) + + assert len(list(telem.credentials)) == 1 + assert list(telem.credentials)[0] == credentials_for_test