From c58d26a5e47c5ca54cc72329ef3076d9bca789cc Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 6 Jul 2022 20:48:54 -0400 Subject: [PATCH] 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)