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):
USERNAME = 1
PASSWORD = 2
NT_HASH = 3
LM_HASH = 4
SSH_KEYPAIR = 5
USERNAME = "username"
PASSWORD = "password"
NT_HASH = "nt_hash"
LM_HASH = "lm_hash"
SSH_KEYPAIR = "ssh_keypair"

View File

@ -6,5 +6,5 @@ from infection_monkey.i_puppet import ICredentialComponent
@dataclass(frozen=True)
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

View File

@ -6,5 +6,5 @@ from infection_monkey.i_puppet import ICredentialComponent
@dataclass(frozen=True)
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

View File

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

View File

@ -6,6 +6,6 @@ from infection_monkey.i_puppet import ICredentialComponent
@dataclass(frozen=True)
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
public_key: str

View File

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

View File

@ -11,14 +11,14 @@ from .secrets.ssh_key_processor import process_ssh_key
logger = logging.getLogger(__name__)
SECRET_PROCESSORS = {
CredentialsType.PASSWORD: process_password,
CredentialsType.NT_HASH: process_nt_hash,
CredentialsType.LM_HASH: process_lm_hash,
CredentialsType.SSH_KEYPAIR: process_ssh_key,
CredentialsType.PASSWORD.value: process_password,
CredentialsType.NT_HASH.value: process_nt_hash,
CredentialsType.LM_HASH.value: process_lm_hash,
CredentialsType.SSH_KEYPAIR.value: process_ssh_key,
}
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"]:
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:
for identity in credential["identities"]:
IDENTITY_PROCESSORS[identity["credential_type"]](identity)
@ -39,6 +39,6 @@ def is_ssh_keypair(credential: dict) -> bool:
[
secret
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"]:
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")
if not ssh_key["public_key"] or not ssh_key["private_key"]: