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

UT: Add fixture to test_exploiter to remove code duplication

This commit is contained in:
Mike Salvatore 2022-02-17 12:41:27 -05:00
parent 7551f254fc
commit 2305a9d413
1 changed files with 15 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py
View File

@ -66,12 +66,20 @@ def get_credentials_for_propagation():
return CREDENTIALS_FOR_PROPAGATION return CREDENTIALS_FOR_PROPAGATION
def test_exploiter(exploiter_config, callback, scan_completed, stop, hosts, hosts_to_exploit): @pytest.fixture
# Set this so that Exploiter() exits once it has processed all victims def run_exploiters(exploiter_config, hosts_to_exploit, callback, scan_completed, stop):
scan_completed.set() def inner(puppet, num_workers):
# Set this so that Exploiter() exits once it has processed all victims
scan_completed.set()
e = Exploiter(MockPuppet(), 2, get_credentials_for_propagation) e = Exploiter(puppet, num_workers, get_credentials_for_propagation)
e.exploit_hosts(exploiter_config, hosts_to_exploit, callback, scan_completed, stop) e.exploit_hosts(exploiter_config, hosts_to_exploit, callback, scan_completed, stop)
return inner
def test_exploiter(callback, hosts, hosts_to_exploit, run_exploiters):
run_exploiters(MockPuppet(), 2)
assert callback.call_count == 5 assert callback.call_count == 5
host_exploit_combos = set() host_exploit_combos = set()
@ -88,15 +96,9 @@ def test_exploiter(exploiter_config, callback, scan_completed, stop, hosts, host
assert ("SSHExploiter", hosts[1]) in host_exploit_combos assert ("SSHExploiter", hosts[1]) in host_exploit_combos
def test_credentials_passed_to_exploiter( def test_credentials_passed_to_exploiter(run_exploiters):
exploiter_config, callback, scan_completed, stop, hosts, hosts_to_exploit
):
mock_puppet = MagicMock() mock_puppet = MagicMock()
# Set this so that Exploiter() exits once it has processed all victims run_exploiters(mock_puppet, 1)
scan_completed.set()
e = Exploiter(mock_puppet, 2, get_credentials_for_propagation)
e.exploit_hosts(exploiter_config, hosts_to_exploit, callback, scan_completed, stop)
for call_args in mock_puppet.exploit_host.call_args_list: for call_args in mock_puppet.exploit_host.call_args_list:
assert call_args[0][2].get("credentials") == CREDENTIALS_FOR_PROPAGATION assert call_args[0][2].get("credentials") == CREDENTIALS_FOR_PROPAGATION