Agent: change agent startup to check if agent should run via control channel

This commit is contained in:
VakarisZ 2021-12-07 13:27:40 +02:00
parent e4280660df
commit ea621b49d1
1 changed files with 7 additions and 4 deletions

View File

@ -8,8 +8,9 @@ import time
import infection_monkey.tunnel as tunnel import infection_monkey.tunnel as tunnel
from common.utils.attack_utils import ScanStatus, UsageEnum from common.utils.attack_utils import ScanStatus, UsageEnum
from common.version import get_version 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.control import ControlClient
from infection_monkey.master.control_channel import ControlChannel
from infection_monkey.master.mock_master import MockMaster from infection_monkey.master.mock_master import MockMaster
from infection_monkey.model import DELAY_DELETE_CMD from infection_monkey.model import DELAY_DELETE_CMD
from infection_monkey.network.firewall import app as firewall from infection_monkey.network.firewall import app as firewall
@ -73,9 +74,11 @@ class InfectionMonkey:
if is_windows_os(): if is_windows_os():
T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send() T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
if InfectionMonkey._is_monkey_alive_by_config(): # TODO move this function
logger.info("Monkey marked 'not alive' from configuration.") should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop()
return logger.info(f"Should monkey stop: {should_stop}")
if should_stop:
sys.exit(1)
if InfectionMonkey._is_upgrade_to_64_needed(): if InfectionMonkey._is_upgrade_to_64_needed():
self._upgrade_to_64() self._upgrade_to_64()