Common: Replace credentials dataclasses with pydantic dataclasses

For pydantic to work properly, the models need to either use pydantic version of dataclass or inherit BaseModel, but not both. https://pydantic-docs.helpmanual.io/usage/dataclasses/
This commit is contained in:
vakarisz 2022-08-30 13:38:59 +03:00 committed by vakaris_zilius
parent 5972f87391
commit 3ac60988a8
9 changed files with 40 additions and 25 deletions

View File

@ -1,8 +1,11 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from pydantic.dataclasses import dataclass
from . import CredentialComponentType from . import CredentialComponentType
@dataclass
class ICredentialComponent(ABC): class ICredentialComponent(ABC):
@property @property
@abstractmethod @abstractmethod

View File

@ -1,6 +1,8 @@
from dataclasses import dataclass, field from dataclasses import field
from typing import ClassVar
from marshmallow import fields from marshmallow import fields
from pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
@ -12,9 +14,9 @@ class LMHashSchema(CredentialComponentSchema):
lm_hash = fields.Str(validate=ntlm_hash_validator) lm_hash = fields.Str(validate=ntlm_hash_validator)
@dataclass(frozen=True) @dataclass
class LMHash(ICredentialComponent): class LMHash(ICredentialComponent):
credential_type: CredentialComponentType = field( credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.LM_HASH, init=False default=CredentialComponentType.LM_HASH, init=False
) )
lm_hash: str lm_hash: str

View File

@ -1,6 +1,8 @@
from dataclasses import dataclass, field from dataclasses import field
from typing import ClassVar
from marshmallow import fields from marshmallow import fields
from pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
@ -12,9 +14,9 @@ class NTHashSchema(CredentialComponentSchema):
nt_hash = fields.Str(validate=ntlm_hash_validator) nt_hash = fields.Str(validate=ntlm_hash_validator)
@dataclass(frozen=True) @dataclass
class NTHash(ICredentialComponent): class NTHash(ICredentialComponent):
credential_type: CredentialComponentType = field( credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.NT_HASH, init=False default=CredentialComponentType.NT_HASH, init=False
) )
nt_hash: str nt_hash: str

View File

@ -1,6 +1,8 @@
from dataclasses import dataclass, field from dataclasses import field
from typing import ClassVar
from marshmallow import fields from marshmallow import fields
from pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
@ -11,9 +13,9 @@ class PasswordSchema(CredentialComponentSchema):
password = fields.Str() password = fields.Str()
@dataclass(frozen=True) @dataclass
class Password(ICredentialComponent): class Password(ICredentialComponent):
credential_type: CredentialComponentType = field( credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.PASSWORD, init=False default=CredentialComponentType.PASSWORD, init=False
) )
password: str password: str

View File

@ -1,22 +1,26 @@
from dataclasses import dataclass, field from dataclasses import field
from typing import ClassVar
from marshmallow import fields from marshmallow import fields
from pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
class SSHKeypairSchema(CredentialComponentSchema): class SSHKeypairSchema(CredentialComponentSchema):
credential_type = CredentialTypeField(CredentialComponentType.SSH_KEYPAIR) credential_type: ClassVar[CredentialComponentType] = CredentialTypeField(
CredentialComponentType.SSH_KEYPAIR
)
# TODO: Find a list of valid formats for ssh keys and add validators. # TODO: Find a list of valid formats for ssh keys and add validators.
# See https://github.com/nemchik/ssh-key-regex # See https://github.com/nemchik/ssh-key-regex
private_key = fields.Str() private_key = fields.Str()
public_key = fields.Str() public_key = fields.Str()
@dataclass(frozen=True) @dataclass
class SSHKeypair(ICredentialComponent): class SSHKeypair(ICredentialComponent):
credential_type: CredentialComponentType = field( credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.SSH_KEYPAIR, init=False default=CredentialComponentType.SSH_KEYPAIR, init=False
) )
private_key: str private_key: str

View File

@ -1,6 +1,8 @@
from dataclasses import dataclass, field from dataclasses import field
from typing import ClassVar
from marshmallow import fields from marshmallow import fields
from pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
@ -11,9 +13,9 @@ class UsernameSchema(CredentialComponentSchema):
username = fields.Str() username = fields.Str()
@dataclass(frozen=True) @dataclass
class Username(ICredentialComponent): class Username(ICredentialComponent):
credential_type: CredentialComponentType = field( credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.USERNAME, init=False default=CredentialComponentType.USERNAME, init=False
) )
username: str username: str

View File

@ -45,7 +45,7 @@ SECRET_DICTS = [
] ]
CREDENTIALS = [ CREDENTIALS = [
Credentials(identity, secret) Credentials(identity=identity, secret=secret)
for identity, secret in product(IDENTITIES, SECRETS) for identity, secret in product(IDENTITIES, SECRETS)
if not (identity is None and secret is None) if not (identity is None and secret is None)
] ]

View File

@ -60,8 +60,8 @@ STOLEN_CREDENTIALS = [
STOLEN_SSH_KEYS_CREDENTIALS = [ STOLEN_SSH_KEYS_CREDENTIALS = [
Credentials( Credentials(
Username(USERNAME), identity=Username(USERNAME),
SSHKeypair(public_key=STOLEN_PUBLIC_KEY_2, private_key=STOLEN_PRIVATE_KEY_2), secret=SSHKeypair(public_key=STOLEN_PUBLIC_KEY_2, private_key=STOLEN_PRIVATE_KEY_2),
) )
] ]

View File

@ -24,12 +24,12 @@ identities = (fake_username,)
secrets = (fake_nt_hash, fake_lm_hash, fake_password, fake_ssh_key) secrets = (fake_nt_hash, fake_lm_hash, fake_password, fake_ssh_key)
fake_credentials = [ fake_credentials = [
Credentials(fake_username, fake_nt_hash), Credentials(identity=fake_username, secret=fake_nt_hash),
Credentials(fake_username, fake_lm_hash), Credentials(identity=fake_username, secret=fake_lm_hash),
Credentials(fake_username, fake_password), Credentials(identity=fake_username, secret=fake_password),
Credentials(fake_username, fake_ssh_key), Credentials(identity=fake_username, secret=fake_ssh_key),
Credentials(None, fake_ssh_key), Credentials(identity=None, secret=fake_ssh_key),
Credentials(fake_username, None), Credentials(identity=fake_username, secret=None),
] ]