From 9286e8690008c62054bda04b533f3bd52a791130 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 22 Jun 2022 17:52:34 +0300 Subject: [PATCH] Agent: Use deserialized in exploiter.py and propagator.py --- monkey/infection_monkey/master/exploiter.py | 34 ++++++++++------- monkey/infection_monkey/master/propagator.py | 13 ++++--- .../infection_monkey/master/test_exploiter.py | 38 +++++++++++++------ 3 files changed, 55 insertions(+), 30 deletions(-) diff --git a/monkey/infection_monkey/master/exploiter.py b/monkey/infection_monkey/master/exploiter.py index 35a212482..46b6799a0 100644 --- a/monkey/infection_monkey/master/exploiter.py +++ b/monkey/infection_monkey/master/exploiter.py @@ -5,9 +5,13 @@ from copy import deepcopy from itertools import chain from queue import Queue from threading import Event -from typing import Callable, Dict, List, Mapping +from typing import Callable, Dict, List from common import OperatingSystems +from common.configuration.agent_sub_configurations import ( + ExploitationConfiguration, + ExploiterConfiguration, +) from infection_monkey.custom_types import PropagationCredentials from infection_monkey.i_puppet import ExploiterResultData, IPuppet from infection_monkey.model import VictimHost @@ -46,7 +50,7 @@ class Exploiter: def exploit_hosts( self, - exploiter_config: Dict, + exploiter_config: ExploitationConfiguration, hosts_to_exploit: Queue, current_depth: int, results_callback: Callback, @@ -56,7 +60,7 @@ class Exploiter: exploiters_to_run = self._process_exploiter_config(exploiter_config) logger.debug( "Agent is configured to run the following exploiters in order: " - f"{', '.join([e['name'] for e in exploiters_to_run])}" + f"{', '.join([e.name for e in exploiters_to_run])}" ) exploit_args = ( @@ -75,24 +79,28 @@ class Exploiter: ) @staticmethod - def _process_exploiter_config(exploiter_config: Mapping) -> List[Mapping]: + def _process_exploiter_config( + exploiter_config: ExploitationConfiguration, + ) -> List[ExploiterConfiguration]: # Run vulnerability exploiters before brute force exploiters to minimize the effect of # account lockout due to invalid credentials - ordered_exploiters = chain( - exploiter_config["vulnerability"], exploiter_config["brute_force"] - ) + ordered_exploiters = chain(exploiter_config.vulnerability, exploiter_config.brute_force) exploiters_to_run = list(deepcopy(ordered_exploiters)) + extended_exploiters = [] for exploiter in exploiters_to_run: # This order allows exploiter-specific options to # override general options for all exploiters. - exploiter["options"] = {**exploiter_config["options"], **exploiter["options"]} + options = {**exploiter_config.options.__dict__, **exploiter.options} + extended_exploiters.append( + ExploiterConfiguration(exploiter.name, options, exploiter.supported_os) + ) - return exploiters_to_run + return extended_exploiters def _exploit_hosts_on_queue( self, - exploiters_to_run: List[Dict], + exploiters_to_run: List[ExploiterConfiguration], hosts_to_exploit: Queue, current_depth: int, results_callback: Callback, @@ -119,7 +127,7 @@ class Exploiter: def _run_all_exploiters( self, - exploiters_to_run: List[Dict], + exploiters_to_run: List[ExploiterConfiguration], victim_host: VictimHost, current_depth: int, results_callback: Callback, @@ -127,7 +135,7 @@ class Exploiter: ): for exploiter in interruptible_iter(exploiters_to_run, stop): - exploiter_name = exploiter["name"] + exploiter_name = exploiter.name victim_os = victim_host.os.get("type") # We want to try all exploiters if the victim's OS is unknown @@ -140,7 +148,7 @@ class Exploiter: continue exploiter_results = self._run_exploiter( - exploiter_name, exploiter["options"], victim_host, current_depth, stop + exploiter_name, exploiter.options, victim_host, current_depth, stop ) results_callback(exploiter_name, victim_host, exploiter_results) diff --git a/monkey/infection_monkey/master/propagator.py b/monkey/infection_monkey/master/propagator.py index a74cf7c86..64edae2ec 100644 --- a/monkey/infection_monkey/master/propagator.py +++ b/monkey/infection_monkey/master/propagator.py @@ -1,10 +1,13 @@ import logging from queue import Queue from threading import Event -from typing import Dict, List +from typing import List -from common.configuration import PropagationConfiguration,\ - NetworkScanConfiguration, ScanTargetConfiguration +from common.configuration import ( + NetworkScanConfiguration, + PropagationConfiguration, + ScanTargetConfiguration, +) from infection_monkey.i_puppet import ( ExploiterResultData, FingerprintData, @@ -139,14 +142,14 @@ class Propagator: def _exploit_hosts( self, - propagation_config: Dict, + propagation_config: PropagationConfiguration, current_depth: int, network_scan_completed: Event, stop: Event, ): logger.info("Exploiting victims") - exploiter_config = propagation_config["exploiters"] + exploiter_config = propagation_config.exploitation self._exploiter.exploit_hosts( exploiter_config, self._hosts_to_exploit, 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 f5d9c2da4..2af0f3809 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py @@ -8,6 +8,10 @@ import pytest from tests.unit_tests.infection_monkey.master.mock_puppet import MockPuppet from common import OperatingSystems +from common.configuration.agent_sub_configurations import ( + ExploitationConfiguration, + ExploiterConfiguration, +) from infection_monkey.master import Exploiter from infection_monkey.model import VictimHost @@ -35,18 +39,28 @@ def callback(): @pytest.fixture -def exploiter_config(): - return { - "options": {"dropper_path_linux": "/tmp/monkey"}, - "brute_force": [ - {"name": "MSSQLExploiter", "options": {"timeout": 10}}, - {"name": "SSHExploiter", "options": {}}, - {"name": "WmiExploiter", "options": {"timeout": 10}}, - ], - "vulnerability": [ - {"name": "ZerologonExploiter", "options": {}}, - ], - } +def exploiter_config(default_agent_config): + brute_force = [ + ExploiterConfiguration( + name="MSSQLExploiter", options={"timeout": 10} + ), + ExploiterConfiguration( + name="SSHExploiter", options={} + ), + ExploiterConfiguration( + name="WmiExploiter", options={"timeout": 10} + ), + ] + vulnerability = [ + ExploiterConfiguration( + name="ZerologonExploiter", options={} + ) + ] + return ExploitationConfiguration( + options=default_agent_config.propagation.exploitation.options, + brute_force=brute_force, + vulnerability=vulnerability, + ) @pytest.fixture