UT: Remove arch stuff from PowerShell exploiter tests

This commit is contained in:
Shreya Malviya 2022-03-11 19:28:13 +05:30 committed by Ilija Lazoroski
parent 7d25bf711a
commit 25f90c84bc
1 changed files with 2 additions and 11 deletions

View File

@ -4,7 +4,6 @@ from unittest.mock import MagicMock
import pytest import pytest
from infection_monkey.exploit import powershell from infection_monkey.exploit import powershell
from infection_monkey.exploit.consts import WIN_ARCH_32, WIN_ARCH_64
from infection_monkey.exploit.powershell_utils.auth_options import AuthOptions from infection_monkey.exploit.powershell_utils.auth_options import AuthOptions
from infection_monkey.exploit.powershell_utils.credentials import Credentials from infection_monkey.exploit.powershell_utils.credentials import Credentials
from infection_monkey.model.host import VictimHost from infection_monkey.model.host import VictimHost
@ -115,26 +114,20 @@ def authenticate(mock_client):
return inner return inner
@pytest.mark.parametrize( def test_successful_copy(monkeypatch, powershell_exploiter):
"dropper_target_path,arch",
[(DROPPER_TARGET_PATH_32, WIN_ARCH_32), (DROPPER_TARGET_PATH_64, WIN_ARCH_64)],
)
def test_successful_copy(monkeypatch, powershell_exploiter, dropper_target_path, arch):
mock_client = MagicMock() mock_client = MagicMock()
mock_client.return_value.get_host_architecture = lambda: arch
mock_client.return_value.copy_file = MagicMock(return_value=True) mock_client.return_value.copy_file = MagicMock(return_value=True)
monkeypatch.setattr(powershell, "PowerShellClient", mock_client) monkeypatch.setattr(powershell, "PowerShellClient", mock_client)
success = powershell_exploiter.exploit_host() success = powershell_exploiter.exploit_host()
assert dropper_target_path in mock_client.return_value.copy_file.call_args[0][1] assert DROPPER_TARGET_PATH_64 in mock_client.return_value.copy_file.call_args[0][1]
assert success assert success
def test_failed_copy(monkeypatch, powershell_exploiter): def test_failed_copy(monkeypatch, powershell_exploiter):
mock_client = MagicMock() mock_client = MagicMock()
mock_client.return_value.get_host_architecture = lambda: WIN_ARCH_32
mock_client.return_value.copy_file = MagicMock(return_value=False) mock_client.return_value.copy_file = MagicMock(return_value=False)
monkeypatch.setattr(powershell, "PowerShellClient", mock_client) monkeypatch.setattr(powershell, "PowerShellClient", mock_client)
@ -145,7 +138,6 @@ def test_failed_copy(monkeypatch, powershell_exploiter):
def test_failed_monkey_execution(monkeypatch, powershell_exploiter): def test_failed_monkey_execution(monkeypatch, powershell_exploiter):
mock_client = MagicMock() mock_client = MagicMock()
mock_client.get_host_architecture = lambda: WIN_ARCH_32
mock_client.copy_file = MagicMock(return_value=True) mock_client.copy_file = MagicMock(return_value=True)
mock_client.execute_cmd_as_detached_process = MagicMock(side_effect=Exception) mock_client.execute_cmd_as_detached_process = MagicMock(side_effect=Exception)
@ -158,7 +150,6 @@ def test_failed_monkey_execution(monkeypatch, powershell_exploiter):
def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter): def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter):
mock_client = MagicMock() mock_client = MagicMock()
mock_client.return_value.get_host_architecture = lambda: WIN_ARCH_32
mock_client.return_value.copy_file = MagicMock(return_value=True) mock_client.return_value.copy_file = MagicMock(return_value=True)
# execute_cmd method will throw exceptions for 5 first calls. # execute_cmd method will throw exceptions for 5 first calls.