From e6399de860a79cd9ceffc23614ca364243909e12 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 1 Sep 2021 13:36:36 -0400 Subject: [PATCH] Agent: Move get_credentials() to credentials.py --- monkey/infection_monkey/exploit/powershell.py | 3 +- .../powershell_utils/credential_generators.py | 41 ------------------- .../exploit/powershell_utils/credentials.py | 40 +++++++++++++++++- ...tial_generators.py => test_credentials.py} | 3 +- 4 files changed, 41 insertions(+), 46 deletions(-) delete mode 100644 monkey/infection_monkey/exploit/powershell_utils/credential_generators.py rename monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/{test_credential_generators.py => test_credentials.py} (93%) diff --git a/monkey/infection_monkey/exploit/powershell.py b/monkey/infection_monkey/exploit/powershell.py index 1765eaa6e..6361165c1 100644 --- a/monkey/infection_monkey/exploit/powershell.py +++ b/monkey/infection_monkey/exploit/powershell.py @@ -19,8 +19,7 @@ from infection_monkey.exploit.powershell_utils.auth_options import ( AuthOptions, get_auth_options, ) -from infection_monkey.exploit.powershell_utils.credential_generators import get_credentials -from infection_monkey.exploit.powershell_utils.credentials import Credentials +from infection_monkey.exploit.powershell_utils.credentials import Credentials, get_credentials from infection_monkey.exploit.powershell_utils.utils import ( IClient, get_client_based_on_auth_options, diff --git a/monkey/infection_monkey/exploit/powershell_utils/credential_generators.py b/monkey/infection_monkey/exploit/powershell_utils/credential_generators.py deleted file mode 100644 index 79840a800..000000000 --- a/monkey/infection_monkey/exploit/powershell_utils/credential_generators.py +++ /dev/null @@ -1,41 +0,0 @@ -from itertools import product -from typing import List - -from infection_monkey.exploit.powershell_utils.credentials import Credentials - - -def get_credentials( - usernames: List[str], passwords: List[str], is_windows: bool -) -> List[Credentials]: - 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[Credentials]: - if is_windows: - return [Credentials(username=None, password=None)] - - return [] - - -def _get_username_only_credentials(usernames: List[str], is_windows: bool) -> List[Credentials]: - credentials = [Credentials(username=username, password="") for username in usernames] - - if is_windows: - credentials.extend( - [Credentials(username=username, password=None) for username in usernames] - ) - - return credentials - - -def _get_username_password_credentials( - usernames: List[str], passwords: List[str] -) -> List[Credentials]: - username_password_pairs = product(usernames, passwords) - - return [Credentials(credentials[0], credentials[1]) for credentials in username_password_pairs] diff --git a/monkey/infection_monkey/exploit/powershell_utils/credentials.py b/monkey/infection_monkey/exploit/powershell_utils/credentials.py index 1a11b3f18..3d09498d7 100644 --- a/monkey/infection_monkey/exploit/powershell_utils/credentials.py +++ b/monkey/infection_monkey/exploit/powershell_utils/credentials.py @@ -1,8 +1,46 @@ from dataclasses import dataclass -from typing import Union +from itertools import product +from typing import List, Union @dataclass class Credentials: username: Union[str, None] password: Union[str, None] + + +def get_credentials( + usernames: List[str], passwords: List[str], is_windows: bool +) -> List[Credentials]: + 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[Credentials]: + if is_windows: + return [Credentials(username=None, password=None)] + + return [] + + +def _get_username_only_credentials(usernames: List[str], is_windows: bool) -> List[Credentials]: + credentials = [Credentials(username=username, password="") for username in usernames] + + if is_windows: + credentials.extend( + [Credentials(username=username, password=None) for username in usernames] + ) + + return credentials + + +def _get_username_password_credentials( + usernames: List[str], passwords: List[str] +) -> List[Credentials]: + username_password_pairs = product(usernames, passwords) + + return [Credentials(credentials[0], credentials[1]) for credentials in username_password_pairs] diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/test_credential_generators.py b/monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/test_credentials.py similarity index 93% rename from monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/test_credential_generators.py rename to monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/test_credentials.py index 7c41827ab..f1913169c 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/test_credential_generators.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/test_credentials.py @@ -1,5 +1,4 @@ -from infection_monkey.exploit.powershell_utils.credential_generators import get_credentials -from infection_monkey.exploit.powershell_utils.credentials import Credentials +from infection_monkey.exploit.powershell_utils.credentials import Credentials, get_credentials TEST_USERNAMES = ["user1", "user2"] TEST_PASSWORDS = ["p1", "p2"]