diff --git a/monkey/agent_plugins/exploiters/hadoop/src/plugin.py b/monkey/agent_plugins/exploiters/hadoop/src/plugin.py index 9e49b4280f1..ec860d53ad5 100644 --- a/monkey/agent_plugins/exploiters/hadoop/src/plugin.py +++ b/monkey/agent_plugins/exploiters/hadoop/src/plugin.py @@ -16,7 +16,7 @@ from common.utils.code_utils import del_key # dependencies to get rid of or internalize -from infection_monkey.exploit import IAgentBinaryRepository +from infection_monkey.exploit import IAgentBinaryRepository, IAgentOTPProvider from infection_monkey.exploit.tools.http_agent_binary_server import start_agent_binary_server from infection_monkey.i_puppet import ExploiterResultData, TargetHost from infection_monkey.network import TCPPortSelector @@ -37,6 +37,7 @@ def __init__( agent_event_publisher: IAgentEventPublisher, agent_binary_repository: IAgentBinaryRepository, tcp_port_selector: TCPPortSelector, + otp_provider: IAgentOTPProvider, **kwargs, ): hadoop_exploit_client = HadoopExploitClient(agent_id, agent_event_publisher) diff --git a/monkey/infection_monkey/exploit/HostExploiter.py b/monkey/infection_monkey/exploit/HostExploiter.py index a6c08978d4d..6f7bdbff6fb 100644 --- a/monkey/infection_monkey/exploit/HostExploiter.py +++ b/monkey/infection_monkey/exploit/HostExploiter.py @@ -8,6 +8,7 @@ from common.event_queue import IAgentEventQueue from common.types import Event from common.utils.exceptions import FailedExploitationError +from infection_monkey.exploit import IAgentOTPProvider from infection_monkey.i_puppet import ExploiterResultData, TargetHost from infection_monkey.network import TCPPortSelector from infection_monkey.utils.ids import get_agent_id @@ -77,6 +78,7 @@ def exploit_host( tcp_port_selector: TCPPortSelector, options: Dict, interrupt: Event, + otp_provider: IAgentOTPProvider, ): self.host = host self.servers = servers @@ -86,6 +88,7 @@ def exploit_host( self.tcp_port_selector = tcp_port_selector self.options = options self.interrupt = interrupt + self.otp_provider = otp_provider self.pre_exploit() try: diff --git a/monkey/infection_monkey/exploit/__init__.py b/monkey/infection_monkey/exploit/__init__.py index 29c6ba75685..b329b8ca87d 100644 --- a/monkey/infection_monkey/exploit/__init__.py +++ b/monkey/infection_monkey/exploit/__init__.py @@ -1,4 +1,5 @@ from .i_agent_binary_repository import IAgentBinaryRepository, RetrievalError from .caching_agent_binary_repository import CachingAgentBinaryRepository -from .exploiter_wrapper import ExploiterWrapper from .island_api_agent_otp_provider import IslandAPIAgentOTPProvider +from .i_agent_otp_provider import IAgentOTPProvider +from .exploiter_wrapper import ExploiterWrapper diff --git a/monkey/infection_monkey/exploit/exploiter_wrapper.py b/monkey/infection_monkey/exploit/exploiter_wrapper.py index e2b42b21942..a9037d82aeb 100644 --- a/monkey/infection_monkey/exploit/exploiter_wrapper.py +++ b/monkey/infection_monkey/exploit/exploiter_wrapper.py @@ -5,7 +5,7 @@ from infection_monkey.i_puppet import TargetHost from infection_monkey.network import TCPPortSelector -from . import IAgentBinaryRepository +from . import IAgentBinaryRepository, IAgentOTPProvider from .HostExploiter import HostExploiter @@ -24,11 +24,13 @@ def __init__( event_queue: IAgentEventQueue, agent_binary_repository: IAgentBinaryRepository, tcp_port_selector: TCPPortSelector, + otp_provider: IAgentOTPProvider, ): self._exploit_class = exploit_class self._event_queue = event_queue self._agent_binary_repository = agent_binary_repository self._tcp_port_selector = tcp_port_selector + self._otp_provider = otp_provider def run( self, @@ -48,6 +50,7 @@ def run( self._tcp_port_selector, options, interrupt, + self._otp_provider, ) def __init__( @@ -55,12 +58,18 @@ def __init__( event_queue: IAgentEventQueue, agent_binary_repository: IAgentBinaryRepository, tcp_port_selector: TCPPortSelector, + otp_provider: IAgentOTPProvider, ): self._event_queue = event_queue self._agent_binary_repository = agent_binary_repository self._tcp_port_selector = tcp_port_selector + self._otp_provider = otp_provider def wrap(self, exploit_class: Type[HostExploiter]): return ExploiterWrapper.Inner( - exploit_class, self._event_queue, self._agent_binary_repository, self._tcp_port_selector + exploit_class, + self._event_queue, + self._agent_binary_repository, + self._tcp_port_selector, + self._otp_provider, ) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index c8ad16c08b7..69923aa40ba 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -46,7 +46,11 @@ MimikatzCredentialCollector, SSHCredentialCollector, ) -from infection_monkey.exploit import CachingAgentBinaryRepository, ExploiterWrapper +from infection_monkey.exploit import ( + IslandAPIAgentOTPProvider, + CachingAgentBinaryRepository, + ExploiterWrapper, +) from infection_monkey.exploit.log4shell import Log4ShellExploiter from infection_monkey.exploit.mssqlexec import MSSQLExploiter from infection_monkey.exploit.powershell import PowerShellExploiter @@ -342,18 +346,21 @@ def _build_puppet(self, operating_system: OperatingSystem) -> IPuppet: manager=self._manager, ) + plugin_source_extractor = PluginSourceExtractor(self._plugin_dir) plugin_loader = PluginLoader( self._plugin_dir, partial(configure_child_process_logger, self._ipc_logger_queue) ) + otp_provider = IslandAPIAgentOTPProvider(self._island_api_client) plugin_registry = PluginRegistry( operating_system, self._island_api_client, - PluginSourceExtractor(self._plugin_dir), + plugin_source_extractor, plugin_loader, agent_binary_repository, self._agent_event_publisher, self._propagation_credentials_repository, - tcp_port_selector=self._tcp_port_selector, + self._tcp_port_selector, + otp_provider, ) plugin_compatability_verifier = PluginCompatabilityVerifier( self._island_api_client, HARD_CODED_EXPLOITER_MANIFESTS @@ -377,7 +384,7 @@ def _build_puppet(self, operating_system: OperatingSystem) -> IPuppet: puppet.load_plugin(AgentPluginType.FINGERPRINTER, "ssh", SSHFingerprinter()) exploit_wrapper = ExploiterWrapper( - self._agent_event_queue, agent_binary_repository, self._tcp_port_selector + self._agent_event_queue, agent_binary_repository, self._tcp_port_selector, otp_provider ) puppet.load_plugin( diff --git a/monkey/infection_monkey/puppet/plugin_registry.py b/monkey/infection_monkey/puppet/plugin_registry.py index 52f7d746e3a..b7a34a2df0a 100644 --- a/monkey/infection_monkey/puppet/plugin_registry.py +++ b/monkey/infection_monkey/puppet/plugin_registry.py @@ -9,7 +9,7 @@ from common import OperatingSystem from common.agent_plugins import AgentPlugin, AgentPluginType from common.event_queue import IAgentEventPublisher -from infection_monkey.exploit import IAgentBinaryRepository +from infection_monkey.exploit import IAgentBinaryRepository, IAgentOTPProvider from infection_monkey.i_puppet import UnknownPluginError from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIRequestError from infection_monkey.network import TCPPortSelector @@ -34,6 +34,7 @@ def __init__( agent_event_publisher: IAgentEventPublisher, propagation_credentials_repository: IPropagationCredentialsRepository, tcp_port_selector: TCPPortSelector, + otp_provider: IAgentOTPProvider, ): """ `self._registry` looks like - @@ -54,6 +55,7 @@ def __init__( self._agent_event_publisher = agent_event_publisher self._propagation_credentials_repository = propagation_credentials_repository self._tcp_port_selector = tcp_port_selector + self._otp_provider = otp_provider self._agent_id = get_agent_id() self._lock = RLock() @@ -81,6 +83,7 @@ def _load_plugin_from_island(self, plugin_name: str, plugin_type: AgentPluginTyp agent_event_publisher=self._agent_event_publisher, propagation_credentials_repository=self._propagation_credentials_repository, tcp_port_selector=self._tcp_port_selector, + otp_provider=self._otp_provider, ) self.load_plugin(plugin_type, plugin_name, multiprocessing_plugin) diff --git a/monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_plugin.py b/monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_plugin.py index fd25503b87a..d875f9e2469 100644 --- a/monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_plugin.py +++ b/monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_plugin.py @@ -46,6 +46,7 @@ def plugin(monkeypatch) -> Plugin: agent_event_publisher=MagicMock(), agent_binary_repository=MagicMock(), tcp_port_selector=MagicMock(), + otp_provider=MagicMock(), ) @@ -86,6 +87,7 @@ def test_run__exploit_host_raises_exception(monkeypatch, plugin: Plugin): agent_event_publisher=MagicMock(), agent_binary_repository=MagicMock(), tcp_port_selector=MagicMock(), + otp_provider=MagicMock(), ) result = plugin.run( host=TARGET_HOST, diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py index ef0c2311bef..679d175959d 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py @@ -61,6 +61,7 @@ def powershell_arguments(host_with_ip_address): "agent_binary_repository": mock_agent_binary_repository, "tcp_port_selector": MagicMock(), "interrupt": threading.Event(), + "otp_provider": MagicMock(), } return arguments diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_plugin_registry.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_plugin_registry.py index 5bdd979481a..a645f1445cc 100644 --- a/monkey/tests/unit_tests/infection_monkey/puppet/test_plugin_registry.py +++ b/monkey/tests/unit_tests/infection_monkey/puppet/test_plugin_registry.py @@ -6,7 +6,7 @@ from common import OperatingSystem from common.agent_plugins import AgentPlugin, AgentPluginManifest, AgentPluginType from common.event_queue import IAgentEventPublisher -from infection_monkey.exploit import IAgentBinaryRepository +from infection_monkey.exploit import IAgentBinaryRepository, IAgentOTPProvider from infection_monkey.i_puppet import UnknownPluginError from infection_monkey.island_api_client import ( IIslandAPIClient, @@ -48,6 +48,11 @@ def dummy_tcp_port_selector() -> TCPPortSelector: return MagicMock(spec=TCPPortSelector) +@pytest.fixture +def dummy_otp_provider() -> IAgentOTPProvider: + return MagicMock(spec=IAgentOTPProvider) + + @pytest.mark.parametrize( "error_raised_by_island_api_client, error_raised_by_plugin_registry", [(IslandAPIRequestError, UnknownPluginError), (IslandAPIError, IslandAPIError)], @@ -59,6 +64,7 @@ def test_get_plugin__error_handling( dummy_agent_event_publisher: IAgentEventPublisher, dummy_propagation_credentials_repository: IPropagationCredentialsRepository, dummy_tcp_port_selector: TCPPortSelector, + dummy_otp_provider: IAgentOTPProvider, error_raised_by_island_api_client: Exception, error_raised_by_plugin_registry: Exception, ): @@ -75,6 +81,7 @@ def test_get_plugin__error_handling( dummy_agent_event_publisher, dummy_propagation_credentials_repository, dummy_tcp_port_selector, + dummy_otp_provider, ) with pytest.raises(error_raised_by_plugin_registry): @@ -128,6 +135,7 @@ def plugin_registry( dummy_agent_event_publisher: IAgentEventPublisher, dummy_propagation_credentials_repository: IPropagationCredentialsRepository, dummy_tcp_port_selector: TCPPortSelector, + dummy_otp_provider: IAgentOTPProvider, ) -> PluginRegistry: return PluginRegistry( OperatingSystem.LINUX, @@ -138,6 +146,7 @@ def plugin_registry( dummy_agent_event_publisher, dummy_propagation_credentials_repository, dummy_tcp_port_selector, + dummy_otp_provider, ) diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py index 9e83d0fcfe1..1b575138c3d 100644 --- a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py +++ b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py @@ -29,6 +29,7 @@ def mock_plugin_registry() -> PluginRegistry: MagicMock(), MagicMock(), MagicMock(), + MagicMock(), )