Merge pull request #2033 from guardicore/fix-depth-logic

Agent: Change the logic of depth to represent current depth

CHANGELOG.md

@@ -38,6 +38,8 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
 - Update MongoDB version to 4.4.x. #1924
 - Endpoint to get agent binaries from "/api/agent/download/<string:os>" to
   "/api/agent-binaries/<string:os>". #1978
+- Depth flag (-d) on the agent now acts the  way you would expect(it represents
+  the current depth of the agent, not hops remaining). #2033
 - Agent configuration structure. #1996, #1998, #1961, #1997, #1994, #1741,
   #1761, #1695, #1605, #2028
 

monkey/common/utils/argparse_types.py

@@ -0,0 +1,9 @@
+import argparse
+
+
+def positive_int(_input: str):
+    int_value = int(_input)
+    if int_value < 0:
+        raise argparse.ArgumentTypeError(f"{_input} is not a positive integer")
+
+    return int_value

monkey/infection_monkey/config.py

@@ -60,7 +60,7 @@ class Configuration(object):
     ###########################
 
     # depth of propagation
-    depth = 2
+    depth = 0
     max_depth = None
 
     keep_tunnel_open_time = 30

monkey/infection_monkey/dropper.py

@@ -9,6 +9,7 @@ import sys
 import time
 from pathlib import PosixPath, WindowsPath
 
