Agent: Extract _get_credentials() into powershell_utils/utils.py

monkey/infection_monkey/exploit/powershell.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import pypsrp
 import spnego
@@ -12,6 +12,7 @@ import infection_monkey.monkeyfs as monkeyfs
 from common.utils.exploit_enum import ExploitType
 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
@@ -48,7 +49,11 @@ class PowerShellExploiter(HostExploiter):
         return self._execute_monkey_agent_on_victim()
 
     def _authenticate_via_brute_force(self) -> Optional[Client]:
-        for username, password in self._get_credentials():
+        credentials = utils.get_credentials(
+            self._config.exploit_user_list, self._config.exploit_password_list, is_windows_os()
+        )
+
+        for username, password in credentials:
             try:
                 client = self._authenticate(username, password)
 
@@ -68,35 +73,6 @@ class PowerShellExploiter(HostExploiter):
 
         return None
 
-    def _get_credentials(self) -> List[Tuple[Optional[str], Optional[str]]]:
-        # When username or password is None, this instructs the powershell client to attempt to use
-        # The current user's credentials. This is only valid if the client is running from a Windows
-        # machine.
-
-        credentials = []
-        credentials.extend(self._get_empty_credentials())
-        credentials.extend(self._get_username_only_credentials())
-        credentials.extend(self._get_username_password_credentials())
-
-        return credentials
-
-    def _get_empty_credentials(self) -> List[Tuple[None, None]]:
-        if is_windows_os():
-            return [(None, None)]
-
-        return []
-
-    def _get_username_only_credentials(self) -> List[Tuple[str, Optional[str]]]:
-        credentials = [(username, "") for username in self._config.exploit_user_list]
-
-        if is_windows_os():
-            credentials.extend([(username, None) for username in self._config.exploit_user_list])
-
-        return credentials
-
-    def _get_username_password_credentials(self) -> List[Tuple[str, str]]:
-        return [credentials for credentials in self._config.get_exploit_user_password_pairs()]
-
     def _authenticate(self, username: Optional[str], password: Optional[str]) -> Client:
         ssl = password != ""
         auth = "negotiate" if password != "" else "basic"

monkey/infection_monkey/exploit/powershell_utils/utils.py

@@ -0,0 +1,43 @@
+from itertools import product
+from typing import List, Optional, Tuple
+
+
+def get_credentials(
+    usernames: List[str], passwords: List[str], is_windows: bool
+) -> List[Tuple[Optional[str], Optional[str]]]:
+    # When username or password is None, this instructs the powershell client to attempt to use
+    # The current user's credentials. This is only valid if the client is running from a Windows
+    # machine.
+
+    credentials = []
+    credentials.extend(_get_empty_credentials(is_windows))
+    credentials.extend(_get_username_only_credentials(usernames, is_windows))
+    credentials.extend(_get_username_password_credentials(usernames, passwords))
+
+    return credentials
+
+
+def _get_empty_credentials(is_windows: bool) -> List[Tuple[None, None]]:
+    if is_windows:
+        return [(None, None)]
+
+    return []
+
+
+def _get_username_only_credentials(
+    usernames: List[str], is_windows: bool
+) -> List[Tuple[str, Optional[str]]]:
+    credentials = [(username, "") for username in usernames]
+
+    if is_windows:
+        credentials.extend([(username, None) for username in usernames])
+
+    return credentials
+
+
+def _get_username_password_credentials(
+    usernames: List[str], passwords: List[str]
+) -> List[Tuple[str, str]]:
+    username_password_pairs = product(usernames, passwords)
+
+    return [credentials for credentials in username_password_pairs]

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

@@ -0,0 +1,44 @@
+from infection_monkey.exploit.powershell_utils import utils
+
+TEST_USERS = ["user1", "user2"]
+TEST_PASSWORDS = ["p1", "p2"]
+
+
+def test_get_credentials__empty_windows_true():
+    credentials = utils.get_credentials([], [], True)
+
+    assert len(credentials) == 1
+    assert credentials[0] == (None, None)
+
+
+def test_get_credentials__empty_windows_false():
+    credentials = utils.get_credentials([], [], False)
+
+    assert len(credentials) == 0
+
+
+def test_get_credentials__username_only_windows_false():
+    credentials = utils.get_credentials(TEST_USERS, [], False)
+
+    assert len(credentials) == 2
+    assert (TEST_USERS[0], "") in credentials
+    assert (TEST_USERS[1], "") in credentials
+
+
+def test_get_credentials__username_only_windows_true():
+    credentials = utils.get_credentials(TEST_USERS, [], True)
+
+    assert len(credentials) == 5
+    assert (TEST_USERS[0], "") in credentials
+    assert (TEST_USERS[1], "") in credentials
+    assert (TEST_USERS[0], None) in credentials
+    assert (TEST_USERS[1], None) in credentials
+
+
+def test_get_credentials__username_password():
+    credentials = utils.get_credentials(TEST_USERS, TEST_PASSWORDS, True)
+
+    assert len(credentials) == 9
+    for user in TEST_USERS:
+        for password in TEST_PASSWORDS:
+            assert (user, password) in credentials