Merge pull request #1764 from guardicore/scan-depth

Scan depth

monkey/infection_monkey/config.py

@@ -28,9 +28,6 @@ class Configuration(object):
                 continue
             if key in LOCAL_CONFIG_VARS:
                 continue
-            if self._depth_from_commandline and key == "depth":
-                self.max_depth = value
-                continue
             if hasattr(self, key):
                 setattr(self, key, value)
             else:
@@ -70,9 +67,6 @@ class Configuration(object):
 
         return result
 
-    # Used to keep track of our depth if manually specified
-    _depth_from_commandline = False
-
     ###########################
     # logging config
     ###########################

monkey/infection_monkey/exploit/HostExploiter.py

@@ -72,11 +72,13 @@ class HostExploiter:
     def exploit_host(
         self,
         host,
+        current_depth: int,
         telemetry_messenger: ITelemetryMessenger,
         agent_repository: IAgentRepository,
         options: Dict,
     ):
         self.host = host
+        self.current_depth = current_depth
         self.telemetry_messenger = telemetry_messenger
         self.agent_repository = agent_repository
         self.options = options

monkey/infection_monkey/exploit/exploiter_wrapper.py

@@ -26,10 +26,10 @@ class ExploiterWrapper:
             self._telemetry_messenger = telemetry_messenger
             self._agent_repository = agent_repository
 
