Agent: Skip exploiter if victim OS is not supported

monkey/infection_monkey/master/exploiter.py

@@ -114,6 +114,16 @@ class Exploiter:
 
         for exploiter in interruptible_iter(exploiters_to_run, stop):
             exploiter_name = exploiter["name"]
+            victim_os = victim_host.os.get("type")
+
+            # We want to try all exploiters if the victim's OS is unknown
+            if victim_os is not None and victim_os not in exploiter["supported_os"]:
+                logger.debug(
+                    f"Skipping {exploiter_name} because it does not support "
+                    f"the victim's OS ({victim_os})"
+                )
+                continue
+
             exploiter_results = self._run_exploiter(
                 exploiter_name, exploiter["options"], victim_host, current_depth, stop
             )

monkey/infection_monkey/puppet/mock_puppet.py

@@ -163,8 +163,8 @@ class MockPuppet(IPuppet):
                 "ssh_key": host,
             },
         ]
-        info_powershell = {
+        info_wmi = {
-            "display_name": "PowerShell",
+            "display_name": "WMI",
             "started": "2021-11-25T15:57:06.307696",
             "finished": "2021-11-25T15:58:33.788238",
             "vulnerable_urls": [],
@@ -189,15 +189,15 @@ class MockPuppet(IPuppet):
 
         successful_exploiters = {
             DOT_1: {
-                "PowerShellExploiter": ExploiterResultData(
-                    True, True, False, os_windows, info_powershell, attempts, None
-                ),
                 "ZerologonExploiter": ExploiterResultData(
                     False, False, False, os_windows, {}, [], "Zerologon failed"
                 ),
                 "SSHExploiter": ExploiterResultData(
                     False, False, False, os_linux, info_ssh, attempts, "Failed exploiting"
                 ),
+                "WmiExploiter": ExploiterResultData(
+                    True, True, False, os_windows, info_wmi, attempts, None
+                ),
             },
             DOT_3: {
                 "PowerShellExploiter": ExploiterResultData(
@@ -205,7 +205,7 @@ class MockPuppet(IPuppet):
                     False,
                     False,
                     os_windows,
-                    info_powershell,
+                    info_wmi,
                     attempts,
                     "PowerShell Exploiter Failed",
                 ),

monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py

@@ -1,6 +1,7 @@
 import logging
 from queue import Queue
 from threading import Barrier, Event
+from typing import Iterable
 from unittest.mock import MagicMock
 
 import pytest
@@ -37,29 +38,47 @@ def exploiter_config():
     return {
         "options": {"dropper_path_linux": "/tmp/monkey"},
         "brute_force": [
-            {"name": "PowerShellExploiter", "options": {"timeout": 10}},
+            {"name": "HadoopExploiter", "supported_os": ["windows"], "options": {"timeout": 10}},
-            {"name": "SSHExploiter", "options": {}},
+            {"name": "SSHExploiter", "supported_os": ["linux"], "options": {}},
+            {"name": "WmiExploiter", "supported_os": ["windows"], "options": {"timeout": 10}},
         ],
         "vulnerability": [
-            {"name": "ZerologonExploiter", "options": {}},
+            {"name": "ZerologonExploiter", "supported_os": ["windows"], "options": {}},
         ],
     }
 
 
 @pytest.fixture
 def hosts():
-    return [VictimHost("10.0.0.1"), VictimHost("10.0.0.3")]
+    host_1 = VictimHost("10.0.0.1")
+    host_2 = VictimHost("10.0.0.3")
+    return [host_1, host_2]
 
 
 @pytest.fixture
 def hosts_to_exploit(hosts):
+    return enqueue_hosts(hosts)
+
+
+def enqueue_hosts(hosts: Iterable[VictimHost]):
     q = Queue()
-    q.put(hosts[0])
+    for h in hosts:
-    q.put(hosts[1])
+        q.put(h)
 
     return q
 
 
+def get_host_exploit_combos_from_call_args_list(call_args_list):
+    host_exploit_combos = set()
+
+    for call_args in call_args_list:
+        victim_host = call_args[0][0]
+        exploiter_name = call_args[0][1]
+        host_exploit_combos.add((victim_host, exploiter_name))
+
+    return host_exploit_combos
+
+
 CREDENTIALS_FOR_PROPAGATION = {"usernames": ["m0nk3y", "user"], "passwords": ["1234", "pword"]}
 
 
@@ -69,12 +88,12 @@ def get_credentials_for_propagation():
 
 @pytest.fixture
 def run_exploiters(exploiter_config, hosts_to_exploit, callback, scan_completed, stop):
-    def inner(puppet, num_workers):
+    def inner(puppet, num_workers, hosts=hosts_to_exploit):
         # Set this so that Exploiter() exits once it has processed all victims
         scan_completed.set()
 
         e = Exploiter(puppet, num_workers, get_credentials_for_propagation)
-        e.exploit_hosts(exploiter_config, hosts_to_exploit, 1, callback, scan_completed, stop)
+        e.exploit_hosts(exploiter_config, hosts, 1, callback, scan_completed, stop)
 
     return inner
 
@@ -82,18 +101,16 @@ def run_exploiters(exploiter_config, hosts_to_exploit, callback, scan_completed,
 def test_exploiter(callback, hosts, hosts_to_exploit, run_exploiters):
     run_exploiters(MockPuppet(), 2)
 
-    assert callback.call_count == 5
+    assert callback.call_count == 8
-    host_exploit_combos = set()
+    host_exploit_combos = get_host_exploit_combos_from_call_args_list(callback.call_args_list)
-
-    for i in range(0, 5):
-        victim_host = callback.call_args_list[i][0][0]
-        exploiter_name = callback.call_args_list[i][0][1]
-        host_exploit_combos.add((victim_host, exploiter_name))
 
     assert ("ZerologonExploiter", hosts[0]) in host_exploit_combos
-    assert ("PowerShellExploiter", hosts[0]) in host_exploit_combos
+    assert ("HadoopExploiter", hosts[0]) in host_exploit_combos
+    assert ("SSHExploiter", hosts[0]) in host_exploit_combos
+    assert ("WmiExploiter", hosts[0]) in host_exploit_combos
     assert ("ZerologonExploiter", hosts[1]) in host_exploit_combos
-    assert ("PowerShellExploiter", hosts[1]) in host_exploit_combos
+    assert ("HadoopExploiter", hosts[1]) in host_exploit_combos
+    assert ("WmiExploiter", hosts[1]) in host_exploit_combos
     assert ("SSHExploiter", hosts[1]) in host_exploit_combos
 
 
@@ -132,10 +149,53 @@ def test_exploiter_raises_exception(callback, hosts, hosts_to_exploit, run_explo
     mock_puppet.exploit_host = MagicMock(side_effect=Exception(error_message))
     run_exploiters(mock_puppet, 3)
 
-    assert callback.call_count == 6
+    assert callback.call_count == 8
 
     for i in range(0, 6):
         exploit_result_data = callback.call_args_list[i][0][2]
         assert exploit_result_data.exploitation_success is False
         assert exploit_result_data.propagation_success is False
         assert error_message in exploit_result_data.error_message
+
+
+def test_windows_exploiters_run_on_windows_host(callback, hosts, hosts_to_exploit, run_exploiters):
+    host = VictimHost("10.0.0.1")
+    host.os["type"] = "windows"
+    q = enqueue_hosts([host])
+    run_exploiters(MockPuppet(), 1, q)
+
+    assert callback.call_count == 3
+    host_exploit_combos = get_host_exploit_combos_from_call_args_list(callback.call_args_list)
+
+    assert ("SSHExploiter", host) not in host_exploit_combos
+
+
+def test_linux_exploiters_run_on_linux_host(callback, hosts, hosts_to_exploit, run_exploiters):
+    host = VictimHost("10.0.0.1")
+    host.os["type"] = "linux"
+    q = enqueue_hosts([host])
+    run_exploiters(MockPuppet(), 1, q)
+
+    assert callback.call_count == 1
+    host_exploit_combos = get_host_exploit_combos_from_call_args_list(callback.call_args_list)
+
+    assert ("SSHExploiter", host) in host_exploit_combos
+
+
+def test_all_exploiters_run_on_unknown_host(callback, hosts, hosts_to_exploit, run_exploiters):
+    host = VictimHost("10.0.0.1")
+    try:
+        del host.os["type"]
+    except KeyError:
+        pass
+
+    q = enqueue_hosts([host])
+    run_exploiters(MockPuppet(), 1, q)
+
+    assert callback.call_count == 4
+    host_exploit_combos = get_host_exploit_combos_from_call_args_list(callback.call_args_list)
+
+    assert ("ZerologonExploiter", hosts[0]) in host_exploit_combos
+    assert ("HadoopExploiter", hosts[0]) in host_exploit_combos
+    assert ("SSHExploiter", host) in host_exploit_combos
+    assert ("WmiExploiter", hosts[0]) in host_exploit_combos