diff --git a/monkey/common/credentials/ssh_keypair.py b/monkey/common/credentials/ssh_keypair.py index 897acba8a..35f074b13 100644 --- a/monkey/common/credentials/ssh_keypair.py +++ b/monkey/common/credentials/ssh_keypair.py @@ -1,6 +1,16 @@ from dataclasses import dataclass, field +from marshmallow import fields + from . import CredentialComponentType, ICredentialComponent +from .credential_component_schema import CredentialComponentSchema, CredentialTypeField + + +class SSHKeypairSchema(CredentialComponentSchema): + credential_type = CredentialTypeField(CredentialComponentType.SSH_KEYPAIR) + # TODO: Find a list of valid formats for ssh keys and add validators + private_key = fields.Str() + public_key = fields.Str() @dataclass(frozen=True) diff --git a/monkey/tests/unit_tests/common/credentials/test_credential_components.py b/monkey/tests/unit_tests/common/credentials/test_credential_components.py index 6be331356..55a2b9279 100644 --- a/monkey/tests/unit_tests/common/credentials/test_credential_components.py +++ b/monkey/tests/unit_tests/common/credentials/test_credential_components.py @@ -3,10 +3,18 @@ from typing import Any, Mapping import pytest from marshmallow.exceptions import ValidationError -from common.credentials import CredentialComponentType, LMHash, NTHash, Password, Username +from common.credentials import ( + CredentialComponentType, + LMHash, + NTHash, + Password, + SSHKeypair, + Username, +) from common.credentials.lm_hash import LMHashSchema from common.credentials.nt_hash import NTHashSchema from common.credentials.password import PasswordSchema +from common.credentials.ssh_keypair import SSHKeypairSchema from common.credentials.username import UsernameSchema PARAMETRIZED_PARAMETER_NAMES = ( @@ -28,6 +36,12 @@ PARAMETRIZED_PARAMETER_VALUES = [ CredentialComponentType.NT_HASH, {"nt_hash": "E52CAC67419A9A224A3B108F3FA6CB6D"}, ), + ( + SSHKeypair, + SSHKeypairSchema, + CredentialComponentType.SSH_KEYPAIR, + {"public_key": "TEST_PUBLIC_KEY", "private_key": "TEST_PRIVATE_KEY"}, + ), ] @@ -48,6 +62,12 @@ INVALID_COMPONENT_DATA = { {"nt_hash": "0123456789012345678901234568901"}, {"nt_hash": "E52GAC67419A9A224A3B108F3FA6CB6D"}, ), + CredentialComponentType.SSH_KEYPAIR: ( + {"public_key": None, "private_key": "TEST_PRIVATE_KEY"}, + {"public_key": "TEST_PUBLIC_KEY", "private_key": None}, + {"public_key": 1, "private_key": "TEST_PRIVATE_KEY"}, + {"public_key": "TEST_PUBLIC_KEY", "private_key": 999}, + ), }