-        def exploit_host(self, host: VictimHost, options: Dict):
+        def exploit_host(self, host: VictimHost, current_depth: int, options: Dict):
             exploiter = self._exploit_class()
             return exploiter.exploit_host(
-                host, self._telemetry_messenger, self._agent_repository, options
+                host, current_depth, self._telemetry_messenger, self._agent_repository, options
             )
 
     def __init__(

monkey/infection_monkey/exploit/hadoop.py

@@ -12,7 +12,6 @@ from random import SystemRandom
 import requests
 
 from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
-from infection_monkey.exploit.tools.helpers import get_monkey_depth
 from infection_monkey.exploit.tools.http_tools import HTTPTools
 from infection_monkey.exploit.web_rce import WebRCE
 from infection_monkey.model import (
@@ -95,7 +94,7 @@ class HadoopExploiter(WebRCE):
 
     def _build_command(self, path, http_path):
         # Build command to execute
-        monkey_cmd = build_monkey_commandline(self.host, get_monkey_depth() - 1)
+        monkey_cmd = build_monkey_commandline(self.host, self.current_depth - 1)
         if "linux" in self.host.os["type"]:
             base_command = HADOOP_LINUX_COMMAND
         else:

monkey/infection_monkey/exploit/log4shell.py

@@ -10,7 +10,6 @@ from infection_monkey.exploit.log4shell_utils import (
     build_exploit_bytecode,
     get_log4shell_service_exploiters,
 )
-from infection_monkey.exploit.tools.helpers import get_monkey_depth
 from infection_monkey.exploit.tools.http_tools import HTTPTools
 from infection_monkey.exploit.web_rce import WebRCE
 from infection_monkey.i_puppet.i_puppet import ExploiterResultData
@@ -114,7 +113,7 @@ class Log4ShellExploiter(WebRCE):
 
     def _build_command(self, path, http_path) -> str:
         # Build command to execute
-        monkey_cmd = build_monkey_commandline(self.host, get_monkey_depth() - 1, location=path)
+        monkey_cmd = build_monkey_commandline(self.host, self.current_depth - 1, location=path)
         if "linux" in self.host.os["type"]:
             base_command = LOG4SHELL_LINUX_COMMAND
         else:

monkey/infection_monkey/exploit/web_rce.py

@@ -5,7 +5,7 @@ from typing import List, Tuple
 
 from common.utils.attack_utils import BITS_UPLOAD_STRING, ScanStatus
 from infection_monkey.exploit.HostExploiter import HostExploiter
-from infection_monkey.exploit.tools.helpers import get_monkey_depth, get_target_monkey
+from infection_monkey.exploit.tools.helpers import get_target_monkey
 from infection_monkey.exploit.tools.http_tools import HTTPTools
 from infection_monkey.model import (
     BITSADMIN_CMDLINE_HTTP,
@@ -371,14 +371,14 @@ class WebRCE(HostExploiter):
             default_path = self.get_default_dropper_path()
             if default_path is False:
                 return False
-            monkey_cmd = build_monkey_commandline(self.host, get_monkey_depth() - 1, default_path)
+            monkey_cmd = build_monkey_commandline(self.host, self.current_depth - 1, default_path)
             command = RUN_MONKEY % {
                 "monkey_path": path,
                 "monkey_type": DROPPER_ARG,
                 "parameters": monkey_cmd,
             }
         else:
-            monkey_cmd = build_monkey_commandline(self.host, get_monkey_depth() - 1)
+            monkey_cmd = build_monkey_commandline(self.host, self.current_depth - 1)
             command = RUN_MONKEY % {
                 "monkey_path": path,
                 "monkey_type": MONKEY_ARG,

monkey/infection_monkey/i_puppet/i_puppet.py

@@ -114,12 +114,18 @@ class IPuppet(metaclass=abc.ABCMeta):
 
     @abc.abstractmethod
     def exploit_host(
-        self, name: str, host: VictimHost, options: Dict, interrupt: threading.Event
+        self,
+        name: str,
+        host: VictimHost,
+        current_depth: int,
+        options: Dict,
+        interrupt: threading.Event,
     ) -> ExploiterResultData:
         """
         Runs an exploiter against a remote host
         :param str name: The name of the exploiter to run
         :param VictimHost host: A VictimHost object representing the target to exploit
+        :param int current_depth: The current propagation depth
         :param Dict options: A dictionary containing options that modify the behavior of the
                              exploiter
         :param threading.Event interrupt: A threading.Event object that signals the exploit to stop

monkey/infection_monkey/master/automated_master.py

@@ -1,7 +1,7 @@
 import logging
 import threading
 import time
-from typing import Any, Callable, Dict, Iterable, List, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
 
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
 from infection_monkey.i_master import IMaster
@@ -30,12 +30,14 @@ logger = logging.getLogger()
 class AutomatedMaster(IMaster):
     def __init__(
         self,
+        current_depth: Optional[int],
         puppet: IPuppet,
         telemetry_messenger: ITelemetryMessenger,
         victim_host_factory: VictimHostFactory,
         control_channel: IControlChannel,
         local_network_interfaces: List[NetworkInterface],
     ):
+        self._current_depth = current_depth
         self._puppet = puppet
         self._telemetry_messenger = telemetry_messenger
         self._control_channel = control_channel
@@ -162,8 +164,11 @@ class AutomatedMaster(IMaster):
         # still running.
         credential_collector_thread.join()
 
-        if self._can_propagate():
+        current_depth = self._current_depth if self._current_depth is not None else config["depth"]
-            self._propagator.propagate(config["propagation"], self._stop)
+        logger.info(f"Current depth is {current_depth}")
+
+        if self._can_propagate() and current_depth > 0:
+            self._propagator.propagate(config["propagation"], current_depth, self._stop)
 
         payload_thread = create_daemon_thread(
             target=self._run_plugins,

monkey/infection_monkey/master/exploiter.py

@@ -34,6 +34,7 @@ class Exploiter:
         self,
         exploiter_config: Dict,
         hosts_to_exploit: Queue,
+        current_depth: int,
         results_callback: Callback,
         scan_completed: Event,
         stop: Event,
@@ -44,7 +45,14 @@ class Exploiter:
             f"{', '.join([e['name'] for e in exploiters_to_run])}"
         )
 
-        exploit_args = (exploiters_to_run, hosts_to_exploit, results_callback, scan_completed, stop)
+        exploit_args = (
+            exploiters_to_run,
+            hosts_to_exploit,
+            current_depth,
+            results_callback,
+            scan_completed,
+            stop,
+        )
         run_worker_threads(
             target=self._exploit_hosts_on_queue, args=exploit_args, num_workers=self._num_workers
         )
@@ -69,6 +77,7 @@ class Exploiter:
         self,
         exploiters_to_run: List[Dict],
         hosts_to_exploit: Queue,
+        current_depth: int,
         results_callback: Callback,
         scan_completed: Event,
         stop: Event,
@@ -78,7 +87,9 @@ class Exploiter:
         while not stop.is_set():
             try:
                 victim_host = hosts_to_exploit.get(timeout=QUEUE_TIMEOUT)
-                self._run_all_exploiters(exploiters_to_run, victim_host, results_callback, stop)
+                self._run_all_exploiters(
+                    exploiters_to_run, victim_host, current_depth, results_callback, stop
+                )
             except queue.Empty:
                 if _all_hosts_have_been_processed(scan_completed, hosts_to_exploit):
                     break
@@ -93,6 +104,7 @@ class Exploiter:
         self,
         exploiters_to_run: List[Dict],
         victim_host: VictimHost,
+        current_depth: int,
         results_callback: Callback,
         stop: Event,
     ):
@@ -100,7 +112,7 @@ class Exploiter:
         for exploiter in interruptable_iter(exploiters_to_run, stop):
             exploiter_name = exploiter["name"]
             exploiter_results = self._run_exploiter(
-                exploiter_name, exploiter["options"], victim_host, stop
+                exploiter_name, exploiter["options"], victim_host, current_depth, stop
             )
             results_callback(exploiter_name, victim_host, exploiter_results)
 
@@ -108,7 +120,12 @@ class Exploiter:
                 break
 
     def _run_exploiter(
-        self, exploiter_name: str, options: Dict, victim_host: VictimHost, stop: Event
+        self,
+        exploiter_name: str,
+        options: Dict,
+        victim_host: VictimHost,
+        current_depth: int,
+        stop: Event,
     ) -> ExploiterResultData:
         logger.debug(f"Attempting to use {exploiter_name} on {victim_host.ip_addr}")
 
@@ -116,7 +133,9 @@ class Exploiter:
         options = {"credentials": credentials, **options}
 
         try:
-            return self._puppet.exploit_host(exploiter_name, victim_host, options, stop)
+            return self._puppet.exploit_host(
+                exploiter_name, victim_host, current_depth, options, stop
+            )
         except Exception as ex:
             msg = (
                 f"An unexpected error occurred while exploiting {victim_host.ip_addr} with "

monkey/infection_monkey/master/mock_master.py

@@ -101,13 +101,13 @@ class MockMaster(IMaster):
 
     def _exploit(self):
         logger.info("Exploiting victims")
-        result = self._puppet.exploit_host("PowerShellExploiter", "10.0.0.1", {}, None)
+        result = self._puppet.exploit_host("PowerShellExploiter", "10.0.0.1", 0, {}, None)
         logger.info(f"Attempts for exploiting {result.attempts}")
         self._telemetry_messenger.send_telemetry(
             ExploitTelem("PowerShellExploiter", self._hosts["10.0.0.1"], result)
         )
 
-        result = self._puppet.exploit_host("SSHExploiter", "10.0.0.3", {}, None)
+        result = self._puppet.exploit_host("SSHExploiter", "10.0.0.3", 0, {}, None)
         logger.info(f"Attempts for exploiting {result.attempts}")
         self._telemetry_messenger.send_telemetry(
             ExploitTelem("SSHExploiter", self._hosts["10.0.0.3"], result)

monkey/infection_monkey/master/propagator.py

@@ -39,7 +39,7 @@ class Propagator:
         self._local_network_interfaces = local_network_interfaces
         self._hosts_to_exploit = None
 
-    def propagate(self, propagation_config: Dict, stop: Event):
+    def propagate(self, propagation_config: Dict, current_depth: int, stop: Event):
         logger.info("Attempting to propagate")
 
         network_scan_completed = Event()
@@ -50,7 +50,7 @@ class Propagator:
         )
         exploit_thread = create_daemon_thread(
             target=self._exploit_hosts,
-            args=(propagation_config, network_scan_completed, stop),
+            args=(propagation_config, current_depth, network_scan_completed, stop),
         )
 
         scan_thread.start()
@@ -134,6 +134,7 @@ class Propagator:
     def _exploit_hosts(
         self,
         propagation_config: Dict,
+        current_depth: int,
         network_scan_completed: Event,
         stop: Event,
     ):
@@ -143,6 +144,7 @@ class Propagator:
         self._exploiter.exploit_hosts(
             exploiter_config,
             self._hosts_to_exploit,
+            current_depth,
             self._process_exploit_attempts,
             network_scan_completed,
             stop,

monkey/infection_monkey/monkey.py

@@ -68,6 +68,7 @@ class InfectionMonkey:
         # TODO used in propogation phase
         self._monkey_inbound_tunnel = None
         self.telemetry_messenger = LegacyTelemetryMessengerAdapter()
+        self._current_depth = self._opts.depth
 
     @staticmethod
     def _get_arguments(args):
@@ -93,7 +94,6 @@ class InfectionMonkey:
 
         logger.info("Monkey is starting...")
 
-        self._set_propagation_depth(self._opts)
         self._add_default_server_to_config(self._opts.server)
         self._connect_to_island()
 
@@ -111,14 +111,6 @@ class InfectionMonkey:
         self._setup()
         self._master.start()
 
-    @staticmethod
-    def _set_propagation_depth(options):
-        if options.depth is not None:
-            WormConfiguration._depth_from_commandline = True
-            WormConfiguration.depth = options.depth
-            logger.debug("Setting propagation depth from command line")
-        logger.debug(f"Set propagation depth to {WormConfiguration.depth}")
-
     @staticmethod
     def _add_default_server_to_config(default_server: str):
         if default_server:
@@ -179,6 +171,7 @@ class InfectionMonkey:
         )
 
         self._master = AutomatedMaster(
+            self._current_depth,
             puppet,
             telemetry_messenger,
             victim_host_factory,

monkey/infection_monkey/puppet/mock_puppet.py

@@ -137,7 +137,12 @@ class MockPuppet(IPuppet):
 
     # TODO: host should be VictimHost, at the moment it can't because of circular dependency
     def exploit_host(
-        self, name: str, host: VictimHost, options: Dict, interrupt: threading.Event
+        self,
+        name: str,
+        host: VictimHost,
+        current_depth: int,
+        options: Dict,
+        interrupt: threading.Event,
     ) -> ExploiterResultData:
         logger.debug(f"exploit_hosts({name}, {host}, {options})")
         attempts = [

monkey/infection_monkey/puppet/puppet.py

@@ -58,10 +58,15 @@ class Puppet(IPuppet):
         return fingerprinter.get_host_fingerprint(host, ping_scan_data, port_scan_data, options)
 
     def exploit_host(
-        self, name: str, host: VictimHost, options: Dict, interrupt: threading.Event
+        self,
+        name: str,
+        host: VictimHost,
+        current_depth: int,
+        options: Dict,
+        interrupt: threading.Event,
     ) -> ExploiterResultData:
         exploiter = self._plugin_registry.get_plugin(name, PluginType.EXPLOITER)
-        return exploiter.exploit_host(host, options)
+        return exploiter.exploit_host(host, current_depth, options)
 
     def run_payload(self, name: str, options: Dict, interrupt: threading.Event):
         payload = self._plugin_registry.get_plugin(name, PluginType.PAYLOAD)

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, MagicMock(), [])
+    m = AutomatedMaster(None, None, None, None, 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, cc, [])
+    m = AutomatedMaster(None, None, None, None, cc, [])
     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, cc, [])
+    m = AutomatedMaster(None, None, None, None, cc, [])
     m.start()
 
     assert cc.should_agent_stop.call_count == CHECK_FOR_STOP_AGENT_COUNT

monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py

@@ -74,7 +74,7 @@ def run_exploiters(exploiter_config, hosts_to_exploit, callback, scan_completed,
         scan_completed.set()
 
         e = Exploiter(puppet, num_workers, get_credentials_for_propagation)
-        e.exploit_hosts(exploiter_config, hosts_to_exploit, callback, scan_completed, stop)
+        e.exploit_hosts(exploiter_config, hosts_to_exploit, 1, callback, scan_completed, stop)
 
     return inner
 
@@ -102,7 +102,7 @@ def test_credentials_passed_to_exploiter(run_exploiters):
     run_exploiters(mock_puppet, 1)
 
     for call_args in mock_puppet.exploit_host.call_args_list:
-        assert call_args[0][2].get("credentials") == CREDENTIALS_FOR_PROPAGATION
+        assert call_args[0][3].get("credentials") == CREDENTIALS_FOR_PROPAGATION
 
 
 def test_stop_after_callback(exploiter_config, callback, scan_completed, stop, hosts_to_exploit):
@@ -121,7 +121,7 @@ def test_stop_after_callback(exploiter_config, callback, scan_completed, stop, h
     # Intentionally NOT setting scan_completed.set(); _callback() will set stop
 
     e = Exploiter(MockPuppet(), callback_barrier_count + 2, get_credentials_for_propagation)
-    e.exploit_hosts(exploiter_config, hosts_to_exploit, stoppable_callback, scan_completed, stop)
+    e.exploit_hosts(exploiter_config, hosts_to_exploit, 1, stoppable_callback, scan_completed, stop)
 
     assert stoppable_callback.call_count == 2
 

monkey/tests/unit_tests/infection_monkey/master/test_propagator.py

@@ -124,7 +124,13 @@ def mock_ip_scanner():
 
 class StubExploiter:
     def exploit_hosts(
-        self, hosts_to_exploit, exploiter_config, results_callback, scan_completed, stop
+        self,
+        exploiters_to_run,
+        hosts_to_exploit,
+        current_depth,
+        results_callback,
+        scan_completed,
+        stop,
     ):
         pass
 
@@ -144,6 +150,7 @@ def test_scan_result_processing(telemetry_messenger_spy, mock_ip_scanner, mock_v
             "network_scan": {},  # This is empty since MockIPscanner ignores it
             "exploiters": {},  # This is empty since StubExploiter ignores it
         },
+        1,
         Event(),
     )
 
@@ -174,7 +181,13 @@ def test_scan_result_processing(telemetry_messenger_spy, mock_ip_scanner, mock_v
 
 class MockExploiter:
     def exploit_hosts(
-        self, exploiter_config, hosts_to_exploit, results_callback, scan_completed, stop
+        self,
+        exploiters_to_run,
+        hosts_to_exploit,
+        current_depth,
+        results_callback,
+        scan_completed,
+        stop,
     ):
         scan_completed.wait()
         hte = []
@@ -240,6 +253,7 @@ def test_exploiter_result_processing(
             "network_scan": {},  # This is empty since MockIPscanner ignores it
             "exploiters": {},  # This is empty since MockExploiter ignores it
         },
+        1,
         Event(),
     )
 
@@ -284,6 +298,7 @@ def test_scan_target_generation(telemetry_messenger_spy, mock_ip_scanner, mock_v
             "network_scan": {},  # This is empty since MockIPscanner ignores it
             "exploiters": {},  # This is empty since MockExploiter ignores it
         },
+        1,
         Event(),
     )
     expected_ip_scan_list = [