Common: Add validation to LMHash and NTHash

This commit is contained in:
Mike Salvatore 2022-07-06 16:00:18 -04:00
parent 0d477cef7c
commit 92416cb079
4 changed files with 68 additions and 3 deletions

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)