From ce492d25f4a6ca598ca51d78c0d65d832d07be49 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Jul 2021 04:22:21 -0400 Subject: [PATCH] Agent: Return boolean from "should_monkey_run()" A function named "should_monkey_run()" should let you know whether or not monkey should run. Before this commit, the function was responsible for flow control, as it raised a PlannedShutdownException, resulting in the shutdown of the monkey agent. --- monkey/infection_monkey/control.py | 7 ++----- monkey/infection_monkey/monkey.py | 7 ++++++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index b71e27a5e..109110498 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -22,7 +22,6 @@ from infection_monkey.config import GUID, WormConfiguration from infection_monkey.network.info import check_internet_access, local_ips from infection_monkey.transport.http import HTTPConnectProxy from infection_monkey.transport.tcp import TcpProxy -from infection_monkey.utils.exceptions.planned_shutdown_exception import PlannedShutdownException requests.packages.urllib3.disable_warnings() @@ -401,10 +400,8 @@ class ControlClient(object): and ControlClient.can_island_see_port(vulnerable_port) and WormConfiguration.started_on_island ): - raise PlannedShutdownException( - "Monkey shouldn't run on current machine " - "(it will be exploited later with more depth)." - ) + return False + return True @staticmethod diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 6059c291a..0c0cca9cc 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -125,7 +125,12 @@ class InfectionMonkey(object): if is_running_on_island(): WormConfiguration.started_on_island = True ControlClient.report_start_on_island() - ControlClient.should_monkey_run(self._opts.vulnerable_port) + + if not ControlClient.should_monkey_run(self._opts.vulnerable_port): + raise PlannedShutdownException( + "Monkey shouldn't run on current machine " + "(it will be exploited later with more depth)." + ) if firewall.is_enabled(): firewall.add_firewall_rule()