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 pydantic.dataclasses import dataclass
from . import CredentialComponentType
@dataclass
class ICredentialComponent(ABC):
@property
@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 pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
@ -12,9 +14,9 @@ class LMHashSchema(CredentialComponentSchema):
lm_hash = fields.Str(validate=ntlm_hash_validator)
@dataclass(frozen=True)
@dataclass
class LMHash(ICredentialComponent):
credential_type: CredentialComponentType = field(
credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.LM_HASH, init=False
)
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 pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
@ -12,9 +14,9 @@ class NTHashSchema(CredentialComponentSchema):
nt_hash = fields.Str(validate=ntlm_hash_validator)
@dataclass(frozen=True)
@dataclass
class NTHash(ICredentialComponent):
credential_type: CredentialComponentType = field(
credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.NT_HASH, init=False
)
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 pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
@ -11,9 +13,9 @@ class PasswordSchema(CredentialComponentSchema):
password = fields.Str()
@dataclass(frozen=True)
@dataclass
class Password(ICredentialComponent):
credential_type: CredentialComponentType = field(
credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.PASSWORD, init=False
)
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 pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
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.
# See https://github.com/nemchik/ssh-key-regex
private_key = fields.Str()
public_key = fields.Str()
@dataclass(frozen=True)
@dataclass
class SSHKeypair(ICredentialComponent):
credential_type: CredentialComponentType = field(
credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.SSH_KEYPAIR, init=False
)
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 pydantic.dataclasses import dataclass
from . import CredentialComponentType, ICredentialComponent
from .credential_component_schema import CredentialComponentSchema, CredentialTypeField
@ -11,9 +13,9 @@ class UsernameSchema(CredentialComponentSchema):
username = fields.Str()
@dataclass(frozen=True)
@dataclass
class Username(ICredentialComponent):
credential_type: CredentialComponentType = field(
credential_type: ClassVar[CredentialComponentType] = field(
default=CredentialComponentType.USERNAME, init=False
)
username: str

View File

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

View File

@ -60,8 +60,8 @@ STOLEN_CREDENTIALS = [
STOLEN_SSH_KEYS_CREDENTIALS = [
Credentials(
Username(USERNAME),
SSHKeypair(public_key=STOLEN_PUBLIC_KEY_2, private_key=STOLEN_PRIVATE_KEY_2),
identity=Username(USERNAME),
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)
fake_credentials = [
Credentials(fake_username, fake_nt_hash),
Credentials(fake_username, fake_lm_hash),
Credentials(fake_username, fake_password),
Credentials(fake_username, fake_ssh_key),
Credentials(None, fake_ssh_key),
Credentials(fake_username, None),
Credentials(identity=fake_username, secret=fake_nt_hash),
Credentials(identity=fake_username, secret=fake_lm_hash),
Credentials(identity=fake_username, secret=fake_password),
Credentials(identity=fake_username, secret=fake_ssh_key),
Credentials(identity=None, secret=fake_ssh_key),
Credentials(identity=fake_username, secret=None),
]