From 92416cb079795e0b78600bd7d86607215e457a5d Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Wed, 6 Jul 2022 16:00:18 -0400
Subject: [PATCH] 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)