From 3f3494e5d401d950e5936dd6f566eb038b0dff6e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 10:01:26 -0400 Subject: [PATCH 01/32] Common: Move DIContainer._del_key() to code_utils --- monkey/common/di_container.py | 28 ++++++------------- monkey/common/utils/code_utils.py | 18 +++++++++++- .../common/utils/test_code_utils.py | 22 ++++++++++++++- 3 files changed, 46 insertions(+), 22 deletions(-) diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py index 0a9748f4b..c1c740e39 100644 --- a/monkey/common/di_container.py +++ b/monkey/common/di_container.py @@ -1,5 +1,7 @@ import inspect -from typing import Any, MutableMapping, Sequence, Type, TypeVar +from typing import Any, Sequence, Type, TypeVar + +from common.utils.code_utils import del_key T = TypeVar("T") @@ -43,7 +45,7 @@ class DIContainer: ) self._type_registry[interface] = concrete_type - DIContainer._del_key(self._instance_registry, interface) + del_key(self._instance_registry, interface) def register_instance(self, interface: Type[T], instance: T): """ @@ -59,7 +61,7 @@ class DIContainer: ) self._instance_registry[interface] = instance - DIContainer._del_key(self._type_registry, interface) + del_key(self._type_registry, interface) def register_convention(self, type_: Type[T], name: str, instance: T): """ @@ -168,8 +170,8 @@ class DIContainer: :param interface: The interface to release """ - DIContainer._del_key(self._type_registry, interface) - DIContainer._del_key(self._instance_registry, interface) + del_key(self._type_registry, interface) + del_key(self._instance_registry, interface) def release_convention(self, type_: Type[T], name: str): """ @@ -179,18 +181,4 @@ class DIContainer: :param name: The name of the dependency parameter """ convention_identifier = (type_, name) - DIContainer._del_key(self._convention_registry, convention_identifier) - - @staticmethod - def _del_key(mapping: MutableMapping[T, Any], key: T): - """ - Deletes key from mapping. Unlike the `del` keyword, this function does not raise a KeyError - if the key does not exist. - - :param mapping: A mapping from which a key will be deleted - :param key: A key to delete from `mapping` - """ - try: - del mapping[key] - except KeyError: - pass + del_key(self._convention_registry, convention_identifier) diff --git a/monkey/common/utils/code_utils.py b/monkey/common/utils/code_utils.py index 251ce9375..a0fe6b5a9 100644 --- a/monkey/common/utils/code_utils.py +++ b/monkey/common/utils/code_utils.py @@ -1,5 +1,7 @@ import queue -from typing import Any, List +from typing import Any, List, MutableMapping, TypeVar + +T = TypeVar("T") class abstractstatic(staticmethod): @@ -30,3 +32,17 @@ def queue_to_list(q: queue.Queue) -> List[Any]: pass return list_ + + +def del_key(mapping: MutableMapping[T, Any], key: T): + """ + Deletes key from mapping. Unlike the `del` keyword, this function does not raise a KeyError + if the key does not exist. + + :param mapping: A mapping from which a key will be deleted + :param key: A key to delete from `mapping` + """ + try: + del mapping[key] + except KeyError: + pass diff --git a/monkey/tests/unit_tests/common/utils/test_code_utils.py b/monkey/tests/unit_tests/common/utils/test_code_utils.py index 411b07a63..e5980723d 100644 --- a/monkey/tests/unit_tests/common/utils/test_code_utils.py +++ b/monkey/tests/unit_tests/common/utils/test_code_utils.py @@ -1,6 +1,6 @@ from queue import Queue -from common.utils.code_utils import queue_to_list +from common.utils.code_utils import del_key, queue_to_list def test_empty_queue_to_empty_list(): @@ -20,3 +20,23 @@ def test_queue_to_list(): list_ = queue_to_list(q) assert list_ == expected_list + + +def test_del_key__deletes_key(): + key_to_delete = "a" + my_dict = {"a": 1, "b": 2} + expected_dict = {k: v for k, v in my_dict.items() if k != key_to_delete} + + del_key(my_dict, key_to_delete) + + assert my_dict == expected_dict + + +def test_del_key__nonexistant_key(): + key_to_delete = "a" + my_dict = {"a": 1, "b": 2} + + del_key(my_dict, key_to_delete) + + # This test passes if the following call does not raise an error + del_key(my_dict, key_to_delete) From 0be43157cfbb4a2bd941665746109147b2c11ed1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 10:46:05 -0400 Subject: [PATCH 02/32] Common: Add PasswordSchema --- monkey/common/credentials/password.py | 17 +++++++ .../common/credentials/test_password.py | 49 +++++++++++++++++++ vulture_allowlist.py | 3 ++ 3 files changed, 69 insertions(+) create mode 100644 monkey/tests/unit_tests/common/credentials/test_password.py diff --git a/monkey/common/credentials/password.py b/monkey/common/credentials/password.py index ff6da0eae..1fee478aa 100644 --- a/monkey/common/credentials/password.py +++ b/monkey/common/credentials/password.py @@ -1,8 +1,25 @@ from dataclasses import dataclass, field +from marshmallow import Schema, fields, post_load, validate +from marshmallow_enum import EnumField + +from common.utils.code_utils import del_key + from . import CredentialComponentType, ICredentialComponent +class PasswordSchema(Schema): + credential_type = EnumField( + CredentialComponentType, validate=validate.Equal(CredentialComponentType.PASSWORD) + ) + password = fields.Str() + + @post_load + def _strip_credential_type(self, data, **kwargs): + del_key(data, "credential_type") + return data + + @dataclass(frozen=True) class Password(ICredentialComponent): credential_type: CredentialComponentType = field( diff --git a/monkey/tests/unit_tests/common/credentials/test_password.py b/monkey/tests/unit_tests/common/credentials/test_password.py new file mode 100644 index 000000000..deb9ca594 --- /dev/null +++ b/monkey/tests/unit_tests/common/credentials/test_password.py @@ -0,0 +1,49 @@ +from copy import deepcopy + +import pytest +from marshmallow.exceptions import ValidationError + +from common.credentials import CredentialComponentType, Password +from common.credentials.password import PasswordSchema + +PASSWORD_VALUE = "123456" +PASSWORD_DICT = { + "credential_type": CredentialComponentType.PASSWORD.name, + "password": PASSWORD_VALUE, +} + + +def test_password_serialize(): + schema = PasswordSchema() + password = Password(PASSWORD_VALUE) + + serialized_password = schema.dump(password) + + assert serialized_password == PASSWORD_DICT + + +def test_password_deserialize(): + schema = PasswordSchema() + + password = Password(**schema.load(PASSWORD_DICT)) + + assert password.credential_type == CredentialComponentType.PASSWORD + assert password.password == PASSWORD_VALUE + + +def test_invalid_credential_type(): + invalid_password_dict = deepcopy(PASSWORD_DICT) + invalid_password_dict["credential_type"] = "INVALID" + schema = PasswordSchema() + + with pytest.raises(ValidationError): + Password(**schema.load(invalid_password_dict)) + + +def test_incorrect_credential_type(): + invalid_password_dict = deepcopy(PASSWORD_DICT) + invalid_password_dict["credential_type"] = CredentialComponentType.USERNAME.name + schema = PasswordSchema() + + with pytest.raises(ValidationError): + Password(**schema.load(invalid_password_dict)) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index a3393325e..b04ec8cca 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -196,6 +196,9 @@ _make_tcp_scan_configuration # unused method (monkey/common/configuration/agent _make_network_scan_configuration # unused method (monkey/common/configuration/agent_configuration.py:110) _make_propagation_configuration # unused method (monkey/common/configuration/agent_configuration.py:167) +# Credentials +_strip_credential_type # unused method(monkey/common/credentials/password.py:18) + # Models _make_simulation # unused method (monkey/monkey_island/cc/models/simulation.py:19 From 0b887a2704488d2bff710fb0d2134724b9994ca9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 10:48:44 -0400 Subject: [PATCH 03/32] UT: Add unit_tests/common/credentials/__init__.py --- monkey/tests/unit_tests/common/credentials/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 monkey/tests/unit_tests/common/credentials/__init__.py diff --git a/monkey/tests/unit_tests/common/credentials/__init__.py b/monkey/tests/unit_tests/common/credentials/__init__.py new file mode 100644 index 000000000..e69de29bb From 037b4ef8c50285d8fb65ac115da1e0975a20ab02 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 10:52:13 -0400 Subject: [PATCH 04/32] Common: Add UsernameSchema --- monkey/common/credentials/username.py | 17 +++++++ .../common/credentials/test_username.py | 49 +++++++++++++++++++ 2 files changed, 66 insertions(+) create mode 100644 monkey/tests/unit_tests/common/credentials/test_username.py diff --git a/monkey/common/credentials/username.py b/monkey/common/credentials/username.py index c3249058e..837ba2b98 100644 --- a/monkey/common/credentials/username.py +++ b/monkey/common/credentials/username.py @@ -1,8 +1,25 @@ from dataclasses import dataclass, field +from marshmallow import Schema, fields, post_load, validate +from marshmallow_enum import EnumField + +from common.utils.code_utils import del_key + from . import CredentialComponentType, ICredentialComponent +class UsernameSchema(Schema): + credential_type = EnumField( + CredentialComponentType, validate=validate.Equal(CredentialComponentType.USERNAME) + ) + username = fields.Str() + + @post_load + def _strip_credential_type(self, data, **kwargs): + del_key(data, "credential_type") + return data + + @dataclass(frozen=True) class Username(ICredentialComponent): credential_type: CredentialComponentType = field( diff --git a/monkey/tests/unit_tests/common/credentials/test_username.py b/monkey/tests/unit_tests/common/credentials/test_username.py new file mode 100644 index 000000000..3bee18978 --- /dev/null +++ b/monkey/tests/unit_tests/common/credentials/test_username.py @@ -0,0 +1,49 @@ +from copy import deepcopy + +import pytest +from marshmallow.exceptions import ValidationError + +from common.credentials import CredentialComponentType, Username +from common.credentials.username import UsernameSchema + +USERNAME_VALUE = "test_user" +USERNAME_DICT = { + "credential_type": CredentialComponentType.USERNAME.name, + "username": USERNAME_VALUE, +} + + +def test_username_serialize(): + schema = UsernameSchema() + username = Username(USERNAME_VALUE) + + serialized_username = schema.dump(username) + + assert serialized_username == USERNAME_DICT + + +def test_username_deserialize(): + schema = UsernameSchema() + + username = Username(**schema.load(USERNAME_DICT)) + + assert username.credential_type == CredentialComponentType.USERNAME + assert username.username == USERNAME_VALUE + + +def test_invalid_credential_type(): + invalid_username_dict = deepcopy(USERNAME_DICT) + invalid_username_dict["credential_type"] = "INVALID" + schema = UsernameSchema() + + with pytest.raises(ValidationError): + Username(**schema.load(invalid_username_dict)) + + +def test_incorrect_credential_type(): + invalid_username_dict = deepcopy(USERNAME_DICT) + invalid_username_dict["credential_type"] = CredentialComponentType.PASSWORD.name + schema = UsernameSchema() + + with pytest.raises(ValidationError): + Username(**schema.load(invalid_username_dict)) From a8747c9d5d8859fcf5395f03783e901f84e8c1ff Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 10:58:53 -0400 Subject: [PATCH 05/32] Common: Add CredentialComponentSchema Extract common _strip_credential_type() post_load function from PasswordSchema and UsernameSchema into a parent class. --- .../credentials/credential_component_schema.py | 10 ++++++++++ monkey/common/credentials/password.py | 12 +++--------- monkey/common/credentials/username.py | 12 +++--------- 3 files changed, 16 insertions(+), 18 deletions(-) create mode 100644 monkey/common/credentials/credential_component_schema.py diff --git a/monkey/common/credentials/credential_component_schema.py b/monkey/common/credentials/credential_component_schema.py new file mode 100644 index 000000000..97f8c10b0 --- /dev/null +++ b/monkey/common/credentials/credential_component_schema.py @@ -0,0 +1,10 @@ +from marshmallow import Schema, post_load + +from common.utils.code_utils import del_key + + +class CredentialComponentSchema(Schema): + @post_load + def _strip_credential_type(self, data, **kwargs): + del_key(data, "credential_type") + return data diff --git a/monkey/common/credentials/password.py b/monkey/common/credentials/password.py index 1fee478aa..7b4d2178f 100644 --- a/monkey/common/credentials/password.py +++ b/monkey/common/credentials/password.py @@ -1,24 +1,18 @@ from dataclasses import dataclass, field -from marshmallow import Schema, fields, post_load, validate +from marshmallow import fields, validate from marshmallow_enum import EnumField -from common.utils.code_utils import del_key - from . import CredentialComponentType, ICredentialComponent +from .credential_component_schema import CredentialComponentSchema -class PasswordSchema(Schema): +class PasswordSchema(CredentialComponentSchema): credential_type = EnumField( CredentialComponentType, validate=validate.Equal(CredentialComponentType.PASSWORD) ) password = fields.Str() - @post_load - def _strip_credential_type(self, data, **kwargs): - del_key(data, "credential_type") - return data - @dataclass(frozen=True) class Password(ICredentialComponent): diff --git a/monkey/common/credentials/username.py b/monkey/common/credentials/username.py index 837ba2b98..ffa863e15 100644 --- a/monkey/common/credentials/username.py +++ b/monkey/common/credentials/username.py @@ -1,24 +1,18 @@ from dataclasses import dataclass, field -from marshmallow import Schema, fields, post_load, validate +from marshmallow import fields, validate from marshmallow_enum import EnumField -from common.utils.code_utils import del_key - from . import CredentialComponentType, ICredentialComponent +from .credential_component_schema import CredentialComponentSchema -class UsernameSchema(Schema): +class UsernameSchema(CredentialComponentSchema): credential_type = EnumField( CredentialComponentType, validate=validate.Equal(CredentialComponentType.USERNAME) ) username = fields.Str() - @post_load - def _strip_credential_type(self, data, **kwargs): - del_key(data, "credential_type") - return data - @dataclass(frozen=True) class Username(ICredentialComponent): From 9a45d777cafcdeb200b336becf01afa0560551be Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 11:07:46 -0400 Subject: [PATCH 06/32] Common: Add CredentialTypeField --- .../credentials/credential_component_schema.py | 12 +++++++++++- monkey/common/credentials/password.py | 9 +++------ monkey/common/credentials/username.py | 9 +++------ 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/monkey/common/credentials/credential_component_schema.py b/monkey/common/credentials/credential_component_schema.py index 97f8c10b0..ff3e657be 100644 --- a/monkey/common/credentials/credential_component_schema.py +++ b/monkey/common/credentials/credential_component_schema.py @@ -1,7 +1,17 @@ -from marshmallow import Schema, post_load +from marshmallow import Schema, post_load, validate +from marshmallow_enum import EnumField from common.utils.code_utils import del_key +from . import CredentialComponentType + + +class CredentialTypeField(EnumField): + def __init__(self, credential_component_type: CredentialComponentType): + super().__init__( + CredentialComponentType, validate=validate.Equal(credential_component_type) + ) + class CredentialComponentSchema(Schema): @post_load diff --git a/monkey/common/credentials/password.py b/monkey/common/credentials/password.py index 7b4d2178f..b7bd1b84c 100644 --- a/monkey/common/credentials/password.py +++ b/monkey/common/credentials/password.py @@ -1,16 +1,13 @@ from dataclasses import dataclass, field -from marshmallow import fields, validate -from marshmallow_enum import EnumField +from marshmallow import fields from . import CredentialComponentType, ICredentialComponent -from .credential_component_schema import CredentialComponentSchema +from .credential_component_schema import CredentialComponentSchema, CredentialTypeField class PasswordSchema(CredentialComponentSchema): - credential_type = EnumField( - CredentialComponentType, validate=validate.Equal(CredentialComponentType.PASSWORD) - ) + credential_type = CredentialTypeField(CredentialComponentType.PASSWORD) password = fields.Str() diff --git a/monkey/common/credentials/username.py b/monkey/common/credentials/username.py index ffa863e15..86fde05ff 100644 --- a/monkey/common/credentials/username.py +++ b/monkey/common/credentials/username.py @@ -1,16 +1,13 @@ from dataclasses import dataclass, field -from marshmallow import fields, validate -from marshmallow_enum import EnumField +from marshmallow import fields from . import CredentialComponentType, ICredentialComponent -from .credential_component_schema import CredentialComponentSchema +from .credential_component_schema import CredentialComponentSchema, CredentialTypeField class UsernameSchema(CredentialComponentSchema): - credential_type = EnumField( - CredentialComponentType, validate=validate.Equal(CredentialComponentType.USERNAME) - ) + credential_type = CredentialTypeField(CredentialComponentType.USERNAME) username = fields.Str() From be9889c9d1fcd4fca678bdfd86ac7ac50ca05603 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 11:31:24 -0400 Subject: [PATCH 07/32] UT: Use parametrization to remove duplicate credential component tests --- .../common/credentials/test_password.py | 49 ------------ ...test_single_value_credential_components.py | 74 +++++++++++++++++++ .../common/credentials/test_username.py | 49 ------------ 3 files changed, 74 insertions(+), 98 deletions(-) delete mode 100644 monkey/tests/unit_tests/common/credentials/test_password.py create mode 100644 monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py delete mode 100644 monkey/tests/unit_tests/common/credentials/test_username.py diff --git a/monkey/tests/unit_tests/common/credentials/test_password.py b/monkey/tests/unit_tests/common/credentials/test_password.py deleted file mode 100644 index deb9ca594..000000000 --- a/monkey/tests/unit_tests/common/credentials/test_password.py +++ /dev/null @@ -1,49 +0,0 @@ -from copy import deepcopy - -import pytest -from marshmallow.exceptions import ValidationError - -from common.credentials import CredentialComponentType, Password -from common.credentials.password import PasswordSchema - -PASSWORD_VALUE = "123456" -PASSWORD_DICT = { - "credential_type": CredentialComponentType.PASSWORD.name, - "password": PASSWORD_VALUE, -} - - -def test_password_serialize(): - schema = PasswordSchema() - password = Password(PASSWORD_VALUE) - - serialized_password = schema.dump(password) - - assert serialized_password == PASSWORD_DICT - - -def test_password_deserialize(): - schema = PasswordSchema() - - password = Password(**schema.load(PASSWORD_DICT)) - - assert password.credential_type == CredentialComponentType.PASSWORD - assert password.password == PASSWORD_VALUE - - -def test_invalid_credential_type(): - invalid_password_dict = deepcopy(PASSWORD_DICT) - invalid_password_dict["credential_type"] = "INVALID" - schema = PasswordSchema() - - with pytest.raises(ValidationError): - Password(**schema.load(invalid_password_dict)) - - -def test_incorrect_credential_type(): - invalid_password_dict = deepcopy(PASSWORD_DICT) - invalid_password_dict["credential_type"] = CredentialComponentType.USERNAME.name - schema = PasswordSchema() - - with pytest.raises(ValidationError): - Password(**schema.load(invalid_password_dict)) diff --git a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py b/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py new file mode 100644 index 000000000..2c138f9d6 --- /dev/null +++ b/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py @@ -0,0 +1,74 @@ +import pytest +from marshmallow.exceptions import ValidationError + +from common.credentials import CredentialComponentType, Password, Username +from common.credentials.password import PasswordSchema +from common.credentials.username import UsernameSchema + +PARAMETRIZED_PARAMETER_NAMES = ( + "credential_component_class, schema_class, credential_component_type, key, value" +) + +PARAMETRIZED_PARAMETER_VALUES = [ + (Password, PasswordSchema, CredentialComponentType.PASSWORD, "password", "123456"), + (Username, UsernameSchema, CredentialComponentType.USERNAME, "username", "test_user"), +] + + +def build_credential_dict(credential_component_type: CredentialComponentType, key: str, value: str): + return {"credential_type": credential_component_type.name, key: value} + + +@pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) +def test_credential_component_serialize( + credential_component_class, schema_class, credential_component_type, key, value +): + schema = schema_class() + print(type(schema)) + constructed_object = credential_component_class(value) + print(type(constructed_object)) + + serialized_object = schema.dump(constructed_object) + + assert serialized_object == build_credential_dict(credential_component_type, key, value) + + +@pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) +def test_credential_component_deserialize( + credential_component_class, schema_class, credential_component_type, key, value +): + schema = schema_class() + credential_dict = build_credential_dict(credential_component_type, key, value) + expected_deserialized_object = credential_component_class(value) + + deserialized_object = credential_component_class(**schema.load(credential_dict)) + + assert deserialized_object == expected_deserialized_object + + +@pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) +def test_invalid_credential_type( + credential_component_class, schema_class, credential_component_type, key, value +): + invalid_component_dict = build_credential_dict(credential_component_type, key, value) + invalid_component_dict["credential_type"] = "INVALID" + schema = schema_class() + + with pytest.raises(ValidationError): + credential_component_class(**schema.load(invalid_component_dict)) + + +@pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) +def test_encorrect_credential_type( + credential_component_class, schema_class, credential_component_type, key, value +): + incorrect_component_dict = build_credential_dict(credential_component_type, key, value) + incorrect_component_dict["credential_type"] = ( + CredentialComponentType.USERNAME.name + if credential_component_type != CredentialComponentType.USERNAME + else CredentialComponentType.PASSWORD + ) + schema = schema_class() + + with pytest.raises(ValidationError): + credential_component_class(**schema.load(incorrect_component_dict)) diff --git a/monkey/tests/unit_tests/common/credentials/test_username.py b/monkey/tests/unit_tests/common/credentials/test_username.py deleted file mode 100644 index 3bee18978..000000000 --- a/monkey/tests/unit_tests/common/credentials/test_username.py +++ /dev/null @@ -1,49 +0,0 @@ -from copy import deepcopy - -import pytest -from marshmallow.exceptions import ValidationError - -from common.credentials import CredentialComponentType, Username -from common.credentials.username import UsernameSchema - -USERNAME_VALUE = "test_user" -USERNAME_DICT = { - "credential_type": CredentialComponentType.USERNAME.name, - "username": USERNAME_VALUE, -} - - -def test_username_serialize(): - schema = UsernameSchema() - username = Username(USERNAME_VALUE) - - serialized_username = schema.dump(username) - - assert serialized_username == USERNAME_DICT - - -def test_username_deserialize(): - schema = UsernameSchema() - - username = Username(**schema.load(USERNAME_DICT)) - - assert username.credential_type == CredentialComponentType.USERNAME - assert username.username == USERNAME_VALUE - - -def test_invalid_credential_type(): - invalid_username_dict = deepcopy(USERNAME_DICT) - invalid_username_dict["credential_type"] = "INVALID" - schema = UsernameSchema() - - with pytest.raises(ValidationError): - Username(**schema.load(invalid_username_dict)) - - -def test_incorrect_credential_type(): - invalid_username_dict = deepcopy(USERNAME_DICT) - invalid_username_dict["credential_type"] = CredentialComponentType.PASSWORD.name - schema = UsernameSchema() - - with pytest.raises(ValidationError): - Username(**schema.load(invalid_username_dict)) From 031fce9fd8b661caebbce2c7d507536a94f23e42 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 11:50:54 -0400 Subject: [PATCH 08/32] UT: Add test_invalid_values() --- ...test_single_value_credential_components.py | 28 +++++++++++++++---- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py b/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py index 2c138f9d6..51d02ae04 100644 --- a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py +++ b/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py @@ -15,7 +15,14 @@ PARAMETRIZED_PARAMETER_VALUES = [ ] -def build_credential_dict(credential_component_type: CredentialComponentType, key: str, value: str): + + +INVALID_VALUES = { + CredentialComponentType.USERNAME: (None, 1, 2.0), + CredentialComponentType.PASSWORD: (None, 1, 2.0), +} + +def build_component_dict(credential_component_type: CredentialComponentType, key: str, value: str): return {"credential_type": credential_component_type.name, key: value} @@ -30,7 +37,7 @@ def test_credential_component_serialize( serialized_object = schema.dump(constructed_object) - assert serialized_object == build_credential_dict(credential_component_type, key, value) + assert serialized_object == build_component_dict(credential_component_type, key, value) @pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) @@ -38,7 +45,7 @@ def test_credential_component_deserialize( credential_component_class, schema_class, credential_component_type, key, value ): schema = schema_class() - credential_dict = build_credential_dict(credential_component_type, key, value) + credential_dict = build_component_dict(credential_component_type, key, value) expected_deserialized_object = credential_component_class(value) deserialized_object = credential_component_class(**schema.load(credential_dict)) @@ -50,7 +57,7 @@ def test_credential_component_deserialize( def test_invalid_credential_type( credential_component_class, schema_class, credential_component_type, key, value ): - invalid_component_dict = build_credential_dict(credential_component_type, key, value) + invalid_component_dict = build_component_dict(credential_component_type, key, value) invalid_component_dict["credential_type"] = "INVALID" schema = schema_class() @@ -62,7 +69,7 @@ def test_invalid_credential_type( def test_encorrect_credential_type( credential_component_class, schema_class, credential_component_type, key, value ): - incorrect_component_dict = build_credential_dict(credential_component_type, key, value) + incorrect_component_dict = build_component_dict(credential_component_type, key, value) incorrect_component_dict["credential_type"] = ( CredentialComponentType.USERNAME.name if credential_component_type != CredentialComponentType.USERNAME @@ -72,3 +79,14 @@ def test_encorrect_credential_type( with pytest.raises(ValidationError): credential_component_class(**schema.load(incorrect_component_dict)) + +@pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) +def test_invalid_values( + credential_component_class, schema_class, credential_component_type, key, value +): + schema = schema_class() + + for invalid_value in INVALID_VALUES[credential_component_type]: + component_dict = build_component_dict(credential_component_type, key, invalid_value) + with pytest.raises(ValidationError): + credential_component_class(**schema.load(component_dict)) From 68e52eb512ec97defecc777c014e67df06ee05b9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 12:02:25 -0400 Subject: [PATCH 09/32] Common: Add ntlm_hash_validator --- monkey/common/credentials/validators.py | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 monkey/common/credentials/validators.py diff --git a/monkey/common/credentials/validators.py b/monkey/common/credentials/validators.py new file mode 100644 index 000000000..3d1b74650 --- /dev/null +++ b/monkey/common/credentials/validators.py @@ -0,0 +1,6 @@ +import re + +from marshmallow import validate + +_ntlm_hash_regex = re.compile(r"^[a-fA-F0-9]{32}$") +ntlm_hash_validator = validate.Regexp(regex=_ntlm_hash_regex) From def2381da6eb9474e3bab9e618df8e0cabaf6e74 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 12:02:49 -0400 Subject: [PATCH 10/32] Common: Add LMHashSchema --- monkey/common/credentials/lm_hash.py | 9 ++++++++ ...test_single_value_credential_components.py | 23 +++++++++++++++---- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/monkey/common/credentials/lm_hash.py b/monkey/common/credentials/lm_hash.py index 818999244..4a1b59f80 100644 --- a/monkey/common/credentials/lm_hash.py +++ b/monkey/common/credentials/lm_hash.py @@ -1,6 +1,15 @@ from dataclasses import dataclass, field +from marshmallow import fields + from . import CredentialComponentType, ICredentialComponent +from .credential_component_schema import CredentialComponentSchema, CredentialTypeField +from .validators import ntlm_hash_validator + + +class LMHashSchema(CredentialComponentSchema): + credential_type = CredentialTypeField(CredentialComponentType.LM_HASH) + lm_hash = fields.Str(validate=ntlm_hash_validator) @dataclass(frozen=True) diff --git a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py b/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py index 51d02ae04..1b61dc329 100644 --- a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py +++ b/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py @@ -1,7 +1,8 @@ import pytest from marshmallow.exceptions import ValidationError -from common.credentials import CredentialComponentType, Password, Username +from common.credentials import CredentialComponentType, LMHash, Password, Username +from common.credentials.lm_hash import LMHashSchema from common.credentials.password import PasswordSchema from common.credentials.username import UsernameSchema @@ -12,16 +13,29 @@ PARAMETRIZED_PARAMETER_NAMES = ( PARAMETRIZED_PARAMETER_VALUES = [ (Password, PasswordSchema, CredentialComponentType.PASSWORD, "password", "123456"), (Username, UsernameSchema, CredentialComponentType.USERNAME, "username", "test_user"), + ( + LMHash, + LMHashSchema, + CredentialComponentType.LM_HASH, + "lm_hash", + "E52CAC67419A9A224A3B108F3FA6CB6D", + ), ] - - INVALID_VALUES = { CredentialComponentType.USERNAME: (None, 1, 2.0), CredentialComponentType.PASSWORD: (None, 1, 2.0), + CredentialComponentType.LM_HASH: ( + None, + 1, + 2.0, + "0123456789012345678901234568901", + "E52GAC67419A9A224A3B108F3FA6CB6D", + ), } + def build_component_dict(credential_component_type: CredentialComponentType, key: str, value: str): return {"credential_type": credential_component_type.name, key: value} @@ -31,9 +45,7 @@ def test_credential_component_serialize( credential_component_class, schema_class, credential_component_type, key, value ): schema = schema_class() - print(type(schema)) constructed_object = credential_component_class(value) - print(type(constructed_object)) serialized_object = schema.dump(constructed_object) @@ -80,6 +92,7 @@ def test_encorrect_credential_type( with pytest.raises(ValidationError): credential_component_class(**schema.load(incorrect_component_dict)) + @pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) def test_invalid_values( credential_component_class, schema_class, credential_component_type, key, value From 58fcc3761ce6011dde8cd6db7f865e50d7394679 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 12:07:44 -0400 Subject: [PATCH 11/32] UT: Add NTHashSchema --- monkey/common/credentials/nt_hash.py | 9 +++++++++ .../test_single_value_credential_components.py | 17 ++++++++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/monkey/common/credentials/nt_hash.py b/monkey/common/credentials/nt_hash.py index 67246f695..838ec4596 100644 --- a/monkey/common/credentials/nt_hash.py +++ b/monkey/common/credentials/nt_hash.py @@ -1,6 +1,15 @@ from dataclasses import dataclass, field +from marshmallow import fields + from . import CredentialComponentType, ICredentialComponent +from .credential_component_schema import CredentialComponentSchema, CredentialTypeField +from .validators import ntlm_hash_validator + + +class NTHashSchema(CredentialComponentSchema): + credential_type = CredentialTypeField(CredentialComponentType.NT_HASH) + nt_hash = fields.Str(validate=ntlm_hash_validator) @dataclass(frozen=True) diff --git a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py b/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py index 1b61dc329..0004b38ee 100644 --- a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py +++ b/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py @@ -1,8 +1,9 @@ import pytest from marshmallow.exceptions import ValidationError -from common.credentials import CredentialComponentType, LMHash, Password, Username +from common.credentials import CredentialComponentType, LMHash, NTHash, Password, Username from common.credentials.lm_hash import LMHashSchema +from common.credentials.nt_hash import NTHashSchema from common.credentials.password import PasswordSchema from common.credentials.username import UsernameSchema @@ -20,6 +21,13 @@ PARAMETRIZED_PARAMETER_VALUES = [ "lm_hash", "E52CAC67419A9A224A3B108F3FA6CB6D", ), + ( + NTHash, + NTHashSchema, + CredentialComponentType.NT_HASH, + "nt_hash", + "E52CAC67419A9A224A3B108F3FA6CB6D", + ), ] @@ -33,6 +41,13 @@ INVALID_VALUES = { "0123456789012345678901234568901", "E52GAC67419A9A224A3B108F3FA6CB6D", ), + CredentialComponentType.NT_HASH: ( + None, + 1, + 2.0, + "0123456789012345678901234568901", + "E52GAC67419A9A224A3B108F3FA6CB6D", + ), } From e92de42da998a77bf4bce4d51ab4785db9a5ed74 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 12:25:01 -0400 Subject: [PATCH 12/32] UT: Use dict for key,value in credential components tests The SSHKeypair credential component has two fields (public, private), not just a single value. This commit modifies the tests to be able to support credential components with multiple fields. --- ...nents.py => test_credential_components.py} | 78 ++++++++++--------- 1 file changed, 43 insertions(+), 35 deletions(-) rename monkey/tests/unit_tests/common/credentials/{test_single_value_credential_components.py => test_credential_components.py} (60%) diff --git a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py b/monkey/tests/unit_tests/common/credentials/test_credential_components.py similarity index 60% rename from monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py rename to monkey/tests/unit_tests/common/credentials/test_credential_components.py index 0004b38ee..6be331356 100644 --- a/monkey/tests/unit_tests/common/credentials/test_single_value_credential_components.py +++ b/monkey/tests/unit_tests/common/credentials/test_credential_components.py @@ -1,3 +1,5 @@ +from typing import Any, Mapping + import pytest from marshmallow.exceptions import ValidationError @@ -8,72 +10,74 @@ from common.credentials.password import PasswordSchema from common.credentials.username import UsernameSchema PARAMETRIZED_PARAMETER_NAMES = ( - "credential_component_class, schema_class, credential_component_type, key, value" + "credential_component_class, schema_class, credential_component_type, credential_component_data" ) PARAMETRIZED_PARAMETER_VALUES = [ - (Password, PasswordSchema, CredentialComponentType.PASSWORD, "password", "123456"), - (Username, UsernameSchema, CredentialComponentType.USERNAME, "username", "test_user"), + (Username, UsernameSchema, CredentialComponentType.USERNAME, {"username": "test_user"}), + (Password, PasswordSchema, CredentialComponentType.PASSWORD, {"password": "123456"}), ( LMHash, LMHashSchema, CredentialComponentType.LM_HASH, - "lm_hash", - "E52CAC67419A9A224A3B108F3FA6CB6D", + {"lm_hash": "E52CAC67419A9A224A3B108F3FA6CB6D"}, ), ( NTHash, NTHashSchema, CredentialComponentType.NT_HASH, - "nt_hash", - "E52CAC67419A9A224A3B108F3FA6CB6D", + {"nt_hash": "E52CAC67419A9A224A3B108F3FA6CB6D"}, ), ] -INVALID_VALUES = { - CredentialComponentType.USERNAME: (None, 1, 2.0), - CredentialComponentType.PASSWORD: (None, 1, 2.0), +INVALID_COMPONENT_DATA = { + CredentialComponentType.USERNAME: ({"username": None}, {"username": 1}, {"username": 2.0}), + CredentialComponentType.PASSWORD: ({"password": None}, {"password": 1}, {"password": 2.0}), CredentialComponentType.LM_HASH: ( - None, - 1, - 2.0, - "0123456789012345678901234568901", - "E52GAC67419A9A224A3B108F3FA6CB6D", + {"lm_hash": None}, + {"lm_hash": 1}, + {"lm_hash": 2.0}, + {"lm_hash": "0123456789012345678901234568901"}, + {"lm_hash": "E52GAC67419A9A224A3B108F3FA6CB6D"}, ), CredentialComponentType.NT_HASH: ( - None, - 1, - 2.0, - "0123456789012345678901234568901", - "E52GAC67419A9A224A3B108F3FA6CB6D", + {"nt_hash": None}, + {"nt_hash": 1}, + {"nt_hash": 2.0}, + {"nt_hash": "0123456789012345678901234568901"}, + {"nt_hash": "E52GAC67419A9A224A3B108F3FA6CB6D"}, ), } -def build_component_dict(credential_component_type: CredentialComponentType, key: str, value: str): - return {"credential_type": credential_component_type.name, key: value} +def build_component_dict( + credential_component_type: CredentialComponentType, credential_component_data: Mapping[str, Any] +): + return {"credential_type": credential_component_type.name, **credential_component_data} @pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) def test_credential_component_serialize( - credential_component_class, schema_class, credential_component_type, key, value + credential_component_class, schema_class, credential_component_type, credential_component_data ): schema = schema_class() - constructed_object = credential_component_class(value) + constructed_object = credential_component_class(**credential_component_data) serialized_object = schema.dump(constructed_object) - assert serialized_object == build_component_dict(credential_component_type, key, value) + assert serialized_object == build_component_dict( + credential_component_type, credential_component_data + ) @pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) def test_credential_component_deserialize( - credential_component_class, schema_class, credential_component_type, key, value + credential_component_class, schema_class, credential_component_type, credential_component_data ): schema = schema_class() - credential_dict = build_component_dict(credential_component_type, key, value) - expected_deserialized_object = credential_component_class(value) + credential_dict = build_component_dict(credential_component_type, credential_component_data) + expected_deserialized_object = credential_component_class(**credential_component_data) deserialized_object = credential_component_class(**schema.load(credential_dict)) @@ -82,9 +86,11 @@ def test_credential_component_deserialize( @pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) def test_invalid_credential_type( - credential_component_class, schema_class, credential_component_type, key, value + credential_component_class, schema_class, credential_component_type, credential_component_data ): - invalid_component_dict = build_component_dict(credential_component_type, key, value) + invalid_component_dict = build_component_dict( + credential_component_type, credential_component_data + ) invalid_component_dict["credential_type"] = "INVALID" schema = schema_class() @@ -94,9 +100,11 @@ def test_invalid_credential_type( @pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) def test_encorrect_credential_type( - credential_component_class, schema_class, credential_component_type, key, value + credential_component_class, schema_class, credential_component_type, credential_component_data ): - incorrect_component_dict = build_component_dict(credential_component_type, key, value) + incorrect_component_dict = build_component_dict( + credential_component_type, credential_component_data + ) incorrect_component_dict["credential_type"] = ( CredentialComponentType.USERNAME.name if credential_component_type != CredentialComponentType.USERNAME @@ -110,11 +118,11 @@ def test_encorrect_credential_type( @pytest.mark.parametrize(PARAMETRIZED_PARAMETER_NAMES, PARAMETRIZED_PARAMETER_VALUES) def test_invalid_values( - credential_component_class, schema_class, credential_component_type, key, value + credential_component_class, schema_class, credential_component_type, credential_component_data ): schema = schema_class() - for invalid_value in INVALID_VALUES[credential_component_type]: - component_dict = build_component_dict(credential_component_type, key, invalid_value) + for invalid_component_data in INVALID_COMPONENT_DATA[credential_component_type]: + component_dict = build_component_dict(credential_component_type, invalid_component_data) with pytest.raises(ValidationError): credential_component_class(**schema.load(component_dict)) From 0d477cef7cd1356a4c03fdfd4a2f1adc30096274 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 12:32:11 -0400 Subject: [PATCH 13/32] Common: Add SSHKeypairSchema --- monkey/common/credentials/ssh_keypair.py | 10 +++++++++ .../credentials/test_credential_components.py | 22 ++++++++++++++++++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/monkey/common/credentials/ssh_keypair.py b/monkey/common/credentials/ssh_keypair.py index 897acba8a..35f074b13 100644 --- a/monkey/common/credentials/ssh_keypair.py +++ b/monkey/common/credentials/ssh_keypair.py @@ -1,6 +1,16 @@ from dataclasses import dataclass, field +from marshmallow import fields + from . import CredentialComponentType, ICredentialComponent +from .credential_component_schema import CredentialComponentSchema, CredentialTypeField + + +class SSHKeypairSchema(CredentialComponentSchema): + credential_type = CredentialTypeField(CredentialComponentType.SSH_KEYPAIR) + # TODO: Find a list of valid formats for ssh keys and add validators + private_key = fields.Str() + public_key = fields.Str() @dataclass(frozen=True) diff --git a/monkey/tests/unit_tests/common/credentials/test_credential_components.py b/monkey/tests/unit_tests/common/credentials/test_credential_components.py index 6be331356..55a2b9279 100644 --- a/monkey/tests/unit_tests/common/credentials/test_credential_components.py +++ b/monkey/tests/unit_tests/common/credentials/test_credential_components.py @@ -3,10 +3,18 @@ from typing import Any, Mapping import pytest from marshmallow.exceptions import ValidationError -from common.credentials import CredentialComponentType, LMHash, NTHash, Password, Username +from common.credentials import ( + CredentialComponentType, + LMHash, + NTHash, + Password, + SSHKeypair, + Username, +) from common.credentials.lm_hash import LMHashSchema from common.credentials.nt_hash import NTHashSchema from common.credentials.password import PasswordSchema +from common.credentials.ssh_keypair import SSHKeypairSchema from common.credentials.username import UsernameSchema PARAMETRIZED_PARAMETER_NAMES = ( @@ -28,6 +36,12 @@ PARAMETRIZED_PARAMETER_VALUES = [ CredentialComponentType.NT_HASH, {"nt_hash": "E52CAC67419A9A224A3B108F3FA6CB6D"}, ), + ( + SSHKeypair, + SSHKeypairSchema, + CredentialComponentType.SSH_KEYPAIR, + {"public_key": "TEST_PUBLIC_KEY", "private_key": "TEST_PRIVATE_KEY"}, + ), ] @@ -48,6 +62,12 @@ INVALID_COMPONENT_DATA = { {"nt_hash": "0123456789012345678901234568901"}, {"nt_hash": "E52GAC67419A9A224A3B108F3FA6CB6D"}, ), + CredentialComponentType.SSH_KEYPAIR: ( + {"public_key": None, "private_key": "TEST_PRIVATE_KEY"}, + {"public_key": "TEST_PUBLIC_KEY", "private_key": None}, + {"public_key": 1, "private_key": "TEST_PRIVATE_KEY"}, + {"public_key": "TEST_PUBLIC_KEY", "private_key": 999}, + ), } From 92416cb079795e0b78600bd7d86607215e457a5d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 16:00:18 -0400 Subject: [PATCH 14/32] Common: Add validation to LMHash and NTHash --- monkey/common/credentials/lm_hash.py | 5 ++- monkey/common/credentials/nt_hash.py | 5 ++- monkey/common/credentials/validators.py | 35 ++++++++++++++++++- .../common/credentials/test_ntlm_hash.py | 26 ++++++++++++++ 4 files changed, 68 insertions(+), 3 deletions(-) create mode 100644 monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py diff --git a/monkey/common/credentials/lm_hash.py b/monkey/common/credentials/lm_hash.py index 4a1b59f80..5a04e8bae 100644 --- a/monkey/common/credentials/lm_hash.py +++ b/monkey/common/credentials/lm_hash.py @@ -4,7 +4,7 @@ from marshmallow import fields from . import CredentialComponentType, ICredentialComponent from .credential_component_schema import CredentialComponentSchema, CredentialTypeField -from .validators import ntlm_hash_validator +from .validators import credential_component_validator, ntlm_hash_validator class LMHashSchema(CredentialComponentSchema): @@ -18,3 +18,6 @@ class LMHash(ICredentialComponent): default=CredentialComponentType.LM_HASH, init=False ) lm_hash: str + + def __post_init__(self): + credential_component_validator(LMHashSchema(), self) diff --git a/monkey/common/credentials/nt_hash.py b/monkey/common/credentials/nt_hash.py index 838ec4596..a7145a5a0 100644 --- a/monkey/common/credentials/nt_hash.py +++ b/monkey/common/credentials/nt_hash.py @@ -4,7 +4,7 @@ from marshmallow import fields from . import CredentialComponentType, ICredentialComponent from .credential_component_schema import CredentialComponentSchema, CredentialTypeField -from .validators import ntlm_hash_validator +from .validators import credential_component_validator, ntlm_hash_validator class NTHashSchema(CredentialComponentSchema): @@ -18,3 +18,6 @@ class NTHash(ICredentialComponent): default=CredentialComponentType.NT_HASH, init=False ) nt_hash: str + + def __post_init__(self): + credential_component_validator(NTHashSchema(), self) diff --git a/monkey/common/credentials/validators.py b/monkey/common/credentials/validators.py index 3d1b74650..aa5fc7735 100644 --- a/monkey/common/credentials/validators.py +++ b/monkey/common/credentials/validators.py @@ -1,6 +1,39 @@ import re +from typing import Type -from marshmallow import validate +from marshmallow import Schema, validate + +from . import ICredentialComponent _ntlm_hash_regex = re.compile(r"^[a-fA-F0-9]{32}$") ntlm_hash_validator = validate.Regexp(regex=_ntlm_hash_regex) + + +class InvalidCredentialComponent(Exception): + def __init__(self, credential_component_class: Type[ICredentialComponent], message: str): + self._credential_component_name = credential_component_class.__name__ + self._message = message + + def __str__(self) -> str: + return ( + f"Cannot construct a {self._credential_component_name} object with the supplied, " + f"invalid data: {self._message}" + ) + + +def credential_component_validator(schema: Schema, credential_component: ICredentialComponent): + """ + Validate a credential component + + :param schema: A marshmallow schema used for validating the component + :param credential_component: A credential component to be validated + :raises InvalidCredentialComponent: if the credential_component contains invalid data + """ + try: + serialized_data = schema.dump(credential_component) + + # This will raise an exception if the object is invalid. Calling this in __post__init() + # makes it impossible to construct an invalid object + schema.load(serialized_data) + except Exception as err: + raise InvalidCredentialComponent(credential_component.__class__, err) diff --git a/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py b/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py new file mode 100644 index 000000000..ee41a2318 --- /dev/null +++ b/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py @@ -0,0 +1,26 @@ +import pytest + +from common.credentials import LMHash, NTHash + +VALID_HASH = "E520AC67419A9A224A3B108F3FA6CB6D" +INVALID_HASHES = ( + 0, + 1, + 2.0, + "invalid", + "0123456789012345678901234568901", + "E52GAC67419A9A224A3B108F3FA6CB6D", +) + + +@pytest.mark.parametrize("ntlm_hash_class", (LMHash, NTHash)) +def test_construct_valid_ntlm_hash(ntlm_hash_class): + # This test will fail if an exception is raised + ntlm_hash_class(VALID_HASH) + + +@pytest.mark.parametrize("ntlm_hash_class", (LMHash, NTHash)) +def test_construct_invalid_ntlm_hash(ntlm_hash_class): + for invalid_hash in INVALID_HASHES: + with pytest.raises(Exception): + ntlm_hash_class(invalid_hash) From 82fb693f064f13619366cf733f6adf998f519e68 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 16:03:12 -0400 Subject: [PATCH 15/32] Common: Simplify raising of InvalidConfigurationError --- .../common/configuration/agent_configuration.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/monkey/common/configuration/agent_configuration.py b/monkey/common/configuration/agent_configuration.py index 097b86382..5b801610d 100644 --- a/monkey/common/configuration/agent_configuration.py +++ b/monkey/common/configuration/agent_configuration.py @@ -17,16 +17,19 @@ from .agent_sub_configurations import ( PropagationConfiguration, ) - -class InvalidConfigurationError(Exception): - pass - - INVALID_CONFIGURATION_ERROR_MESSAGE = ( "Cannot construct an AgentConfiguration object with the supplied, invalid data:" ) +class InvalidConfigurationError(Exception): + def __init__(self, message: str): + self._message = message + + def __str__(self) -> str: + return f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {self._message}" + + @dataclass(frozen=True) class AgentConfiguration: keep_tunnel_open_time: float @@ -42,7 +45,7 @@ class AgentConfiguration: try: AgentConfigurationSchema().dump(self) except Exception as err: - raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}") + raise InvalidConfigurationError(err) @staticmethod def from_mapping(config_mapping: Mapping[str, Any]) -> AgentConfiguration: From d3859debbe889c0ed63d2a3fe76d1b424b799718 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 19:28:36 -0400 Subject: [PATCH 16/32] UT: Use valid NTLM hashes in test_mimikatz_collector.py --- .../test_mimikatz_collector.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py index f87dea93c..62142f6e9 100644 --- a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py +++ b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py @@ -56,14 +56,16 @@ def test_pypykatz_result_parsing_duplicates(monkeypatch): def test_pypykatz_result_parsing_defaults(monkeypatch): win_creds = [ - WindowsCredentials(username="user2", password="secret2", lm_hash="lm_hash"), + WindowsCredentials( + username="user2", password="secret2", lm_hash="0182BD0BD4444BF8FC83B5D9042EED2E" + ), ] patch_pypykatz(win_creds, monkeypatch) # Expected credentials username = Username("user2") password = Password("secret2") - lm_hash = LMHash("lm_hash") + lm_hash = LMHash("0182BD0BD4444BF8FC83B5D9042EED2E") expected_credentials = Credentials([username], [password, lm_hash]) collected_credentials = collect_credentials() @@ -73,12 +75,17 @@ def test_pypykatz_result_parsing_defaults(monkeypatch): def test_pypykatz_result_parsing_no_identities(monkeypatch): win_creds = [ - WindowsCredentials(username="", password="", ntlm_hash="ntlm_hash", lm_hash="lm_hash"), + WindowsCredentials( + username="", + password="", + ntlm_hash="E9F85516721DDC218359AD5280DB4450", + lm_hash="0182BD0BD4444BF8FC83B5D9042EED2E", + ), ] patch_pypykatz(win_creds, monkeypatch) - lm_hash = LMHash("lm_hash") - nt_hash = NTHash("ntlm_hash") + lm_hash = LMHash("0182BD0BD4444BF8FC83B5D9042EED2E") + nt_hash = NTHash("E9F85516721DDC218359AD5280DB4450") expected_credentials = Credentials([], [lm_hash, nt_hash]) collected_credentials = collect_credentials() From 82ce091063c50fe3df2d98239f9f0b2b52e45fa0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 19:32:18 -0400 Subject: [PATCH 17/32] Common: Export InvalidCredentialComponent from common.credentials --- monkey/common/credentials/__init__.py | 6 +++++- .../tests/unit_tests/common/credentials/test_ntlm_hash.py | 4 ++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/monkey/common/credentials/__init__.py b/monkey/common/credentials/__init__.py index 92a778886..f49f6af03 100644 --- a/monkey/common/credentials/__init__.py +++ b/monkey/common/credentials/__init__.py @@ -1,8 +1,12 @@ from .credential_component_type import CredentialComponentType from .i_credential_component import ICredentialComponent -from .credentials import Credentials + +from .validators import InvalidCredentialComponent + from .lm_hash import LMHash from .nt_hash import NTHash from .password import Password from .ssh_keypair import SSHKeypair from .username import Username + +from .credentials import Credentials diff --git a/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py b/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py index ee41a2318..28f7bcaae 100644 --- a/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py +++ b/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py @@ -1,6 +1,6 @@ import pytest -from common.credentials import LMHash, NTHash +from common.credentials import InvalidCredentialComponent, LMHash, NTHash VALID_HASH = "E520AC67419A9A224A3B108F3FA6CB6D" INVALID_HASHES = ( @@ -22,5 +22,5 @@ def test_construct_valid_ntlm_hash(ntlm_hash_class): @pytest.mark.parametrize("ntlm_hash_class", (LMHash, NTHash)) def test_construct_invalid_ntlm_hash(ntlm_hash_class): for invalid_hash in INVALID_HASHES: - with pytest.raises(Exception): + with pytest.raises(InvalidCredentialComponent): ntlm_hash_class(invalid_hash) From e4d38631b82fa2dd63b3959bfabb4366bfa5309f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 20:45:49 -0400 Subject: [PATCH 18/32] Common: Serialize/Deserialize Credentials --- monkey/common/credentials/credentials.py | 125 +++++++++++++++++- .../common/credentials/test_credentials.py | 57 ++++++++ 2 files changed, 180 insertions(+), 2 deletions(-) create mode 100644 monkey/tests/unit_tests/common/credentials/test_credentials.py diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index d5591f6d7..f60c71fbe 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -1,10 +1,131 @@ -from dataclasses import dataclass -from typing import Tuple +from __future__ import annotations +from dataclasses import dataclass +from typing import Any, Mapping, Tuple + +from marshmallow import INCLUDE, Schema, fields, post_load, pre_dump +from marshmallow_enum import EnumField + +from . import CredentialComponentType, LMHash, NTHash, Password, SSHKeypair, Username from .i_credential_component import ICredentialComponent +from .lm_hash import LMHashSchema +from .nt_hash import NTHashSchema +from .password import PasswordSchema +from .ssh_keypair import SSHKeypairSchema +from .username import UsernameSchema + +CREDENTIL_COMPINENT_TYPE_TO_CLASS = { + CredentialComponentType.LM_HASH: LMHash, + CredentialComponentType.NT_HASH: NTHash, + CredentialComponentType.PASSWORD: Password, + CredentialComponentType.SSH_KEYPAIR: SSHKeypair, + CredentialComponentType.USERNAME: Username, +} +CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA = { + CredentialComponentType.LM_HASH: LMHashSchema(), + CredentialComponentType.NT_HASH: NTHashSchema(), + CredentialComponentType.PASSWORD: PasswordSchema(), + CredentialComponentType.SSH_KEYPAIR: SSHKeypairSchema(), + CredentialComponentType.USERNAME: UsernameSchema(), +} + + +class GenericCredentialComponentSchema(Schema): + class Meta: + unknown = INCLUDE + + credential_type = EnumField(CredentialComponentType) + + """ + @post_load + def _string_to_enum(self, data, **kwargs) -> Mapping[str, Any]: + data["credential_type"] = CredentialComponentType[data["credential_type"]] + """ + + +class CredentialsSchema(Schema): + # Use fields.List instead of fields.Tuple because marshmallow requires fields.Tuple to have a + # fixed length. + # identities = fields.List(fields.Nested(GenericCredentialComponentSchema)) + # secrets = fields.List(fields.Nested(GenericCredentialComponentSchema)) + identities = fields.List(fields.Mapping()) + secrets = fields.List(fields.Mapping()) + + @post_load + def _make_credentials(self, data, **kwargs) -> Mapping[str, Tuple(Mapping)]: + from pprint import pprint + + for component in data["identities"]: + pprint(component) + data["identities"] = tuple( + [ + CredentialsSchema._build_credential_component(component) + for component in data["identities"] + ] + ) + data["secrets"] = tuple( + [ + CredentialsSchema._build_credential_component(component) + for component in data["secrets"] + ] + ) + + return data + + @staticmethod + def _build_credential_component(data: Mapping[str, Any]): + credential_component_type = CredentialComponentType[data["credential_type"]] + credential_component_class = CREDENTIL_COMPINENT_TYPE_TO_CLASS[credential_component_type] + credential_component_schema = CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ + credential_component_type + ] + + data["credential_type"] = data["credential_type"] + return credential_component_class(**credential_component_schema.load(data)) + + @pre_dump + def _serialize_credentials( + self, credentials: Credentials, **kwargs + ) -> Mapping[str, Tuple[Mapping[str, Any]]]: + data = {} + data["identities"] = tuple( + [ + CredentialsSchema._serialize_credential_component(component) + for component in credentials.identities + ] + ) + data["secrets"] = tuple( + [ + CredentialsSchema._serialize_credential_component(component) + for component in credentials.secrets + ] + ) + + return data + + @staticmethod + def _serialize_credential_component(credential_component: ICredentialComponent): + credential_component_schema = CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ + credential_component.credential_type + ] + return credential_component_schema.dump(credential_component) @dataclass(frozen=True) class Credentials: identities: Tuple[ICredentialComponent] secrets: Tuple[ICredentialComponent] + + @staticmethod + def from_mapping(credentials: Mapping) -> Credentials: + deserialized_data = CredentialsSchema().load(credentials) + return Credentials(**deserialized_data) + + @staticmethod + def from_json(credentials: str) -> Credentials: + deserialized_data = CredentialsSchema().loads(credentials) + return Credentials(**deserialized_data) + + @staticmethod + def to_json(credentials: Credentials) -> str: + return CredentialsSchema().dumps(credentials) diff --git a/monkey/tests/unit_tests/common/credentials/test_credentials.py b/monkey/tests/unit_tests/common/credentials/test_credentials.py new file mode 100644 index 000000000..4496a2f05 --- /dev/null +++ b/monkey/tests/unit_tests/common/credentials/test_credentials.py @@ -0,0 +1,57 @@ +import json + +from common.credentials import Credentials, LMHash, NTHash, Password, SSHKeypair, Username + +USER1 = "test_user_1" +USER2 = "test_user_2" +PASSWORD = "12435" +LM_HASH = "AEBD4DE384C7EC43AAD3B435B51404EE" +NT_HASH = "7A21990FCD3D759941E45C490F143D5F" +PUBLIC_KEY = "MY_PUBLIC_KEY" +PRIVATE_KEY = "MY_PRIVATE_KEY" + +CREDENTIALS_DICT = { + "identities": [ + {"credential_type": "USERNAME", "username": USER1}, + {"credential_type": "USERNAME", "username": USER2}, + ], + "secrets": [ + {"credential_type": "PASSWORD", "password": PASSWORD}, + {"credential_type": "LM_HASH", "lm_hash": LM_HASH}, + {"credential_type": "NT_HASH", "nt_hash": NT_HASH}, + { + "credential_type": "SSH_KEYPAIR", + "public_key": PUBLIC_KEY, + "private_key": PRIVATE_KEY, + }, + ], +} + +CREDENTIALS_JSON = json.dumps(CREDENTIALS_DICT) + +IDENTITIES = (Username(USER1), Username(USER2)) +SECRETS = ( + Password(PASSWORD), + LMHash(LM_HASH), + NTHash(NT_HASH), + SSHKeypair(PRIVATE_KEY, PUBLIC_KEY), +) +CREDENTIALS_OBJECT = Credentials(IDENTITIES, SECRETS) + + +def test_credentials_serialization_json(): + serialized_credentials = Credentials.to_json(CREDENTIALS_OBJECT) + + assert json.loads(serialized_credentials) == CREDENTIALS_DICT + + +def test_credentials_deserialization__from_mapping(): + deserialized_credentials = Credentials.from_mapping(CREDENTIALS_DICT) + + assert deserialized_credentials == CREDENTIALS_OBJECT + + +def test_credentials_deserialization__from_json(): + deserialized_credentials = Credentials.from_json(CREDENTIALS_JSON) + + assert deserialized_credentials == CREDENTIALS_OBJECT From c58d26a5e47c5ca54cc72329ef3076d9bca789cc Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 20:48:54 -0400 Subject: [PATCH 19/32] Common: Clean up credentials.py --- monkey/common/credentials/credentials.py | 38 ++++++------------------ 1 file changed, 9 insertions(+), 29 deletions(-) diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index f60c71fbe..66f891aa1 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -1,10 +1,9 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Mapping, Tuple +from typing import Any, Mapping, MutableMapping, Tuple -from marshmallow import INCLUDE, Schema, fields, post_load, pre_dump -from marshmallow_enum import EnumField +from marshmallow import Schema, fields, post_load, pre_dump from . import CredentialComponentType, LMHash, NTHash, Password, SSHKeypair, Username from .i_credential_component import ICredentialComponent @@ -14,14 +13,14 @@ from .password import PasswordSchema from .ssh_keypair import SSHKeypairSchema from .username import UsernameSchema -CREDENTIL_COMPINENT_TYPE_TO_CLASS = { +CREDENTIAL_COMPINENT_TYPE_TO_CLASS = { CredentialComponentType.LM_HASH: LMHash, CredentialComponentType.NT_HASH: NTHash, CredentialComponentType.PASSWORD: Password, CredentialComponentType.SSH_KEYPAIR: SSHKeypair, CredentialComponentType.USERNAME: Username, } -CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA = { +CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA = { CredentialComponentType.LM_HASH: LMHashSchema(), CredentialComponentType.NT_HASH: NTHashSchema(), CredentialComponentType.PASSWORD: PasswordSchema(), @@ -30,33 +29,14 @@ CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA = { } -class GenericCredentialComponentSchema(Schema): - class Meta: - unknown = INCLUDE - - credential_type = EnumField(CredentialComponentType) - - """ - @post_load - def _string_to_enum(self, data, **kwargs) -> Mapping[str, Any]: - data["credential_type"] = CredentialComponentType[data["credential_type"]] - """ - - class CredentialsSchema(Schema): # Use fields.List instead of fields.Tuple because marshmallow requires fields.Tuple to have a # fixed length. - # identities = fields.List(fields.Nested(GenericCredentialComponentSchema)) - # secrets = fields.List(fields.Nested(GenericCredentialComponentSchema)) identities = fields.List(fields.Mapping()) secrets = fields.List(fields.Mapping()) @post_load - def _make_credentials(self, data, **kwargs) -> Mapping[str, Tuple(Mapping)]: - from pprint import pprint - - for component in data["identities"]: - pprint(component) + def _make_credentials(self, data, **kwargs) -> Mapping[str, Tuple[Mapping[str, Any]]]: data["identities"] = tuple( [ CredentialsSchema._build_credential_component(component) @@ -73,10 +53,10 @@ class CredentialsSchema(Schema): return data @staticmethod - def _build_credential_component(data: Mapping[str, Any]): + def _build_credential_component(data: MutableMapping[str, Any]): credential_component_type = CredentialComponentType[data["credential_type"]] - credential_component_class = CREDENTIL_COMPINENT_TYPE_TO_CLASS[credential_component_type] - credential_component_schema = CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ + credential_component_class = CREDENTIAL_COMPINENT_TYPE_TO_CLASS[credential_component_type] + credential_component_schema = CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ credential_component_type ] @@ -105,7 +85,7 @@ class CredentialsSchema(Schema): @staticmethod def _serialize_credential_component(credential_component: ICredentialComponent): - credential_component_schema = CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ + credential_component_schema = CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ credential_component.credential_type ] return credential_component_schema.dump(credential_component) From 59a9aa8a535e8f0a3569d3fec12c0b0926bef085 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 07:37:15 -0400 Subject: [PATCH 20/32] Project: Ignore decorated Schema methods in credentials.py --- vulture_allowlist.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index b04ec8cca..049e487fc 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -197,7 +197,10 @@ _make_network_scan_configuration # unused method (monkey/common/configuration/a _make_propagation_configuration # unused method (monkey/common/configuration/agent_configuration.py:167) # Credentials -_strip_credential_type # unused method(monkey/common/credentials/password.py:18) +_strip_credential_type # unused method (monkey/common/credentials/password.py:18) +_make_credentials # unused method (monkey/common/credentials/credentials:39) +_serialize_credentials # unused method (monkey/common/credentials/credentials:67) + # Models _make_simulation # unused method (monkey/monkey_island/cc/models/simulation.py:19 From 008428e3187d6a05a4bedd752749f8b0bca0e62a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 07:43:12 -0400 Subject: [PATCH 21/32] Common: Fix type hints in credentials.py --- monkey/common/credentials/credentials.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index 66f891aa1..c36879799 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Mapping, MutableMapping, Tuple +from typing import Any, Mapping, MutableMapping, Sequence, Tuple from marshmallow import Schema, fields, post_load, pre_dump @@ -20,6 +20,7 @@ CREDENTIAL_COMPINENT_TYPE_TO_CLASS = { CredentialComponentType.SSH_KEYPAIR: SSHKeypair, CredentialComponentType.USERNAME: Username, } + CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA = { CredentialComponentType.LM_HASH: LMHashSchema(), CredentialComponentType.NT_HASH: NTHashSchema(), @@ -36,7 +37,7 @@ class CredentialsSchema(Schema): secrets = fields.List(fields.Mapping()) @post_load - def _make_credentials(self, data, **kwargs) -> Mapping[str, Tuple[Mapping[str, Any]]]: + def _make_credentials(self, data, **kwargs) -> Mapping[str, Sequence[Mapping[str, Any]]]: data["identities"] = tuple( [ CredentialsSchema._build_credential_component(component) @@ -66,7 +67,7 @@ class CredentialsSchema(Schema): @pre_dump def _serialize_credentials( self, credentials: Credentials, **kwargs - ) -> Mapping[str, Tuple[Mapping[str, Any]]]: + ) -> Mapping[str, Sequence[Mapping[str, Any]]]: data = {} data["identities"] = tuple( [ @@ -84,7 +85,9 @@ class CredentialsSchema(Schema): return data @staticmethod - def _serialize_credential_component(credential_component: ICredentialComponent): + def _serialize_credential_component( + credential_component: ICredentialComponent, + ) -> Mapping[str, Any]: credential_component_schema = CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ credential_component.credential_type ] From 907b35990d723c1005b664356204fda48ec7facc Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 07:45:28 -0400 Subject: [PATCH 22/32] Common: Fix COMPINENT -> COMPONENT --- monkey/common/credentials/credentials.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index c36879799..471d57d09 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -13,7 +13,7 @@ from .password import PasswordSchema from .ssh_keypair import SSHKeypairSchema from .username import UsernameSchema -CREDENTIAL_COMPINENT_TYPE_TO_CLASS = { +CREDENTIAL_COMPONENT_TYPE_TO_CLASS = { CredentialComponentType.LM_HASH: LMHash, CredentialComponentType.NT_HASH: NTHash, CredentialComponentType.PASSWORD: Password, @@ -21,7 +21,7 @@ CREDENTIAL_COMPINENT_TYPE_TO_CLASS = { CredentialComponentType.USERNAME: Username, } -CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA = { +CREDENTIAL_COMPONENT_TYPE_TO_CLASS_SCHEMA = { CredentialComponentType.LM_HASH: LMHashSchema(), CredentialComponentType.NT_HASH: NTHashSchema(), CredentialComponentType.PASSWORD: PasswordSchema(), @@ -56,8 +56,8 @@ class CredentialsSchema(Schema): @staticmethod def _build_credential_component(data: MutableMapping[str, Any]): credential_component_type = CredentialComponentType[data["credential_type"]] - credential_component_class = CREDENTIAL_COMPINENT_TYPE_TO_CLASS[credential_component_type] - credential_component_schema = CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ + credential_component_class = CREDENTIAL_COMPONENT_TYPE_TO_CLASS[credential_component_type] + credential_component_schema = CREDENTIAL_COMPONENT_TYPE_TO_CLASS_SCHEMA[ credential_component_type ] @@ -88,7 +88,7 @@ class CredentialsSchema(Schema): def _serialize_credential_component( credential_component: ICredentialComponent, ) -> Mapping[str, Any]: - credential_component_schema = CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ + credential_component_schema = CREDENTIAL_COMPONENT_TYPE_TO_CLASS_SCHEMA[ credential_component.credential_type ] return credential_component_schema.dump(credential_component) From e76d72e07e4682886282fb69361a3d3697d6923c Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 07:45:44 -0400 Subject: [PATCH 23/32] Common: Remove NOOP assignment --- monkey/common/credentials/credentials.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index 471d57d09..a300405be 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -61,7 +61,6 @@ class CredentialsSchema(Schema): credential_component_type ] - data["credential_type"] = data["credential_type"] return credential_component_class(**credential_component_schema.load(data)) @pre_dump From 3f61ddd584ea3313ed7e8ff55749af59fb078f46 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 07:48:54 -0400 Subject: [PATCH 24/32] Common: Fix type hints in credentials.py --- monkey/common/credentials/credentials.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index a300405be..75775f22d 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -37,7 +37,9 @@ class CredentialsSchema(Schema): secrets = fields.List(fields.Mapping()) @post_load - def _make_credentials(self, data, **kwargs) -> Mapping[str, Sequence[Mapping[str, Any]]]: + def _make_credentials( + self, data: MutableMapping, **kwargs: Mapping[str, Any] + ) -> Mapping[str, Sequence[Mapping[str, Any]]]: data["identities"] = tuple( [ CredentialsSchema._build_credential_component(component) @@ -54,7 +56,7 @@ class CredentialsSchema(Schema): return data @staticmethod - def _build_credential_component(data: MutableMapping[str, Any]): + def _build_credential_component(data: Mapping[str, Any]) -> ICredentialComponent: credential_component_type = CredentialComponentType[data["credential_type"]] credential_component_class = CREDENTIAL_COMPONENT_TYPE_TO_CLASS[credential_component_type] credential_component_schema = CREDENTIAL_COMPONENT_TYPE_TO_CLASS_SCHEMA[ @@ -68,6 +70,7 @@ class CredentialsSchema(Schema): self, credentials: Credentials, **kwargs ) -> Mapping[str, Sequence[Mapping[str, Any]]]: data = {} + data["identities"] = tuple( [ CredentialsSchema._serialize_credential_component(component) @@ -90,6 +93,7 @@ class CredentialsSchema(Schema): credential_component_schema = CREDENTIAL_COMPONENT_TYPE_TO_CLASS_SCHEMA[ credential_component.credential_type ] + return credential_component_schema.dump(credential_component) From a18eb1cb73e48194fb18ecd832b00394ac38388d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 08:31:28 -0400 Subject: [PATCH 25/32] Common: Add error trapping to Credentials deserialization --- monkey/common/credentials/__init__.py | 2 +- monkey/common/credentials/credentials.py | 35 +++++++++++++++---- monkey/common/credentials/validators.py | 15 ++++++-- .../common/credentials/test_credentials.py | 24 ++++++++++++- .../common/credentials/test_ntlm_hash.py | 4 +-- 5 files changed, 68 insertions(+), 12 deletions(-) diff --git a/monkey/common/credentials/__init__.py b/monkey/common/credentials/__init__.py index f49f6af03..6275e0985 100644 --- a/monkey/common/credentials/__init__.py +++ b/monkey/common/credentials/__init__.py @@ -1,7 +1,7 @@ from .credential_component_type import CredentialComponentType from .i_credential_component import ICredentialComponent -from .validators import InvalidCredentialComponent +from .validators import InvalidCredentialComponentError, InvalidCredentialsError from .lm_hash import LMHash from .nt_hash import NTHash diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index 75775f22d..1837fd8b6 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -5,7 +5,16 @@ from typing import Any, Mapping, MutableMapping, Sequence, Tuple from marshmallow import Schema, fields, post_load, pre_dump -from . import CredentialComponentType, LMHash, NTHash, Password, SSHKeypair, Username +from . import ( + CredentialComponentType, + InvalidCredentialComponentError, + InvalidCredentialsError, + LMHash, + NTHash, + Password, + SSHKeypair, + Username, +) from .i_credential_component import ICredentialComponent from .lm_hash import LMHashSchema from .nt_hash import NTHashSchema @@ -57,7 +66,11 @@ class CredentialsSchema(Schema): @staticmethod def _build_credential_component(data: Mapping[str, Any]) -> ICredentialComponent: - credential_component_type = CredentialComponentType[data["credential_type"]] + try: + credential_component_type = CredentialComponentType[data["credential_type"]] + except KeyError as err: + raise InvalidCredentialsError(f"Unknown credential component type {err}") + credential_component_class = CREDENTIAL_COMPONENT_TYPE_TO_CLASS[credential_component_type] credential_component_schema = CREDENTIAL_COMPONENT_TYPE_TO_CLASS_SCHEMA[ credential_component_type @@ -104,13 +117,23 @@ class Credentials: @staticmethod def from_mapping(credentials: Mapping) -> Credentials: - deserialized_data = CredentialsSchema().load(credentials) - return Credentials(**deserialized_data) + try: + deserialized_data = CredentialsSchema().load(credentials) + return Credentials(**deserialized_data) + except (InvalidCredentialsError, InvalidCredentialComponentError) as err: + raise err + except Exception as err: + raise InvalidCredentialsError(str(err)) @staticmethod def from_json(credentials: str) -> Credentials: - deserialized_data = CredentialsSchema().loads(credentials) - return Credentials(**deserialized_data) + try: + deserialized_data = CredentialsSchema().loads(credentials) + return Credentials(**deserialized_data) + except (InvalidCredentialsError, InvalidCredentialComponentError) as err: + raise err + except Exception as err: + raise InvalidCredentialsError(str(err)) @staticmethod def to_json(credentials: Credentials) -> str: diff --git a/monkey/common/credentials/validators.py b/monkey/common/credentials/validators.py index aa5fc7735..2e0e2e93c 100644 --- a/monkey/common/credentials/validators.py +++ b/monkey/common/credentials/validators.py @@ -9,7 +9,7 @@ _ntlm_hash_regex = re.compile(r"^[a-fA-F0-9]{32}$") ntlm_hash_validator = validate.Regexp(regex=_ntlm_hash_regex) -class InvalidCredentialComponent(Exception): +class InvalidCredentialComponentError(Exception): def __init__(self, credential_component_class: Type[ICredentialComponent], message: str): self._credential_component_name = credential_component_class.__name__ self._message = message @@ -21,6 +21,17 @@ class InvalidCredentialComponent(Exception): ) +class InvalidCredentialsError(Exception): + def __init__(self, message: str): + self._message = message + + def __str__(self) -> str: + return ( + f"Cannot construct a Credentials object with the supplied, " + f"invalid data: {self._message}" + ) + + def credential_component_validator(schema: Schema, credential_component: ICredentialComponent): """ Validate a credential component @@ -36,4 +47,4 @@ def credential_component_validator(schema: Schema, credential_component: ICreden # makes it impossible to construct an invalid object schema.load(serialized_data) except Exception as err: - raise InvalidCredentialComponent(credential_component.__class__, err) + raise InvalidCredentialComponentError(credential_component.__class__, err) diff --git a/monkey/tests/unit_tests/common/credentials/test_credentials.py b/monkey/tests/unit_tests/common/credentials/test_credentials.py index 4496a2f05..d6fd55ef5 100644 --- a/monkey/tests/unit_tests/common/credentials/test_credentials.py +++ b/monkey/tests/unit_tests/common/credentials/test_credentials.py @@ -1,6 +1,16 @@ import json -from common.credentials import Credentials, LMHash, NTHash, Password, SSHKeypair, Username +import pytest + +from common.credentials import ( + Credentials, + InvalidCredentialsError, + LMHash, + NTHash, + Password, + SSHKeypair, + Username, +) USER1 = "test_user_1" USER2 = "test_user_2" @@ -55,3 +65,15 @@ def test_credentials_deserialization__from_json(): deserialized_credentials = Credentials.from_json(CREDENTIALS_JSON) assert deserialized_credentials == CREDENTIALS_OBJECT + + +def test_credentials_deserialization__invalid_credentials(): + invalid_data = {"secrets": [], "unknown_key": []} + with pytest.raises(InvalidCredentialsError): + Credentials.from_mapping(invalid_data) + + +def test_credentials_deserialization__invalid_component_type(): + invalid_data = {"secrets": [], "identities": [{"credential_type": "FAKE", "username": "user1"}]} + with pytest.raises(InvalidCredentialsError): + Credentials.from_mapping(invalid_data) diff --git a/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py b/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py index 28f7bcaae..5f50110e8 100644 --- a/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py +++ b/monkey/tests/unit_tests/common/credentials/test_ntlm_hash.py @@ -1,6 +1,6 @@ import pytest -from common.credentials import InvalidCredentialComponent, LMHash, NTHash +from common.credentials import InvalidCredentialComponentError, LMHash, NTHash VALID_HASH = "E520AC67419A9A224A3B108F3FA6CB6D" INVALID_HASHES = ( @@ -22,5 +22,5 @@ def test_construct_valid_ntlm_hash(ntlm_hash_class): @pytest.mark.parametrize("ntlm_hash_class", (LMHash, NTHash)) def test_construct_invalid_ntlm_hash(ntlm_hash_class): for invalid_hash in INVALID_HASHES: - with pytest.raises(InvalidCredentialComponent): + with pytest.raises(InvalidCredentialComponentError): ntlm_hash_class(invalid_hash) From 6bb6aa5250ffabb0618ae78bb8b11db2487c5fc6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 08:38:39 -0400 Subject: [PATCH 26/32] Common: Remove INVALID_CONFIGURATION_ERROR_MESSAGE --- .../common/configuration/agent_configuration.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/monkey/common/configuration/agent_configuration.py b/monkey/common/configuration/agent_configuration.py index 5b801610d..d6e61ddd1 100644 --- a/monkey/common/configuration/agent_configuration.py +++ b/monkey/common/configuration/agent_configuration.py @@ -17,17 +17,16 @@ from .agent_sub_configurations import ( PropagationConfiguration, ) -INVALID_CONFIGURATION_ERROR_MESSAGE = ( - "Cannot construct an AgentConfiguration object with the supplied, invalid data:" -) - class InvalidConfigurationError(Exception): def __init__(self, message: str): self._message = message def __str__(self) -> str: - return f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {self._message}" + return ( + f"Cannot construct an AgentConfiguration object with the supplied, invalid data: " + f"{self._message}" + ) @dataclass(frozen=True) @@ -45,7 +44,7 @@ class AgentConfiguration: try: AgentConfigurationSchema().dump(self) except Exception as err: - raise InvalidConfigurationError(err) + raise InvalidConfigurationError(str(err)) @staticmethod def from_mapping(config_mapping: Mapping[str, Any]) -> AgentConfiguration: @@ -62,7 +61,7 @@ class AgentConfiguration: config_dict = AgentConfigurationSchema().load(config_mapping) return AgentConfiguration(**config_dict) except MarshmallowError as err: - raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}") + raise InvalidConfigurationError(str(err)) @staticmethod def from_json(config_json: str) -> AgentConfiguration: @@ -78,7 +77,7 @@ class AgentConfiguration: config_dict = AgentConfigurationSchema().loads(config_json) return AgentConfiguration(**config_dict) except MarshmallowError as err: - raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}") + raise InvalidConfigurationError(str(err)) @staticmethod def to_json(config: AgentConfiguration) -> str: From 06fc4aaad6b197f635a9cb6c2da6d681e4f5b117 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 08:47:41 -0400 Subject: [PATCH 27/32] Common: Catch MarshmallowError instead of Exception --- monkey/common/credentials/credentials.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index 1837fd8b6..732818783 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Any, Mapping, MutableMapping, Sequence, Tuple from marshmallow import Schema, fields, post_load, pre_dump +from marshmallow.exceptions import MarshmallowError from . import ( CredentialComponentType, @@ -122,7 +123,7 @@ class Credentials: return Credentials(**deserialized_data) except (InvalidCredentialsError, InvalidCredentialComponentError) as err: raise err - except Exception as err: + except MarshmallowError as err: raise InvalidCredentialsError(str(err)) @staticmethod @@ -132,7 +133,7 @@ class Credentials: return Credentials(**deserialized_data) except (InvalidCredentialsError, InvalidCredentialComponentError) as err: raise err - except Exception as err: + except MarshmallowError as err: raise InvalidCredentialsError(str(err)) @staticmethod From 9ea0fb87ea3704ae6153c751adb6c411bf6c0613 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 08:52:33 -0400 Subject: [PATCH 28/32] Common: Raise InvalidCredentialComponentError from Credentials --- monkey/common/credentials/credentials.py | 5 ++++- .../unit_tests/common/credentials/test_credentials.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py index 732818783..e04be27a5 100644 --- a/monkey/common/credentials/credentials.py +++ b/monkey/common/credentials/credentials.py @@ -77,7 +77,10 @@ class CredentialsSchema(Schema): credential_component_type ] - return credential_component_class(**credential_component_schema.load(data)) + try: + return credential_component_class(**credential_component_schema.load(data)) + except MarshmallowError as err: + raise InvalidCredentialComponentError(credential_component_class, str(err)) @pre_dump def _serialize_credentials( diff --git a/monkey/tests/unit_tests/common/credentials/test_credentials.py b/monkey/tests/unit_tests/common/credentials/test_credentials.py index d6fd55ef5..c68e1c813 100644 --- a/monkey/tests/unit_tests/common/credentials/test_credentials.py +++ b/monkey/tests/unit_tests/common/credentials/test_credentials.py @@ -4,6 +4,7 @@ import pytest from common.credentials import ( Credentials, + InvalidCredentialComponentError, InvalidCredentialsError, LMHash, NTHash, @@ -77,3 +78,12 @@ def test_credentials_deserialization__invalid_component_type(): invalid_data = {"secrets": [], "identities": [{"credential_type": "FAKE", "username": "user1"}]} with pytest.raises(InvalidCredentialsError): Credentials.from_mapping(invalid_data) + + +def test_credentials_deserialization__invalid_component(): + invalid_data = { + "secrets": [], + "identities": [{"credential_type": "USERNAME", "unknown_field": "user1"}], + } + with pytest.raises(InvalidCredentialComponentError): + Credentials.from_mapping(invalid_data) From e921f90e002d2fb402622e1220b17f98bc0bf8be Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 09:15:01 -0400 Subject: [PATCH 29/32] Agent: Use Credentials.to_json() for CredentialsTelem serialization --- .../telemetry/credentials_telem.py | 28 +++---------------- .../telemetry/test_credentials_telem.py | 3 +- 2 files changed, 5 insertions(+), 26 deletions(-) diff --git a/monkey/infection_monkey/telemetry/credentials_telem.py b/monkey/infection_monkey/telemetry/credentials_telem.py index 11e2bbb6d..7504b9c65 100644 --- a/monkey/infection_monkey/telemetry/credentials_telem.py +++ b/monkey/infection_monkey/telemetry/credentials_telem.py @@ -1,9 +1,8 @@ -import enum import json -from typing import Dict, Iterable +from typing import Iterable from common.common_consts.telem_categories import TelemCategoryEnum -from common.credentials import Credentials, ICredentialComponent +from common.credentials import Credentials from infection_monkey.telemetry.base_telem import BaseTelem @@ -24,24 +23,5 @@ class CredentialsTelem(BaseTelem): def send(self, log_data=True): super().send(log_data=False) - def get_data(self) -> Dict: - # TODO: At a later time we can consider factoring this into a Serializer class or similar. - return json.loads(json.dumps(self._credentials, default=_serialize)) - - -def _serialize(obj): - if isinstance(obj, enum.Enum): - return obj.name - - if isinstance(obj, ICredentialComponent): - # This is a workaround for ICredentialComponents that are implemented as dataclasses. If the - # credential_type attribute is populated with `field(init=False, ...)`, then credential_type - # is not added to the object's __dict__ attribute. The biggest risk of this workaround is - # that we might change the name of the credential_type field in ICredentialComponents, but - # automated refactoring tools would not detect that this string needs to change. This is - # mittigated by the call to getattr() below, which will raise an AttributeException if the - # attribute name changes and a unit test will fail under these conditions. - credential_type = getattr(obj, "credential_type") - return dict(obj.__dict__, **{"credential_type": credential_type}) - - return getattr(obj, "__dict__", str(obj)) + def get_data(self): + return [json.loads(Credentials.to_json(c)) for c in self._credentials] diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py b/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py index 5fca80eff..701071fbb 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py @@ -38,8 +38,7 @@ def test_credential_telem_send(spy_send_telemetry, credentials_for_test): telem = CredentialsTelem([credentials_for_test]) telem.send() - expected_data = json.dumps(expected_data, cls=telem.json_encoder) - assert spy_send_telemetry.data == expected_data + assert json.loads(spy_send_telemetry.data) == expected_data assert spy_send_telemetry.telem_category == "credentials" From 08bb49af0f59842424c316a19b01cb89b8884111 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 11:25:33 -0400 Subject: [PATCH 30/32] Common: Add ssh-key-regex note to TODO --- monkey/common/credentials/ssh_keypair.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/common/credentials/ssh_keypair.py b/monkey/common/credentials/ssh_keypair.py index 35f074b13..6b8dcded2 100644 --- a/monkey/common/credentials/ssh_keypair.py +++ b/monkey/common/credentials/ssh_keypair.py @@ -8,7 +8,8 @@ from .credential_component_schema import CredentialComponentSchema, CredentialTy class SSHKeypairSchema(CredentialComponentSchema): credential_type = CredentialTypeField(CredentialComponentType.SSH_KEYPAIR) - # TODO: Find a list of valid formats for ssh keys and add validators + # TODO: Find a list of valid formats for ssh keys and add validators. + # See https://github.com/nemchik/ssh-key-regex private_key = fields.Str() public_key = fields.Str() From d0fa9a7dcf1822fc9278744ef23c0886d894fdf1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 11:29:08 -0400 Subject: [PATCH 31/32] Common: Use the imperative in del_key() docstring --- monkey/common/utils/code_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/common/utils/code_utils.py b/monkey/common/utils/code_utils.py index a0fe6b5a9..a7625e0c6 100644 --- a/monkey/common/utils/code_utils.py +++ b/monkey/common/utils/code_utils.py @@ -36,7 +36,7 @@ def queue_to_list(q: queue.Queue) -> List[Any]: def del_key(mapping: MutableMapping[T, Any], key: T): """ - Deletes key from mapping. Unlike the `del` keyword, this function does not raise a KeyError + Delete a key from mapping. Unlike the `del` keyword, this function does not raise a KeyError if the key does not exist. :param mapping: A mapping from which a key will be deleted From 5211045194a3eeeaf06b3dd353198454661e1aef Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 7 Jul 2022 11:30:05 -0400 Subject: [PATCH 32/32] Common: Reformat docstring for del_key() --- monkey/common/utils/code_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/common/utils/code_utils.py b/monkey/common/utils/code_utils.py index a7625e0c6..21c0ce175 100644 --- a/monkey/common/utils/code_utils.py +++ b/monkey/common/utils/code_utils.py @@ -36,7 +36,9 @@ def queue_to_list(q: queue.Queue) -> List[Any]: def del_key(mapping: MutableMapping[T, Any], key: T): """ - Delete a key from mapping. Unlike the `del` keyword, this function does not raise a KeyError + Delete a key from a mapping. + + Unlike the `del` keyword, this function does not raise a KeyError if the key does not exist. :param mapping: A mapping from which a key will be deleted