p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Common: Fix type hints in credentials.py

This commit is contained in:
Mike Salvatore 2022-07-07 07:43:12 -04:00
parent 59a9aa8a53
commit 008428e318
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,7 +1,7 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Mapping, MutableMapping, Tuple from typing import Any, Mapping, MutableMapping, Sequence, Tuple
from marshmallow import Schema, fields, post_load, pre_dump from marshmallow import Schema, fields, post_load, pre_dump
@ -20,6 +20,7 @@ CREDENTIAL_COMPINENT_TYPE_TO_CLASS = {
CredentialComponentType.SSH_KEYPAIR: SSHKeypair, CredentialComponentType.SSH_KEYPAIR: SSHKeypair,
CredentialComponentType.USERNAME: Username, CredentialComponentType.USERNAME: Username,
} }
CREDENTIAL_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(),
@ -36,7 +37,7 @@ class CredentialsSchema(Schema):
secrets = fields.List(fields.Mapping()) secrets = fields.List(fields.Mapping())
@post_load @post_load
def _make_credentials(self, data, **kwargs) -> Mapping[str, Tuple[Mapping[str, Any]]]: def _make_credentials(self, data, **kwargs) -> Mapping[str, Sequence[Mapping[str, Any]]]:
data["identities"] = tuple( data["identities"] = tuple(
[ [
CredentialsSchema._build_credential_component(component) CredentialsSchema._build_credential_component(component)
@ -66,7 +67,7 @@ class CredentialsSchema(Schema):
@pre_dump @pre_dump
def _serialize_credentials( def _serialize_credentials(
self, credentials: Credentials, **kwargs self, credentials: Credentials, **kwargs
) -> Mapping[str, Tuple[Mapping[str, Any]]]: ) -> Mapping[str, Sequence[Mapping[str, Any]]]:
data = {} data = {}
data["identities"] = tuple( data["identities"] = tuple(
[ [
@ -84,7 +85,9 @@ class CredentialsSchema(Schema):
return data return data
@staticmethod @staticmethod
def _serialize_credential_component(credential_component: ICredentialComponent): def _serialize_credential_component(
credential_component: ICredentialComponent,
) -> Mapping[str, Any]:
credential_component_schema = CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA[ credential_component_schema = CREDENTIAL_COMPINENT_TYPE_TO_CLASS_SCHEMA[
credential_component.credential_type credential_component.credential_type
] ]