From 6e254c71d7ef594a0e66c67be8cdc4a14d72143c Mon Sep 17 00:00:00 2001
From: vakaris_zilius <vakarisz@yahoo.com>
Date: Wed, 7 Sep 2022 08:15:12 +0000
Subject: [PATCH] Common: Fix failing UT and add more for get_plain_text

---
 monkey/common/credentials/credentials.py      |  2 +-
 .../data_for_tests/propagation_credentials.py | 12 +++++---
 .../common/credentials/test_credentials.py    | 30 ++++++++++++++++++-
 3 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/monkey/common/credentials/credentials.py b/monkey/common/credentials/credentials.py
index ae8de63d1..d924a2562 100644
--- a/monkey/common/credentials/credentials.py
+++ b/monkey/common/credentials/credentials.py
@@ -12,7 +12,7 @@ Identity = Username
 
 
 def get_plain_text(secret: Union[SecretStr, SecretBytes, None, str]) -> Optional[str]:
-    if secret:
+    if isinstance(secret, (SecretStr, SecretBytes)):
         return secret.get_secret_value()
     else:
         return secret
diff --git a/monkey/tests/data_for_tests/propagation_credentials.py b/monkey/tests/data_for_tests/propagation_credentials.py
index a23334e14..35521e7b6 100644
--- a/monkey/tests/data_for_tests/propagation_credentials.py
+++ b/monkey/tests/data_for_tests/propagation_credentials.py
@@ -6,13 +6,17 @@ from common.credentials import Credentials, LMHash, NTHash, Password, SSHKeypair
 
 USERNAME = "m0nk3y_user"
 SPECIAL_USERNAME = "m0nk3y.user"
-NT_HASH = SecretStr("C1C58F96CDF212B50837BC11A00BE47C")
-LM_HASH = SecretStr("299BD128C1101FD6299BD128C1101FD6")
-PASSWORD_1 = SecretStr("trytostealthis")
+PLAINTEXT_NT_HASH = "C1C58F96CDF212B50837BC11A00BE47C"
+PLAINTEXT_LM_HASH = "299BD128C1101FD6299BD128C1101FD6"
+PLAINTEXT_PASSWORD = "trytostealthis"
+PLAINTEXT_PRIVATE_KEY = "MY_PRIVATE_KEY"
+NT_HASH = SecretStr(PLAINTEXT_NT_HASH)
+LM_HASH = SecretStr(PLAINTEXT_LM_HASH)
+PASSWORD_1 = SecretStr(PLAINTEXT_PASSWORD)
 PASSWORD_2 = SecretStr("password!")
 PASSWORD_3 = SecretStr("rubberbabybuggybumpers")
 PUBLIC_KEY = "MY_PUBLIC_KEY"
-PRIVATE_KEY = SecretStr("MY_PRIVATE_KEY")
+PRIVATE_KEY = SecretStr(PLAINTEXT_PRIVATE_KEY)
 
 IDENTITIES = [Username(username=USERNAME), None, Username(username=SPECIAL_USERNAME)]
 IDENTITY_DICTS = [{"username": USERNAME}, None]
diff --git a/monkey/tests/unit_tests/common/credentials/test_credentials.py b/monkey/tests/unit_tests/common/credentials/test_credentials.py
index 30f27bbdb..79b727a15 100644
--- a/monkey/tests/unit_tests/common/credentials/test_credentials.py
+++ b/monkey/tests/unit_tests/common/credentials/test_credentials.py
@@ -1,12 +1,23 @@
 import logging
+from pathlib import Path
 
 import pytest
 from pydantic import SecretBytes
 from pydantic.types import SecretStr
-from tests.data_for_tests.propagation_credentials import CREDENTIALS, CREDENTIALS_DICTS
+from tests.data_for_tests.propagation_credentials import (
+    CREDENTIALS,
+    CREDENTIALS_DICTS,
+    LM_HASH,
+    PASSWORD_1,
+    PLAINTEXT_LM_HASH,
+    PLAINTEXT_PASSWORD,
+    PLAINTEXT_PRIVATE_KEY,
+    PRIVATE_KEY,
+)
 
 from common.base_models import InfectionMonkeyBaseModel
 from common.credentials import Credentials
+from common.credentials.credentials import get_plain_text
 
 
 @pytest.mark.parametrize(
@@ -42,3 +53,20 @@ def test_credentials_secrets_not_logged(caplog):
     )
 
     assert sensitive not in caplog.text
+
+
+_plaintext = [
+    PLAINTEXT_PASSWORD,
+    PLAINTEXT_PRIVATE_KEY,
+    PLAINTEXT_LM_HASH,
+    "",
+    "already_plaintext",
+    Path("C:\\jolly_fella"),
+    None,
+]
+_hidden = [PASSWORD_1, PRIVATE_KEY, LM_HASH, "", "already_plaintext", Path("C:\\jolly_fella"), None]
+
+
+@pytest.mark.parametrize("expected, hidden", list(zip(_plaintext, _hidden)))
+def test_get_plain_text(expected, hidden):
+    assert expected == get_plain_text(hidden)