diff --git a/monkey/infection_monkey/credential_collectors/credential_components/lm_hash.py b/monkey/infection_monkey/credential_collectors/credential_components/lm_hash.py index 603422dca..4236a5247 100644 --- a/monkey/infection_monkey/credential_collectors/credential_components/lm_hash.py +++ b/monkey/infection_monkey/credential_collectors/credential_components/lm_hash.py @@ -1,9 +1,10 @@ +from dataclasses import dataclass, field + from ..credential_type import CredentialType from .i_credential_component import ICredentialComponent +@dataclass(frozen=True) class LMHash(ICredentialComponent): - type = CredentialType.LM_HASH - - def __init__(self, lm_hash: str): - self.lm_hash = lm_hash + type: CredentialType = field(default=CredentialType.LM_HASH, init=False) + lm_hash: str diff --git a/monkey/infection_monkey/credential_collectors/credential_components/nt_hash.py b/monkey/infection_monkey/credential_collectors/credential_components/nt_hash.py index 9ca0f888d..6f90f37ff 100644 --- a/monkey/infection_monkey/credential_collectors/credential_components/nt_hash.py +++ b/monkey/infection_monkey/credential_collectors/credential_components/nt_hash.py @@ -1,9 +1,10 @@ +from dataclasses import dataclass, field + from ..credential_type import CredentialType from .i_credential_component import ICredentialComponent +@dataclass(frozen=True) class NTHash(ICredentialComponent): - type = CredentialType.NT_HASH - - def __init__(self, nt_hash: str): - self.nt_hash = nt_hash + type: CredentialType = field(default=CredentialType.NT_HASH, init=False) + nt_hash: str diff --git a/monkey/infection_monkey/credential_collectors/credential_components/password.py b/monkey/infection_monkey/credential_collectors/credential_components/password.py index cf546aac6..e3ed03fcf 100644 --- a/monkey/infection_monkey/credential_collectors/credential_components/password.py +++ b/monkey/infection_monkey/credential_collectors/credential_components/password.py @@ -1,9 +1,10 @@ +from dataclasses import dataclass, field + from ..credential_type import CredentialType from .i_credential_component import ICredentialComponent +@dataclass(frozen=True) class Password(ICredentialComponent): - type = CredentialType.PASSWORD - - def __init__(self, password: str): - self.password = password + type: CredentialType = field(default=CredentialType.PASSWORD, init=False) + password: str diff --git a/monkey/infection_monkey/credential_collectors/credential_components/username.py b/monkey/infection_monkey/credential_collectors/credential_components/username.py index 578551489..d822e0e20 100644 --- a/monkey/infection_monkey/credential_collectors/credential_components/username.py +++ b/monkey/infection_monkey/credential_collectors/credential_components/username.py @@ -1,9 +1,10 @@ +from dataclasses import dataclass, field + from ..credential_type import CredentialType from .i_credential_component import ICredentialComponent +@dataclass(frozen=True) class Username(ICredentialComponent): - type = CredentialType.USERNAME - - def __init__(self, username: str): - self.username = username + type: CredentialType = field(default=CredentialType.USERNAME, init=False) + username: str diff --git a/monkey/infection_monkey/credential_collectors/credentials.py b/monkey/infection_monkey/credential_collectors/credentials.py index 1e3d05844..6688e393f 100644 --- a/monkey/infection_monkey/credential_collectors/credentials.py +++ b/monkey/infection_monkey/credential_collectors/credentials.py @@ -1,11 +1,10 @@ -from typing import Iterable +from dataclasses import dataclass +from typing import Tuple from .credential_components.i_credential_component import ICredentialComponent +@dataclass(frozen=True) class Credentials: - def __init__( - self, identities: Iterable[ICredentialComponent], secrets: Iterable[ICredentialComponent] - ): - self.identities = tuple(identities) - self.secrets = tuple(secrets) + identities: Tuple[ICredentialComponent] + secrets: Tuple[ICredentialComponent] diff --git a/monkey/tests/unit_tests/infection_monkey/credential_collectors/mimikatz_collector/test_mimikatz_collector.py b/monkey/tests/unit_tests/infection_monkey/credential_collectors/mimikatz_collector/test_mimikatz_collector.py index 9e45fd4c6..00589eb4b 100644 --- a/monkey/tests/unit_tests/infection_monkey/credential_collectors/mimikatz_collector/test_mimikatz_collector.py +++ b/monkey/tests/unit_tests/infection_monkey/credential_collectors/mimikatz_collector/test_mimikatz_collector.py @@ -1,4 +1,4 @@ -from infection_monkey.credential_collectors import LMHash, NTHash, Password, Username +from infection_monkey.credential_collectors import Credentials, LMHash, NTHash, Password, Username from infection_monkey.credential_collectors.mimikatz_collector.mimikatz_cred_collector import ( MimikatzCredentialCollector, ) @@ -18,27 +18,26 @@ def patch_pypykatz(win_creds: [WindowsCredentials], monkeypatch): def test_empty_results(monkeypatch): win_creds = [WindowsCredentials(username="", password="", ntlm_hash="", lm_hash="")] patch_pypykatz(win_creds, monkeypatch) - expected = [] - collected = MimikatzCredentialCollector().collect_credentials() - assert expected == collected + expected_credentials = [] + collected_credentials = MimikatzCredentialCollector().collect_credentials() + assert expected_credentials == collected_credentials patch_pypykatz([], monkeypatch) - collected = MimikatzCredentialCollector().collect_credentials() - assert [] == collected + collected_credentials = MimikatzCredentialCollector().collect_credentials() + assert not collected_credentials def test_pypykatz_result_parsing(monkeypatch): win_creds = [WindowsCredentials(username="user", password="secret", ntlm_hash="", lm_hash="")] patch_pypykatz(win_creds, monkeypatch) - # Expected credentials username = Username("user") password = Password("secret") + expected_credentials = Credentials([username], [password]) - collected = MimikatzCredentialCollector().collect_credentials() - assert len(list(collected)) == 1 - assert list(collected)[0].identities[0].__dict__ == username.__dict__ - assert list(collected)[0].secrets[0].__dict__ == password.__dict__ + collected_credentials = list(MimikatzCredentialCollector().collect_credentials()) + assert len(collected_credentials) == 1 + assert collected_credentials[0] == expected_credentials def test_pypykatz_result_parsing_duplicates(monkeypatch): @@ -48,8 +47,8 @@ def test_pypykatz_result_parsing_duplicates(monkeypatch): ] patch_pypykatz(win_creds, monkeypatch) - collected = MimikatzCredentialCollector().collect_credentials() - assert len(list(collected)) == 2 + collected_credentials = list(MimikatzCredentialCollector().collect_credentials()) + assert len(collected_credentials) == 2 def test_pypykatz_result_parsing_defaults(monkeypatch): @@ -62,11 +61,11 @@ def test_pypykatz_result_parsing_defaults(monkeypatch): username = Username("user2") password = Password("secret2") lm_hash = LMHash("lm_hash") + expected_credentials = Credentials([username], [password, lm_hash]) - collected = MimikatzCredentialCollector().collect_credentials() - assert list(collected)[0].identities[0].__dict__ == username.__dict__ - assert list(collected)[0].secrets[0].__dict__ == password.__dict__ - assert list(collected)[0].secrets[1].__dict__ == lm_hash.__dict__ + collected_credentials = list(MimikatzCredentialCollector().collect_credentials()) + assert len(collected_credentials) == 1 + assert collected_credentials[0] == expected_credentials def test_pypykatz_result_parsing_no_identities(monkeypatch): @@ -75,10 +74,10 @@ def test_pypykatz_result_parsing_no_identities(monkeypatch): ] patch_pypykatz(win_creds, monkeypatch) - # Expected credentials - nt_hash = NTHash("ntlm_hash") lm_hash = LMHash("lm_hash") + nt_hash = NTHash("ntlm_hash") + expected_credentials = Credentials([], [lm_hash, nt_hash]) - collected = MimikatzCredentialCollector().collect_credentials() - assert list(collected)[0].secrets[0].__dict__ == lm_hash.__dict__ - assert list(collected)[0].secrets[1].__dict__ == nt_hash.__dict__ + collected_credentials = list(MimikatzCredentialCollector().collect_credentials()) + assert len(collected_credentials) == 1 + assert collected_credentials[0] == expected_credentials