Agent: Extract _build_monkey_execution_command() into powershell_utils

monkey/infection_monkey/exploit/powershell.py

@@ -14,8 +14,7 @@ from infection_monkey.exploit.consts import WIN_ARCH_32, WIN_ARCH_64
 from infection_monkey.exploit.HostExploiter import HostExploiter
 from infection_monkey.exploit.powershell_utils import utils
 from infection_monkey.exploit.tools.helpers import get_monkey_depth, get_target_monkey_by_os
-from infection_monkey.model import DROPPER_ARG, GET_ARCH_WINDOWS, RUN_MONKEY, VictimHost
-from infection_monkey.utils.commands import build_monkey_commandline
+from infection_monkey.model import GET_ARCH_WINDOWS, VictimHost
 from infection_monkey.utils.environment import is_windows_os
 
 LOG = logging.getLogger(__name__)
@@ -137,7 +136,9 @@ class PowerShellExploiter(HostExploiter):
             os.remove(TEMP_MONKEY_BINARY_FILEPATH)
 
     def _run_monkey_executable_on_victim(self, executable_path) -> None:
-        monkey_execution_command = self._build_monkey_execution_command(executable_path)
+        monkey_execution_command = utils.build_monkey_execution_command(
+            self.host, get_monkey_depth() - 1, executable_path
+        )
 
         with self.client.wsman, RunspacePool(self.client.wsman) as pool:
             ps = PowerShell(pool)
@@ -145,17 +146,3 @@ class PowerShellExploiter(HostExploiter):
                 "name", "create"
             ).add_parameter("ArgumentList", monkey_execution_command)
             ps.invoke()
-
-    def _build_monkey_execution_command(self, executable_path) -> str:
-        monkey_params = build_monkey_commandline(
-            target_host=self.host,
-            depth=get_monkey_depth() - 1,
-            vulnerable_port=None,
-            location=executable_path,
-        )
-
-        return RUN_MONKEY % {
-            "monkey_path": executable_path,
-            "monkey_type": DROPPER_ARG,
-            "parameters": monkey_params,
-        }

monkey/infection_monkey/exploit/powershell_utils/utils.py

@@ -1,6 +1,9 @@
 from itertools import product
 from typing import List, Optional, Tuple
 
+from infection_monkey.model import DROPPER_ARG, RUN_MONKEY, VictimHost
+from infection_monkey.utils.commands import build_monkey_commandline
+
 AUTH_BASIC = "basic"
 AUTH_NEGOTIATE = "negotiate"
 ENCRYPTION_AUTO = "auto"
@@ -54,3 +57,18 @@ def get_powershell_client_params(password: str) -> Tuple[bool, str, str]:
     encryption = ENCRYPTION_AUTO if password != "" else ENCRYPTION_NEVER
 
     return (ssl, auth, encryption)
+
+
+def build_monkey_execution_command(host: VictimHost, depth: int, executable_path: str) -> str:
+    monkey_params = build_monkey_commandline(
+        target_host=host,
+        depth=depth,
+        vulnerable_port=None,
+        location=executable_path,
+    )
+
+    return RUN_MONKEY % {
+        "monkey_path": executable_path,
+        "monkey_type": DROPPER_ARG,
+        "parameters": monkey_params,
+    }

monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/test_utils.py

@@ -1,4 +1,5 @@
 from infection_monkey.exploit.powershell_utils import utils
+from infection_monkey.model.host import VictimHost
 
 TEST_USERS = ["user1", "user2"]
 TEST_PASSWORDS = ["p1", "p2"]
@@ -66,3 +67,14 @@ def test_get_powershell_client_params__password_empty():
     assert ssl is False
     assert auth == utils.AUTH_BASIC
     assert encryption == utils.ENCRYPTION_NEVER
+
+
+def test_build_monkey_execution_command():
+    host = VictimHost("127.0.0.1")
+    depth = 2
+    executable_path = "/tmp/test-monkey"
+
+    cmd = utils.build_monkey_execution_command(host, depth, executable_path)
+
+    assert f"-d {depth}" in cmd
+    assert executable_path in cmd