From cc1c049ee91cb5f8edb374222c9e0e1c08eb8933 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Thu, 9 Sep 2021 11:34:38 +0300
Subject: [PATCH] Refactor test_login_attemps_correctly_reported in
 test_powershell.py to address the changes in the flow of powershell and
 powershell client

---
 .../exploit/test_powershell.py                | 71 +++++--------------
 1 file changed, 17 insertions(+), 54 deletions(-)

diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py
index b9254c1d8..fa24ee7ed 100644
--- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py
+++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py
@@ -157,65 +157,28 @@ def test_failed_monkey_execution(monkeypatch, powershell_exploiter):
 
 def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter):
     mock_client = MagicMock()
-    mock_client.get_host_architecture = lambda: WIN_ARCH_32
-    mock_client.copy_file = MagicMock(return_value=True)
-    mock_client.execute_cmd_as_detached_process = MagicMock(side_effect=Exception)
+    mock_client.return_value.get_host_architecture = lambda: WIN_ARCH_32
+    mock_client.return_value.copy_file = MagicMock(return_value=True)
 
-    def allow_ntlm(_, credentials: Credentials, auth_options: AuthOptions):
-        if credentials.username == USER_LIST[1] and credentials.secret == NT_HASH_LIST[1]:
-            return mock_client
+    # execute_cmd method will throw exceptions for 5 first calls.
+    # 6-th call doesn't throw an exception == credentials successful
+    execute_cmd_returns = [Exception, Exception, Exception, Exception, Exception, True]
+    mock_client.return_value.execute_cmd = MagicMock(side_effect=execute_cmd_returns)
 
-        raise TestAuthenticationError
-
-    mock_powershell_client = MagicMock(side_effect=allow_ntlm)
-    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_client)
 
     powershell_exploiter.exploit_host()
 
-    assert {
-        "result": False,
-        "user": USER_LIST[1],
-        "password": None,
-        "lm_hash": "",
-        "ntlm_hash": "",
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
-
-    assert {
-        "result": False,
-        "user": USER_LIST[1],
-        "password": PASSWORD_LIST[0],
-        "lm_hash": "",
-        "ntlm_hash": "",
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
-
-    assert {
-        "result": False,
-        "user": USER_LIST[0],
-        "password": "",
-        "lm_hash": LM_HASH_LIST[0],
-        "ntlm_hash": "",
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
-
-    assert {
-        "result": False,
-        "user": USER_LIST[1],
-        "password": "",
-        "lm_hash": "",
-        "ntlm_hash": NT_HASH_LIST[0],
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
-
-    assert {
-        "result": True,
-        "user": USER_LIST[1],
-        "password": "",
-        "lm_hash": "",
-        "ntlm_hash": NT_HASH_LIST[1],
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
+    # Total 6 attempts reported, 5 failed and 1 succeeded
+    assert len(powershell_exploiter.exploit_attempts) == len(execute_cmd_returns)
+    assert (
+        len([attempt for attempt in powershell_exploiter.exploit_attempts if not attempt["result"]])
+        == 5
+    )
+    assert (
+        len([attempt for attempt in powershell_exploiter.exploit_attempts if attempt["result"]])
+        == 1
+    )
 
 
 def test_build_monkey_execution_command():