Agent, Island, Common: change code to process CredentialType value

Island: rename credentials_type.py
This commit is contained in:
vakaris_zilius 2022-02-21 12:40:11 +00:00 committed by vakarisz
parent bb760c7e8a
commit 4b3750076a
8 changed files with 18 additions and 18 deletions

View File

@ -2,8 +2,8 @@ from enum import Enum
class CredentialsType(Enum): class CredentialsType(Enum):
USERNAME = 1 USERNAME = "username"
PASSWORD = 2 PASSWORD = "password"
NT_HASH = 3 NT_HASH = "nt_hash"
LM_HASH = 4 LM_HASH = "lm_hash"
SSH_KEYPAIR = 5 SSH_KEYPAIR = "ssh_keypair"

View File

@ -6,5 +6,5 @@ from infection_monkey.i_puppet import ICredentialComponent
@dataclass(frozen=True) @dataclass(frozen=True)
class LMHash(ICredentialComponent): class LMHash(ICredentialComponent):
credential_type: CredentialsType = field(default=CredentialsType.LM_HASH, init=False) credential_type: CredentialsType = field(default=CredentialsType.LM_HASH.value, init=False)
lm_hash: str lm_hash: str

View File

@ -6,5 +6,5 @@ from infection_monkey.i_puppet import ICredentialComponent
@dataclass(frozen=True) @dataclass(frozen=True)
class NTHash(ICredentialComponent): class NTHash(ICredentialComponent):
credential_type: CredentialsType = field(default=CredentialsType.NT_HASH, init=False) credential_type: CredentialsType = field(default=CredentialsType.NT_HASH.value, init=False)
nt_hash: str nt_hash: str

View File

@ -6,5 +6,5 @@ from infection_monkey.i_puppet import ICredentialComponent
@dataclass(frozen=True) @dataclass(frozen=True)
class Password(ICredentialComponent): class Password(ICredentialComponent):
credential_type: CredentialsType = field(default=CredentialsType.PASSWORD, init=False) credential_type: CredentialsType = field(default=CredentialsType.PASSWORD.value, init=False)
password: str password: str

View File

@ -6,6 +6,6 @@ from infection_monkey.i_puppet import ICredentialComponent
@dataclass(frozen=True) @dataclass(frozen=True)
class SSHKeypair(ICredentialComponent): class SSHKeypair(ICredentialComponent):
credential_type: CredentialsType = field(default=CredentialsType.SSH_KEYPAIR, init=False) credential_type: CredentialsType = field(default=CredentialsType.SSH_KEYPAIR.value, init=False)
private_key: str private_key: str
public_key: str public_key: str

View File

@ -6,5 +6,5 @@ from infection_monkey.i_puppet import ICredentialComponent
@dataclass(frozen=True) @dataclass(frozen=True)
class Username(ICredentialComponent): class Username(ICredentialComponent):
credential_type: CredentialsType = field(default=CredentialsType.USERNAME, init=False) credential_type: CredentialsType = field(default=CredentialsType.USERNAME.value, init=False)
username: str username: str

View File

@ -11,14 +11,14 @@ from .secrets.ssh_key_processor import process_ssh_key
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
SECRET_PROCESSORS = { SECRET_PROCESSORS = {
CredentialsType.PASSWORD: process_password, CredentialsType.PASSWORD.value: process_password,
CredentialsType.NT_HASH: process_nt_hash, CredentialsType.NT_HASH.value: process_nt_hash,
CredentialsType.LM_HASH: process_lm_hash, CredentialsType.LM_HASH.value: process_lm_hash,
CredentialsType.SSH_KEYPAIR: process_ssh_key, CredentialsType.SSH_KEYPAIR.value: process_ssh_key,
} }
IDENTITY_PROCESSORS = { IDENTITY_PROCESSORS = {
CredentialsType.USERNAME: process_username, CredentialsType.USERNAME.value: process_username,
} }
@ -26,7 +26,7 @@ def parse_credentials(credentials: dict):
for credential in credentials["data"]: for credential in credentials["data"]:
if is_ssh_keypair(credential): if is_ssh_keypair(credential):
SECRET_PROCESSORS[CredentialsType.SSH_KEYPAIR](credential, credentials["monkey_guid"]) SECRET_PROCESSORS[CredentialsType.SSH_KEYPAIR.value](credential, credentials["monkey_guid"])
else: else:
for identity in credential["identities"]: for identity in credential["identities"]:
IDENTITY_PROCESSORS[identity["credential_type"]](identity) IDENTITY_PROCESSORS[identity["credential_type"]](identity)
@ -39,6 +39,6 @@ def is_ssh_keypair(credential: dict) -> bool:
[ [
secret secret
for secret in credential["secrets"] for secret in credential["secrets"]
if secret["credential_type"] == CredentialsType.SSH_KEYPAIR if secret["credential_type"] == CredentialsType.SSH_KEYPAIR.value
] ]
) )

View File

@ -17,7 +17,7 @@ def process_ssh_key(credentials: dict, monkey_guid: str):
) )
for ssh_key in credentials["secrets"]: for ssh_key in credentials["secrets"]:
if not ssh_key["credential_type"] == CredentialsType.SSH_KEYPAIR.name: if not ssh_key["credential_type"] == CredentialsType.SSH_KEYPAIR.value:
raise SSHKeyProcessingError("SSH credentials contain secrets that are not keypairs") raise SSHKeyProcessingError("SSH credentials contain secrets that are not keypairs")
if not ssh_key["public_key"] or not ssh_key["private_key"]: if not ssh_key["public_key"] or not ssh_key["private_key"]: