From ea621b49d144834eca5082e7b62d9050a9af7eab Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Tue, 7 Dec 2021 13:27:40 +0200 Subject: [PATCH] Agent: change agent startup to check if agent should run via control channel --- monkey/infection_monkey/monkey.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e63484e70..a1a11b925 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -8,8 +8,9 @@ import time import infection_monkey.tunnel as tunnel from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version -from infection_monkey.config import WormConfiguration +from infection_monkey.config import GUID, WormConfiguration from infection_monkey.control import ControlClient +from infection_monkey.master.control_channel import ControlChannel from infection_monkey.master.mock_master import MockMaster from infection_monkey.model import DELAY_DELETE_CMD from infection_monkey.network.firewall import app as firewall @@ -73,9 +74,11 @@ class InfectionMonkey: if is_windows_os(): T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send() - if InfectionMonkey._is_monkey_alive_by_config(): - logger.info("Monkey marked 'not alive' from configuration.") - return + # TODO move this function + should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop() + logger.info(f"Should monkey stop: {should_stop}") + if should_stop: + sys.exit(1) if InfectionMonkey._is_upgrade_to_64_needed(): self._upgrade_to_64()