p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Common: Clean up credentials.py

This commit is contained in:
Mike Salvatore 2022-07-06 20:48:54 -04:00
parent e4d38631b8
commit c58d26a5e4
1 changed files with 9 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
monkey/common/credentials/credentials.py
View File

@ -1,10 +1,9 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass 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 import Schema, fields, post_load, pre_dump
from marshmallow_enum import EnumField
from . import CredentialComponentType, LMHash, NTHash, Password, SSHKeypair, Username from . import CredentialComponentType, LMHash, NTHash, Password, SSHKeypair, Username
from .i_credential_component import ICredentialComponent from .i_credential_component import ICredentialComponent
@ -14,14 +13,14 @@ from .password import PasswordSchema
from .ssh_keypair import SSHKeypairSchema from .ssh_keypair import SSHKeypairSchema
from .username import UsernameSchema from .username import UsernameSchema
CREDENTIL_COMPINENT_TYPE_TO_CLASS = { CREDENTIAL_COMPINENT_TYPE_TO_CLASS = {
CredentialComponentType.LM_HASH: LMHash, CredentialComponentType.LM_HASH: LMHash,
CredentialComponentType.NT_HASH: NTHash, CredentialComponentType.NT_HASH: NTHash,
CredentialComponentType.PASSWORD: Password, CredentialComponentType.PASSWORD: Password,
CredentialComponentType.SSH_KEYPAIR: SSHKeypair, CredentialComponentType.SSH_KEYPAIR: SSHKeypair,
CredentialComponentType.USERNAME: Username, CredentialComponentType.USERNAME: Username,
} }
CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA = { CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA = {
CredentialComponentType.LM_HASH: LMHashSchema(), CredentialComponentType.LM_HASH: LMHashSchema(),
CredentialComponentType.NT_HASH: NTHashSchema(), CredentialComponentType.NT_HASH: NTHashSchema(),
CredentialComponentType.PASSWORD: PasswordSchema(), 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): class CredentialsSchema(Schema):
# Use fields.List instead of fields.Tuple because marshmallow requires fields.Tuple to have a # Use fields.List instead of fields.Tuple because marshmallow requires fields.Tuple to have a
# fixed length. # fixed length.
# identities = fields.List(fields.Nested(GenericCredentialComponentSchema))
# secrets = fields.List(fields.Nested(GenericCredentialComponentSchema))
identities = fields.List(fields.Mapping()) identities = fields.List(fields.Mapping())
secrets = fields.List(fields.Mapping()) secrets = fields.List(fields.Mapping())
@post_load @post_load
def _make_credentials(self, data, **kwargs) -> Mapping[str, Tuple(Mapping)]: def _make_credentials(self, data, **kwargs) -> Mapping[str, Tuple[Mapping[str, Any]]]:
from pprint import pprint
for component in data["identities"]:
pprint(component)
data["identities"] = tuple( data["identities"] = tuple(
[ [
CredentialsSchema._build_credential_component(component) CredentialsSchema._build_credential_component(component)
@ -73,10 +53,10 @@ class CredentialsSchema(Schema):
return data return data
@staticmethod @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_type = CredentialComponentType[data["credential_type"]]
credential_component_class = CREDENTIL_COMPINENT_TYPE_TO_CLASS[credential_component_type] credential_component_class = CREDENTIAL_COMPINENT_TYPE_TO_CLASS[credential_component_type]
credential_component_schema = CREDENTIL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ credential_component_schema = CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA[
credential_component_type credential_component_type
] ]
@ -105,7 +85,7 @@ class CredentialsSchema(Schema):
@staticmethod @staticmethod
def _serialize_credential_component(credential_component: ICredentialComponent): 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 credential_component.credential_type
] ]
return credential_component_schema.dump(credential_component) return credential_component_schema.dump(credential_component)