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