SNMP: Add class to generate community strings

monkey/agent_plugins/exploiters/snmp/src/community_string_generator.py

@@ -0,0 +1,62 @@
+from itertools import chain
+from typing import Callable, Iterable, List
+
+from common.credentials import (
+    Credentials,
+    CredentialsComponent,
+    Identity,
+    Password,
+    Secret,
+    Username,
+)
+from infection_monkey.exploit.tools import identity_type_filter, secret_type_filter
+
+
+def generate_community_strings(credentials: Iterable[Credentials]) -> Iterable[str]:
+    """
+    Yields community strings from credentials
+
+    :param credentials: The credentials from which to generate community strings
+    :return: An iterable of potential community strings
+    """
+    for credential in _select_credentials_components(
+        credentials,
+        identity_type_filter=identity_type_filter([Username]),
+        secret_type_filter=secret_type_filter([Password]),
+    ):
+        yield _credentials_component_to_community_string(credential)
+
+
+def _select_credentials_components(
+    input_credentials: Iterable[Credentials],
+    identity_type_filter: Callable[[Identity], bool] = lambda identity: True,
+    secret_type_filter: Callable[[Identity], bool] = lambda secret: True,
+) -> Iterable[CredentialsComponent]:
+    identities: List[Identity] = []
+    secrets: List[Secret] = []
+
+    for credentials in input_credentials:
+        if credentials.identity is not None:
+            identities.append(credentials.identity)
+        if credentials.secret is not None:
+            secrets.append(credentials.secret)
+
+    for credential in chain(
+        filter(identity_type_filter, _deduplicate(identities)),
+        filter(secret_type_filter, _deduplicate(secrets)),
+    ):
+        yield credential
+
+
+def _deduplicate(iterable: Iterable[CredentialsComponent]) -> Iterable[CredentialsComponent]:
+    # Using dict.fromkeys() instead of set() because dicts preserve order
+    return dict.fromkeys(iterable).keys()
+
+
+def _credentials_component_to_community_string(credential: CredentialsComponent) -> str:
+    if isinstance(credential, Username):
+        return credential.username
+    elif isinstance(credential, Password):
+        return credential.password.get_secret_value()
+    else:
+        raise TypeError(f"Unexpected credential type: {type(credential)}")

monkey/common/credentials/__init__.py

@@ -5,4 +5,4 @@
 from .username import Username
 from .encoding import get_plaintext, SecretEncodingConfig
 
-from .credentials import Credentials, Identity, Secret
+from .credentials import Credentials, CredentialsComponent, Identity, Secret

monkey/common/credentials/credentials.py

@@ -6,8 +6,9 @@
 from . import LMHash, NTHash, Password, SSHKeypair, Username
 from .encoding import SecretEncodingConfig
 
-Secret = Union[Password, LMHash, NTHash, SSHKeypair]
 Identity = Username
+Secret = Union[Password, LMHash, NTHash, SSHKeypair]
+CredentialsComponent = Union[Identity, Secret]
 
 
 class Credentials(InfectionMonkeyBaseModel):

monkey/tests/unit_tests/agent_plugins/exploiters/snmp/test_community_string_generator.py

@@ -0,0 +1,100 @@
+from agent_plugins.exploiters.snmp.src.community_string_generator import generate_community_strings
+
+from common.credentials import Credentials, LMHash, Password, SSHKeypair, Username
+
+USERNAME_A_STR = "a"
+USERNAME_B_STR = "b"
+USERNAME_C_STR = "c"
+USERNAME_A = Username(username=USERNAME_A_STR)
+USERNAME_B = Username(username=USERNAME_B_STR)
+USERNAME_C = Username(username=USERNAME_C_STR)
+LMHASH_A_STR = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+LMHASH_A = LMHash(lm_hash=LMHASH_A_STR)
+NTHASH_A_STR = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
+NTHASH_A = LMHash(lm_hash=NTHASH_A_STR)
+PASSWORD_A_STR = "a"
+PASSWORD_B_STR = "b"
+PASSWORD_C_STR = "c"
+PASSWORD_A = Password(password=PASSWORD_A_STR)
+PASSWORD_B = Password(password=PASSWORD_B_STR)
+PASSWORD_C = Password(password=PASSWORD_C_STR)
+SSHKEYPAIR_A = SSHKeypair(private_key="a", public_key="b")
+
+
+def test_ordering():
+    # All identities should come before all secrets
+    # All identities should be in the same order as the input credentials
+    # All secrets should be in the same order as the input credentials
+    credentials = [
+        Credentials(identity=USERNAME_A, secret=PASSWORD_A),
+        Credentials(identity=USERNAME_B, secret=None),
+        Credentials(identity=None, secret=PASSWORD_B),
+        Credentials(identity=USERNAME_C, secret=PASSWORD_C),
+    ]
+    expected_result = [
+        USERNAME_A_STR,
+        USERNAME_B_STR,
+        USERNAME_C_STR,
+        PASSWORD_A_STR,
+        PASSWORD_B_STR,
+        PASSWORD_C_STR,
+    ]
+
+    assert list(generate_community_strings(credentials)) == expected_result
+
+
+def test_identities_are_filtered():
+    # The only identity type allowed is Username
+    credentials = [
+        Credentials(identity=USERNAME_A, secret=PASSWORD_A),
+        Credentials(identity=None, secret=PASSWORD_B),
+    ]
+    expected_result = [USERNAME_A_STR, PASSWORD_A_STR, PASSWORD_B_STR]
+
+    result = list(generate_community_strings(credentials))
+
+    assert result == expected_result
+
+
+def test_secrets_are_filtered():
+    # The only secret type allowed is Password
+    credentials = [
+        Credentials(identity=None, secret=PASSWORD_A),
+        Credentials(identity=USERNAME_A, secret=LMHASH_A),
+        Credentials(identity=USERNAME_B, secret=PASSWORD_B),
+        Credentials(identity=None, secret=NTHASH_A),
+        Credentials(identity=None, secret=SSHKEYPAIR_A),
+    ]
+    expected_result = [USERNAME_A_STR, USERNAME_B_STR, PASSWORD_A_STR, PASSWORD_B_STR]
+
+    result = list(generate_community_strings(credentials))
+
+    assert result == expected_result
+
+
+def test_identities_are_unique():
+    credentials = [
+        Credentials(identity=USERNAME_A, secret=None),
+        Credentials(identity=USERNAME_A, secret=PASSWORD_A),
+        Credentials(identity=USERNAME_B, secret=None),
+        Credentials(identity=USERNAME_B, secret=None),
+    ]
+    expected_result = [USERNAME_A_STR, USERNAME_B_STR, PASSWORD_A_STR]
+
+    result = list(generate_community_strings(credentials))
+
+    assert result == expected_result
+
+
+def test_secrets_are_unique():
+    credentials = [
+        Credentials(identity=None, secret=PASSWORD_A),
+        Credentials(identity=USERNAME_A, secret=PASSWORD_A),
+        Credentials(identity=None, secret=PASSWORD_B),
+        Credentials(identity=None, secret=PASSWORD_B),
+    ]
+    expected_result = [USERNAME_A_STR, PASSWORD_A_STR, PASSWORD_B_STR]
+
+    result = list(generate_community_strings(credentials))
+
+    assert result == expected_result