From 7551f254fc7df05e484c249d0b5a4817a2a2d895 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 17 Feb 2022 12:36:17 -0500 Subject: [PATCH 1/5] Agent: Query for updated credentials in Exploiter Allows exploiters to be run with the most up-to-date configured and stolen credentials from the Island. --- .../master/automated_master.py | 4 ++- monkey/infection_monkey/master/exploiter.py | 25 ++++++++++++++++--- .../master/test_automated_master.py | 2 +- .../infection_monkey/master/test_exploiter.py | 25 +++++++++++++++++-- 4 files changed, 49 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 9fbb5f200..d78d5aafe 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -41,7 +41,9 @@ class AutomatedMaster(IMaster): self._control_channel = control_channel ip_scanner = IPScanner(self._puppet, NUM_SCAN_THREADS) - exploiter = Exploiter(self._puppet, NUM_EXPLOIT_THREADS) + exploiter = Exploiter( + self._puppet, NUM_EXPLOIT_THREADS, self._control_channel.get_credentials_for_propagation + ) self._propagator = Propagator( self._telemetry_messenger, ip_scanner, diff --git a/monkey/infection_monkey/master/exploiter.py b/monkey/infection_monkey/master/exploiter.py index 09f6ebf4b..9d5fe4f00 100644 --- a/monkey/infection_monkey/master/exploiter.py +++ b/monkey/infection_monkey/master/exploiter.py @@ -3,7 +3,7 @@ import queue import threading from queue import Queue from threading import Event -from typing import Callable, Dict, List +from typing import Callable, Dict, List, Mapping from infection_monkey.i_puppet import ExploiterResultData, IPuppet from infection_monkey.model import VictimHost @@ -18,9 +18,15 @@ Callback = Callable[[ExploiterName, VictimHost, ExploiterResultData], None] class Exploiter: - def __init__(self, puppet: IPuppet, num_workers: int): + def __init__( + self, + puppet: IPuppet, + num_workers: int, + get_updated_credentials_for_propagation: Callable[[], Mapping], + ): self._puppet = puppet self._num_workers = num_workers + self._get_updated_credentials_for_propagation = get_updated_credentials_for_propagation def exploit_hosts( self, @@ -74,6 +80,7 @@ class Exploiter: results_callback: Callback, stop: Event, ): + for exploiter in interruptable_iter(exploiters_to_run, stop): exploiter_name = exploiter["name"] exploiter_results = self._run_exploiter(exploiter_name, victim_host, stop) @@ -86,7 +93,19 @@ class Exploiter: self, exploiter_name: str, victim_host: VictimHost, stop: Event ) -> ExploiterResultData: logger.debug(f"Attempting to use {exploiter_name} on {victim_host}") - return self._puppet.exploit_host(exploiter_name, victim_host.ip_addr, {}, stop) + + credentials = self._get_credentials_for_propagation() + options = {"credentials": credentials} + + return self._puppet.exploit_host(exploiter_name, victim_host.ip_addr, options, stop) + + def _get_credentials_for_propagation(self) -> Mapping: + try: + return self._get_updated_credentials_for_propagation() + except Exception as ex: + logger.error(f"Error while attempting to retrieve credentials for propagation: {ex}") + + return {} def _all_hosts_have_been_processed(scan_completed: Event, hosts_to_exploit: Queue): 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 d08a4465a..c7023e525 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 @@ -14,7 +14,7 @@ INTERVAL = 0.001 def test_terminate_without_start(): - m = AutomatedMaster(None, None, None, None, []) + m = AutomatedMaster(None, None, None, MagicMock(), []) # Test that call to terminate does not raise exception m.terminate() diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py b/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py index 5b9297fe6..b2c42f1ec 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py @@ -59,11 +59,18 @@ def hosts_to_exploit(hosts): return q +CREDENTIALS_FOR_PROPAGATION = {"usernames": ["m0nk3y", "user"], "passwords": ["1234", "pword"]} + + +def get_credentials_for_propagation(): + return CREDENTIALS_FOR_PROPAGATION + + def test_exploiter(exploiter_config, callback, scan_completed, stop, hosts, hosts_to_exploit): # Set this so that Exploiter() exits once it has processed all victims scan_completed.set() - e = Exploiter(MockPuppet(), 2) + e = Exploiter(MockPuppet(), 2, get_credentials_for_propagation) e.exploit_hosts(exploiter_config, hosts_to_exploit, callback, scan_completed, stop) assert callback.call_count == 5 @@ -81,6 +88,20 @@ def test_exploiter(exploiter_config, callback, scan_completed, stop, hosts, host assert ("SSHExploiter", hosts[1]) in host_exploit_combos +def test_credentials_passed_to_exploiter( + exploiter_config, callback, scan_completed, stop, hosts, hosts_to_exploit +): + mock_puppet = MagicMock() + # Set this so that Exploiter() exits once it has processed all victims + 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: + assert call_args[0][2].get("credentials") == CREDENTIALS_FOR_PROPAGATION + + def test_stop_after_callback(exploiter_config, callback, scan_completed, stop, hosts_to_exploit): callback_barrier_count = 2 @@ -96,7 +117,7 @@ def test_stop_after_callback(exploiter_config, callback, scan_completed, stop, h # Intentionally NOT setting scan_completed.set(); _callback() will set stop - e = Exploiter(MockPuppet(), callback_barrier_count + 2) + e = Exploiter(MockPuppet(), callback_barrier_count + 2, get_credentials_for_propagation) e.exploit_hosts(exploiter_config, hosts_to_exploit, stoppable_callback, scan_completed, stop) assert stoppable_callback.call_count == 2 From 2305a9d413b10ecac0626cc6c822f2e33e3062b7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 17 Feb 2022 12:41:27 -0500 Subject: [PATCH 2/5] UT: Add fixture to test_exploiter to remove code duplication --- .../infection_monkey/master/test_exploiter.py | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py b/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py index b2c42f1ec..26067ab22 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py @@ -66,12 +66,20 @@ def get_credentials_for_propagation(): return CREDENTIALS_FOR_PROPAGATION -def test_exploiter(exploiter_config, callback, scan_completed, stop, hosts, hosts_to_exploit): - # Set this so that Exploiter() exits once it has processed all victims - scan_completed.set() +@pytest.fixture +def run_exploiters(exploiter_config, hosts_to_exploit, callback, scan_completed, stop): + 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.exploit_hosts(exploiter_config, hosts_to_exploit, callback, scan_completed, stop) + e = Exploiter(puppet, num_workers, get_credentials_for_propagation) + 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 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 -def test_credentials_passed_to_exploiter( - exploiter_config, callback, scan_completed, stop, hosts, hosts_to_exploit -): +def test_credentials_passed_to_exploiter(run_exploiters): mock_puppet = MagicMock() - # Set this so that Exploiter() exits once it has processed all victims - scan_completed.set() - - e = Exploiter(mock_puppet, 2, get_credentials_for_propagation) - e.exploit_hosts(exploiter_config, hosts_to_exploit, callback, scan_completed, stop) + run_exploiters(mock_puppet, 1) for call_args in mock_puppet.exploit_host.call_args_list: assert call_args[0][2].get("credentials") == CREDENTIALS_FOR_PROPAGATION From c3e9690280df25a916a89f0b744dd5fc431bec14 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 17 Feb 2022 14:25:03 -0500 Subject: [PATCH 3/5] Agent: Add request_cache decorator --- monkey/infection_monkey/utils/decorators.py | 43 ++++++++++ .../infection_monkey/utils/test_decorators.py | 78 +++++++++++++++++++ 2 files changed, 121 insertions(+) create mode 100644 monkey/infection_monkey/utils/decorators.py create mode 100644 monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py diff --git a/monkey/infection_monkey/utils/decorators.py b/monkey/infection_monkey/utils/decorators.py new file mode 100644 index 000000000..7a93a7c7a --- /dev/null +++ b/monkey/infection_monkey/utils/decorators.py @@ -0,0 +1,43 @@ +from functools import wraps + +from .timer import Timer + + +def request_cache(ttl: float): + """ + This is a decorator that allows a single response of a function to be cached with an expiration + time (TTL). The first call to the function is executed and the response is cached. Subsequent + calls to the function result in the cached value being returned until the TTL elapses. Once the + TTL elapses, the cache is considered stale and the decorated function will be called, its + response cached, and the TTL reset. + + An example usage of this decorator is to wrap a function that makes frequent slow calls to an + external resource, such as an HTTP request to a remote endpoint. If the most up-to-date + information is not need, this decorator provides a simple way to cache the response for a + certain amount of time. + + Example: + @request_cache(600) + def raining_outside(): + return requests.get(f"https://weather.service.api/check_for_rain/{MY_ZIP_CODE}") + + :param ttl: The time-to-live in seconds for the cached return value + :return: The return value of the decorated function, or the cached return value if the TTL has + not elapsed. + """ + + def decorator(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if wrapper.timer.is_expired(): + wrapper.cached_value = fn(*args, **kwargs) + wrapper.timer.set(ttl) + + return wrapper.cached_value + + wrapper.cached_value = None + wrapper.timer = Timer() + + return wrapper + + return decorator diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py b/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py new file mode 100644 index 000000000..30af5837a --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py @@ -0,0 +1,78 @@ +import time +from unittest.mock import MagicMock + +import pytest + +from infection_monkey.utils.decorators import request_cache +from infection_monkey.utils.timer import Timer + + +class MockTimer(Timer): + def __init__(self): + self._time_remaining = 0 + self._set_time = 0 + + def set(self, timeout_sec: float): + self._time_remaining = timeout_sec + self._set_time = timeout_sec + + def set_expired(self): + self._time_remaining = 0 + + @property + def time_remaining(self) -> float: + return self._time_remaining + + def reset(self): + """ + Reset the timer without changing the timeout + """ + self._time_remaining = self._set_time + + +class MockTimerFactory: + def __init__(self): + self._instance = None + + def __call__(self): + if self._instance is None: + mt = MockTimer() + self._instance = mt + + return self._instance + + def reset(self): + self._instance = None + + +mock_timer_factory = MockTimerFactory() + + +@pytest.fixture +def mock_timer(monkeypatch): + mock_timer_factory.reset + + monkeypatch.setattr("infection_monkey.utils.decorators.Timer", mock_timer_factory) + + return mock_timer_factory() + + +def test_request_cache(mock_timer): + mock_request = MagicMock(side_effect=lambda: time.time()) + + @request_cache(10) + def make_request(): + return mock_request() + + t1 = make_request() + t2 = make_request() + + assert t1 == t2 + + mock_timer.set_expired() + + t3 = make_request() + t4 = make_request() + + assert t3 != t1 + assert t3 == t4 From 4005ea2924e2c254941ad062485488e3f08fb9f1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 17 Feb 2022 14:29:07 -0500 Subject: [PATCH 4/5] Agent: Add caching to ControlChannel.get_credentials_for_propagation() --- monkey/infection_monkey/master/control_channel.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 52b565d55..d6781af7f 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -7,11 +7,14 @@ from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from infection_monkey.config import WormConfiguration from infection_monkey.control import ControlClient from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError +from infection_monkey.utils.decorators import request_cache requests.packages.urllib3.disable_warnings() logger = logging.getLogger(__name__) +CREDENTIALS_POLL_PERIOD_SEC = 30 + class ControlChannel(IControlChannel): def __init__(self, server: str, agent_id: str): @@ -66,18 +69,21 @@ class ControlChannel(IControlChannel): ) as e: raise IslandCommunicationError(e) + @request_cache(CREDENTIALS_POLL_PERIOD_SEC) def get_credentials_for_propagation(self) -> dict: + propagation_credentials_url = ( + f"https://{self._control_channel_server}/api/propagation-credentials/{self._agent_id}" + ) try: response = requests.get( # noqa: DUO123 - f"{self._control_channel_server}/api/propagation-credentials/{self._agent_id}", + propagation_credentials_url, verify=False, proxies=ControlClient.proxies, timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() - response = json.loads(response.content.decode())["propagation_credentials"] - return response + return json.loads(response.content.decode())["propagation_credentials"] except ( json.JSONDecodeError, requests.exceptions.ConnectionError, From e2d116fdf1fc703fab5491da3c9661489778b769 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 17 Feb 2022 14:40:07 -0500 Subject: [PATCH 5/5] Agent: Make request_cache() decorator thread-safe --- monkey/infection_monkey/utils/decorators.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/utils/decorators.py b/monkey/infection_monkey/utils/decorators.py index 7a93a7c7a..31ac0661b 100644 --- a/monkey/infection_monkey/utils/decorators.py +++ b/monkey/infection_monkey/utils/decorators.py @@ -1,3 +1,4 @@ +import threading from functools import wraps from .timer import Timer @@ -29,14 +30,16 @@ def request_cache(ttl: float): def decorator(fn): @wraps(fn) def wrapper(*args, **kwargs): - if wrapper.timer.is_expired(): - wrapper.cached_value = fn(*args, **kwargs) - wrapper.timer.set(ttl) + with wrapper.lock: + if wrapper.timer.is_expired(): + wrapper.cached_value = fn(*args, **kwargs) + wrapper.timer.set(ttl) return wrapper.cached_value wrapper.cached_value = None wrapper.timer = Timer() + wrapper.lock = threading.Lock() return wrapper