Agent: Add unit tests for AutomatedMaster island comms retry

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)

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"
+        ]
+    }
+}

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")

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