+from common.utils.argparse_types import positive_int
 from common.utils.attack_utils import UsageEnum
 from infection_monkey.utils.commands import (
     build_monkey_commandline_explicitly,
@@ -45,7 +46,7 @@ class MonkeyDrops(object):
         arg_parser.add_argument("-p", "--parent")
         arg_parser.add_argument("-t", "--tunnel")
         arg_parser.add_argument("-s", "--server")
-        arg_parser.add_argument("-d", "--depth", type=int)
+        arg_parser.add_argument("-d", "--depth", type=positive_int, default=0)
         arg_parser.add_argument("-l", "--location")
         arg_parser.add_argument("-vp", "--vulnerable-port")
         self.opts = arg_parser.parse_args(args)

monkey/infection_monkey/exploit/hadoop.py

@@ -104,7 +104,7 @@ class HadoopExploiter(WebRCE):
 
     def _build_command(self, path, http_path):
         # Build command to execute
-        monkey_cmd = build_monkey_commandline(self.host, self.current_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

@@ -114,7 +114,7 @@ class Log4ShellExploiter(WebRCE):
 
     def _build_command(self, path: PurePath, http_path) -> str:
         # Build command to execute
-        monkey_cmd = build_monkey_commandline(self.host, self.current_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/mssqlexec.py

@@ -178,7 +178,7 @@ class MSSQLExploiter(HostExploiter):
 
     def _build_agent_launch_command(self, agent_path_on_victim: PureWindowsPath) -> str:
         agent_args = build_monkey_commandline(
-            self.host, self.current_depth - 1, agent_path_on_victim
+            self.host, self.current_depth + 1, agent_path_on_victim
         )
 
         return f"{agent_path_on_victim} {DROPPER_ARG} {agent_args}"

monkey/infection_monkey/exploit/powershell.py

@@ -168,7 +168,7 @@ class PowerShellExploiter(HostExploiter):
 
     def _run_monkey_executable_on_victim(self, executable_path):
         monkey_execution_command = build_monkey_execution_command(
-            self.host, self.current_depth - 1, executable_path
+            self.host, self.current_depth + 1, executable_path
         )
 
         logger.info(

monkey/infection_monkey/exploit/smbexec.py

@@ -91,13 +91,13 @@ class SMBExploiter(HostExploiter):
                 "dropper_path": remote_full_path
             } + build_monkey_commandline(
                 self.host,
-                self.current_depth - 1,
+                self.current_depth + 1,
                 str(dest_path),
             )
         else:
             cmdline = MONKEY_CMDLINE_DETACHED_WINDOWS % {
                 "monkey_path": remote_full_path
-            } + build_monkey_commandline(self.host, self.current_depth - 1)
+            } + build_monkey_commandline(self.host, self.current_depth + 1)
 
         smb_conn = None
         for str_bind_format, port in SMBExploiter.KNOWN_PROTOCOLS.values():

monkey/infection_monkey/exploit/sshexec.py

@@ -242,7 +242,7 @@ class SSHExploiter(HostExploiter):
 
         try:
             cmdline = f"{monkey_path_on_victim} {MONKEY_ARG}"
-            cmdline += build_monkey_commandline(self.host, self.current_depth - 1)
+            cmdline += build_monkey_commandline(self.host, self.current_depth + 1)
             cmdline += " > /dev/null 2>&1 &"
             ssh.exec_command(cmdline, timeout=SSH_EXEC_TIMEOUT)
 

monkey/infection_monkey/exploit/web_rce.py

@@ -369,14 +369,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, self.current_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, self.current_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/exploit/wmiexec.py

@@ -96,13 +96,13 @@ class WmiExploiter(HostExploiter):
                     "dropper_path": remote_full_path
                 } + build_monkey_commandline(
                     self.host,
-                    self.current_depth - 1,
+                    self.current_depth + 1,
                     DROPPER_TARGET_PATH_WIN64,
                 )
             else:
                 cmdline = MONKEY_CMDLINE_WINDOWS % {
                     "monkey_path": remote_full_path
-                } + build_monkey_commandline(self.host, self.current_depth - 1)
+                } + build_monkey_commandline(self.host, self.current_depth + 1)
 
             # execute the remote monkey
             result = WmiTools.get_object(wmi_connection, "Win32_Process").Create(

monkey/infection_monkey/master/automated_master.py

@@ -13,6 +13,7 @@ from infection_monkey.network import NetworkInterface
 from infection_monkey.telemetry.credentials_telem import CredentialsTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
+from infection_monkey.utils.propagation import should_propagate
 from infection_monkey.utils.threading import create_daemon_thread, interruptible_iter
 
 from . import Exploiter, IPScanner, Propagator
@@ -169,11 +170,13 @@ class AutomatedMaster(IMaster):
         # still running.
         credential_collector_thread.join()
 
-        current_depth = self._current_depth if self._current_depth is not None else config["depth"]
+        current_depth = self._current_depth if self._current_depth is not None else 0
         logger.info(f"Current depth is {current_depth}")
 
-        if self._can_propagate() and current_depth > 0:
+        if should_propagate(self._control_channel.get_config(), self._current_depth):
             self._propagator.propagate(config["propagation"], current_depth, self._stop)
+        else:
+            logger.info("Skipping propagation: maximum depth reached")
 
         payload_thread = create_daemon_thread(
             target=self._run_plugins,
@@ -200,9 +203,6 @@ class AutomatedMaster(IMaster):
         for pba_data in self._puppet.run_pba(name, options):
             self._telemetry_messenger.send_telemetry(PostBreachTelem(pba_data))
 
-    def _can_propagate(self) -> bool:
-        return True
-
     def _run_payload(self, payload: Tuple[str, Dict]):
         name = payload[0]
         options = payload[1]

monkey/infection_monkey/monkey.py

@@ -8,6 +8,7 @@ from typing import List
 
 import infection_monkey.tunnel as tunnel
 from common.network.network_utils import address_to_ip_port
+from common.utils.argparse_types import positive_int
 from common.utils.attack_utils import ScanStatus, UsageEnum
 from common.version import get_version
 from infection_monkey.config import GUID
@@ -77,6 +78,7 @@ from infection_monkey.utils.monkey_dir import (
     remove_monkey_dir,
 )
 from infection_monkey.utils.monkey_log_path import get_agent_log_path
+from infection_monkey.utils.propagation import should_propagate
 from infection_monkey.utils.signal_handler import register_signal_handlers, reset_signal_handlers
 
 logger = logging.getLogger(__name__)
@@ -97,6 +99,7 @@ class InfectionMonkey:
         self._telemetry_messenger = LegacyTelemetryMessengerAdapter()
         self._current_depth = self._opts.depth
         self._master = None
+        self._inbound_tunnel_opened = False
 
     @staticmethod
     def _get_arguments(args):
@@ -104,7 +107,7 @@ class InfectionMonkey:
         arg_parser.add_argument("-p", "--parent")
         arg_parser.add_argument("-t", "--tunnel")
         arg_parser.add_argument("-s", "--server")
-        arg_parser.add_argument("-d", "--depth", type=int)
+        arg_parser.add_argument("-d", "--depth", type=positive_int, default=0)
         opts = arg_parser.parse_args(args)
         InfectionMonkey._log_arguments(opts)
 
@@ -166,7 +169,11 @@ class InfectionMonkey:
             firewall.add_firewall_rule()
 
         self._monkey_inbound_tunnel = self._control_client.create_control_tunnel()
-        if self._monkey_inbound_tunnel and self._propagation_enabled():
+        config = ControlChannel(
+            self._control_client.server_address, GUID, self._control_client.proxies
+        ).get_config()
+        if self._monkey_inbound_tunnel and should_propagate(config, self._current_depth):
+            self._inbound_tunnel_opened = True
             self._monkey_inbound_tunnel.start()
 
         StateTelem(is_done=False, version=get_version()).send()
@@ -353,7 +360,7 @@ class InfectionMonkey:
 
             reset_signal_handlers()
 
-            if self._monkey_inbound_tunnel and self._propagation_enabled():
+            if self._inbound_tunnel_opened:
                 self._monkey_inbound_tunnel.stop()
                 self._monkey_inbound_tunnel.join()
 
@@ -378,12 +385,6 @@ class InfectionMonkey:
 
         logger.info("Monkey is shutting down")
 
-    def _propagation_enabled(self) -> bool:
-        # If self._current_depth is None, assume that propagation is desired.
-        # The Master will ignore this value if it is None and pull the actual
-        # maximum depth from the server
-        return self._current_depth is None or self._current_depth > 0
-
     def _close_tunnel(self):
         tunnel_address = (
             self._control_client.proxies.get("https", "").replace("http://", "").split(":")[0]

monkey/infection_monkey/utils/commands.py

@@ -40,8 +40,6 @@ def build_monkey_commandline_explicitly(
         cmdline.append("-s")
         cmdline.append(str(server))
     if depth is not None:
-        if int(depth) < 0:
-            depth = 0
         cmdline.append("-d")
         cmdline.append(str(depth))
     if location is not None:

monkey/infection_monkey/utils/propagation.py

@@ -0,0 +1,2 @@
+def should_propagate(config: dict, current_depth: int) -> bool:
+    return config["config"]["depth"] > current_depth

monkey/tests/unit_tests/infection_monkey/utils/test_commands.py

@@ -28,16 +28,6 @@ def test_build_monkey_commandline_explicitly_arguments():
     assert expected == actual
 
 
-def test_build_monkey_commandline_explicitly_depth_condition_less():
-    expected = [
-        "-d",
-        "0",
-    ]
-    actual = build_monkey_commandline_explicitly(depth=-50)
-
-    assert expected == actual
-
-
 def test_build_monkey_commandline_explicitly_depth_condition_greater():
     expected = [
         "-d",

monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py

@@ -0,0 +1,32 @@
+from infection_monkey.utils.propagation import should_propagate
+
+
+def get_config(max_depth):
+    return {"config": {"depth": max_depth}}
+
+
+def test_should_propagate_current_less_than_max():
+    max_depth = 2
+    current_depth = 1
+
+    config = get_config(max_depth)
+
+    assert should_propagate(config, current_depth) is True
+
+
+def test_should_propagate_current_greater_than_max():
+    max_depth = 2
+    current_depth = 3
+
+    config = get_config(max_depth)
+
+    assert should_propagate(config, current_depth) is False
+
+
+def test_should_propagate_current_equal_to_max():
+    max_depth = 2
+    current_depth = max_depth
+
+    config = get_config(max_depth)
+
+    assert should_propagate(config, current_depth) is False