p15693087
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests 1 Projects Releases Wiki Activity

UT: Add mock_agent_event_queue in zerologon tests

This commit is contained in:
Ilija Lazoroski 2022-10-05 17:04:30 +02:00
parent 9a261296e4
commit 1db5e9adc5
1 changed files with 13 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
monkey/tests/unit_tests/infection_monkey/exploit/test_zerologon.py
View File

@ -16,7 +16,12 @@ NT_HASHES = ["def456", "765vut"]
@pytest.fixture @pytest.fixture
def zerologon_exploiter_object(monkeypatch): def mock_agent_event_queue():
return MagicMock(spec=IAgentEventQueue)
@pytest.fixture
def zerologon_exploiter_object(monkeypatch, mock_agent_event_queue):
from infection_monkey.exploit.zerologon import ZerologonExploiter from infection_monkey.exploit.zerologon import ZerologonExploiter
def mock_report_login_attempt(**kwargs): def mock_report_login_attempt(**kwargs):
@ -26,22 +31,26 @@ def zerologon_exploiter_object(monkeypatch):
monkeypatch.setattr(obj, "dc_name", NETBIOS_NAME, raising=False) monkeypatch.setattr(obj, "dc_name", NETBIOS_NAME, raising=False)
monkeypatch.setattr(obj, "report_login_attempt", mock_report_login_attempt) monkeypatch.setattr(obj, "report_login_attempt", mock_report_login_attempt)
monkeypatch.setattr(obj, "host", VictimHost(ip_addr="1.1.1.1")) monkeypatch.setattr(obj, "host", VictimHost(ip_addr="1.1.1.1"))
monkeypatch.setattr(obj, "agent_event_queue", MagicMock(spec=IAgentEventQueue)) monkeypatch.setattr(obj, "agent_event_queue", mock_agent_event_queue)
return obj return obj
@pytest.mark.slow @pytest.mark.slow
def test_assess_exploit_attempt_result_no_error(zerologon_exploiter_object): def test_assess_exploit_attempt_result_no_error(zerologon_exploiter_object, mock_agent_event_queue):
dummy_exploit_attempt_result = {"ErrorCode": 0} dummy_exploit_attempt_result = {"ErrorCode": 0}
assert zerologon_exploiter_object.assess_exploit_attempt_result(dummy_exploit_attempt_result) assert zerologon_exploiter_object.assess_exploit_attempt_result(dummy_exploit_attempt_result)
assert mock_agent_event_queue.publish.call_count == 1
@pytest.mark.slow @pytest.mark.slow
def test_assess_exploit_attempt_result_with_error(zerologon_exploiter_object): def test_assess_exploit_attempt_result_with_error(
zerologon_exploiter_object, mock_agent_event_queue
):
dummy_exploit_attempt_result = {"ErrorCode": 1} dummy_exploit_attempt_result = {"ErrorCode": 1}
assert not zerologon_exploiter_object.assess_exploit_attempt_result( assert not zerologon_exploiter_object.assess_exploit_attempt_result(
dummy_exploit_attempt_result dummy_exploit_attempt_result
) )
assert mock_agent_event_queue.publish.call_count == 1
@pytest.mark.slow @pytest.mark.slow