From fdaa454c5978caf3ef795b11b6ce8f8c42bc2b78 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 15 Dec 2021 13:12:20 -0500 Subject: [PATCH] Agent: Add unit tests for AutomatedMaster island comms retry --- .../master/automated_master.py | 10 -- .../automated_master_config.json | 112 ++++++++++++++++++ .../unit_tests/infection_monkey/conftest.py | 5 + .../master/test_automated_master.py | 59 +++++++++ 4 files changed, 176 insertions(+), 10 deletions(-) create mode 100644 monkey/tests/data_for_tests/monkey_configs/automated_master_config.json diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 0a2c2841e..8c95d529b 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -87,8 +87,6 @@ class AutomatedMaster(IMaster): while self._master_thread_should_run(): if timer.is_expired(): - # TODO: Handle exceptions in _check_for_stop() once - # ControlChannel.should_agent_stop() is refactored. self._check_for_stop() timer.reset() @@ -164,14 +162,6 @@ class AutomatedMaster(IMaster): pba_thread.join() - # TODO: This code is just for testing in development. Remove when - # implementation of AutomatedMaster is finished. - while True: - time.sleep(2) - logger.debug("Simulation thread is finished sleeping") - if self._stop.is_set(): - break - def _collect_system_info(self, collector: str): system_info_telemetry = {} system_info_telemetry[collector] = self._puppet.run_sys_info_collector(collector) diff --git a/monkey/tests/data_for_tests/monkey_configs/automated_master_config.json b/monkey/tests/data_for_tests/monkey_configs/automated_master_config.json new file mode 100644 index 000000000..4a7816301 --- /dev/null +++ b/monkey/tests/data_for_tests/monkey_configs/automated_master_config.json @@ -0,0 +1,112 @@ +{ + "config": { + "propagation": { + "network_scan": { + "tcp": { + "timeout_ms": 3000, + "ports": [ + 22, + 2222, + 445, + 135, + 3389, + 80, + 8080, + 443, + 8008, + 3306, + 7001, + 8088, + 9200 + ] + }, + "icmp": { + "timeout_ms": 1000 + }, + "fingerprinters": [ + "SMBFinger", + "SSHFinger", + "HTTPFinger", + "MySQLFinger", + "MSSQLFinger", + "ElasticFinger" + ] + }, + "targets": { + "blocked_ips": ["192.168.1.1", "192.168.1.100"], + "inaccessible_subnets": ["10.0.0.0/24", "10.0.10.0/24"], + "local_network_scan": true, + "subnet_scan_list": [ + "192.168.1.50", + "192.168.56.0/24", + "10.0.33.0/30", + "10.0.0.1", + "10.0.0.2" + ] + }, + "exploiters": { + "brute_force": [ + {"name": "MSSQLExploiter", "propagator": true}, + {"name": "PowerShellExploiter", "propagator": true}, + {"name": "SmbExploiter", "propagator": true}, + {"name": "SSHExploiter", "propagator": true}, + {"name": "WmiExploiter", "propagator": true} + ], + "vulnerability": [ + {"name": "DrupalExploiter", "propagator": true}, + {"name": "ElasticGroovyExploiter", "propagator": true}, + {"name": "HadoopExploiter", "propagator": true}, + {"name": "ShellShockExploiter", "propagator": true}, + {"name": "Struts2Exploiter", "propagator": true}, + {"name": "WebLogicExploiter", "propagator": true}, + {"name": "ZerologonExploiter", "propagator": false} + ] + } + }, + "PBA_linux_filename": "", + "PBA_windows_filename": "", + "command_servers": ["10.197.94.72:5000"], + "current_server": "localhost:5000", + "custom_PBA_linux_cmd": "", + "custom_PBA_windows_cmd": "", + "depth": 2, + "dropper_set_date": true, + "exploit_lm_hash_list": ["DEADBEEF", "FACADE"], + "exploit_ntlm_hash_list": ["BEADED", "ACCEDE", "DECADE"], + "exploit_password_list": ["p1", "p2", "p3"], + "exploit_ssh_keys": "hidden", + "exploit_user_list": ["u1", "u2", "u3"], + "exploiter_classes": [], + "max_depth": 2, + "post_breach_actions": { + "CommunicateAsBackdoorUser": {}, + "ModifyShellStartupFiles": {}, + "HiddenFiles": {}, + "TrapCommand": {}, + "ChangeSetuidSetgid": {}, + "ScheduleJobs": {}, + "Timestomping": {}, + "AccountDiscovery": {}, + "Custom": { + "linux_command": "chmod u+x my_exec && ./my_exec", + "windows_cmd": "powershell test_driver.ps1", + "linux_filename": "my_exec", + "windows_filename": "test_driver.ps1" + } + }, + "payloads": { + "ransomware": { + "encryption": { + "directories": {"linux_target_dir": "", "windows_target_dir": ""}, + "enabled": true + }, + "other_behaviors": {"readme": true} + } + }, + "system_info_collector_classes": [ + "AwsCollector", + "ProcessListCollector", + "MimikatzCollector" + ] + } +} diff --git a/monkey/tests/unit_tests/infection_monkey/conftest.py b/monkey/tests/unit_tests/infection_monkey/conftest.py index 533572f98..14a193112 100644 --- a/monkey/tests/unit_tests/infection_monkey/conftest.py +++ b/monkey/tests/unit_tests/infection_monkey/conftest.py @@ -15,3 +15,8 @@ class TelemetryMessengerSpy(ITelemetryMessenger): @pytest.fixture def telemetry_messenger_spy(): return TelemetryMessengerSpy() + + +@pytest.fixture +def automated_master_config(load_monkey_config): + return load_monkey_config("automated_master_config.json") diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py b/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py index 0584ca1cd..378eab883 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py @@ -1,4 +1,14 @@ +import time +from unittest.mock import MagicMock + from infection_monkey.master import AutomatedMaster +from infection_monkey.master.automated_master import ( + CHECK_FOR_CONFIG_COUNT, + CHECK_FOR_STOP_AGENT_COUNT, +) +from infection_monkey.master.control_channel import IslandCommunicationError + +INTERVAL = 0.001 def test_terminate_without_start(): @@ -6,3 +16,52 @@ def test_terminate_without_start(): # Test that call to terminate does not raise exception m.terminate() + + +def test_stop_if_cant_get_config_from_island(monkeypatch): + cc = MagicMock() + cc.should_agent_stop = MagicMock(return_value=False) + cc.get_config = MagicMock( + side_effect=IslandCommunicationError("Failed to communicate with island") + ) + + monkeypatch.setattr( + "infection_monkey.master.automated_master.CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC", + INTERVAL, + ) + monkeypatch.setattr( + "infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL + ) + m = AutomatedMaster(None, None, None, cc) + m.start() + + assert cc.get_config.call_count == CHECK_FOR_CONFIG_COUNT + + +# NOTE: This test is a little bit brittle, and probably needs too much knowlegde of the internals +# of AutomatedMaster. For now, it works and it runs quickly. In the future, if we find that +# this test isn't valuable or it starts causing issues, we can just remove it. +def test_stop_if_cant_get_stop_signal_from_island(monkeypatch, automated_master_config): + cc = MagicMock() + cc.should_agent_stop = MagicMock( + side_effect=IslandCommunicationError("Failed to communicate with island") + ) + # Ensure that should_agent_stop times out before get_config() returns to prevent the + # Propagator's sub-threads from hanging + cc.get_config = MagicMock( + return_value=automated_master_config, + side_effect=lambda: time.sleep(INTERVAL * (CHECK_FOR_STOP_AGENT_COUNT + 1)), + ) + + monkeypatch.setattr( + "infection_monkey.master.automated_master.CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC", + INTERVAL, + ) + monkeypatch.setattr( + "infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL + ) + + m = AutomatedMaster(None, None, None, cc) + m.start() + + assert cc.should_agent_stop.call_count == CHECK_FOR_STOP_AGENT_COUNT