Agent: Refactor generate_brute_force_combinations

monkey/infection_monkey/exploit/wmiexec.py

@@ -12,7 +12,8 @@ from infection_monkey.exploit.tools.wmi_tools import AccessDeniedException, WmiT
 from infection_monkey.i_puppet import ExploiterResultData
 from infection_monkey.model import DROPPER_CMDLINE_WINDOWS, MONKEY_CMDLINE_WINDOWS
 from infection_monkey.utils.brute_force import (
-    get_credential_string, generate_brute_force_combinations,
+    generate_brute_force_combinations,
+    get_credential_string,
 )
 from infection_monkey.utils.commands import build_monkey_commandline
 
@@ -28,7 +29,7 @@ class WmiExploiter(HostExploiter):
     @WmiTools.dcom_wrap
     def _exploit_host(self) -> ExploiterResultData:
 
-        creds = generate_brute_force_combinations(self.options)
+        creds = generate_brute_force_combinations(self.options["credentials"])
 
         for user, password, lm_hash, ntlm_hash in creds:
             creds_for_log = get_credential_string([user, password, lm_hash, ntlm_hash])
@@ -96,13 +97,13 @@ class WmiExploiter(HostExploiter):
                     "dropper_path": remote_full_path
                 } + build_monkey_commandline(
                     self.host,
-                    self.current_depth-1,
+                    self.current_depth - 1,
                     self.options["dropper_target_path_win_64"],
                 )
             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/utils/brute_force.py

@@ -1,5 +1,5 @@
 from itertools import chain, product
-from typing import Any, Iterable, List, Tuple
+from typing import Any, Iterable, List, Mapping, Sequence, Tuple
 
 
 def generate_identity_secret_pairs(
@@ -40,12 +40,12 @@ def generate_username_password_or_ntlm_hash_combinations(
     )
 
 
-def generate_brute_force_combinations(options: dict):
+def generate_brute_force_combinations(credentials: Mapping[str, Sequence[str]]):
     return generate_username_password_or_ntlm_hash_combinations(
-        usernames=options["credentials"]["exploit_user_list"],
-        passwords=options["credentials"]["exploit_password_list"],
-        lm_hashes=options["credentials"]["exploit_lm_hash_list"],
-        nt_hashes=options["credentials"]["exploit_ntlm_hash_list"],
+        usernames=credentials["exploit_user_list"],
+        passwords=credentials["exploit_password_list"],
+        lm_hashes=credentials["exploit_lm_hash_list"],
+        nt_hashes=credentials["exploit_ntlm_hash_list"],
     )