From fb0fea6f6af99e10f6edbc40a505ae29d4710727 Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Wed, 1 Jan 2020 15:33:02 +0200 Subject: [PATCH] Improved the monkey start function structure a bit, extracted to functions Prep work for changing system info collection to modular system --- monkey/infection_monkey/monkey.py | 269 ++++++++++-------- .../post_breach/post_breach_handler.py | 2 +- 2 files changed, 147 insertions(+), 124 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 80d2d8642..8543505a7 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -32,11 +32,17 @@ from infection_monkey.telemetry.attack.t1106_telem import T1106Telem from common.utils.attack_utils import ScanStatus, UsageEnum from infection_monkey.exploit.HostExploiter import HostExploiter +MAX_DEPTH_REACHED_MESSAGE = "Reached max depth, shutting down" + __author__ = 'itamar' LOG = logging.getLogger(__name__) +class PlannedShutdown(Exception): + pass + + class InfectionMonkey(object): def __init__(self, args): self._keep_running = False @@ -87,143 +93,158 @@ class InfectionMonkey(object): LOG.debug("Default server: %s is already in command servers list" % self._default_server) def start(self): - LOG.info("Monkey is running...") + try: + LOG.info("Monkey is starting...") - # Sets island's IP and port for monkey to communicate to - if not self.set_default_server(): - return - self.set_default_port() + LOG.debug("Starting the setup phase.") + # Sets island's IP and port for monkey to communicate to + self.set_default_server() + self.set_default_port() - # Create a dir for monkey files if there isn't one - create_monkey_dir() + # Create a dir for monkey files if there isn't one + create_monkey_dir() - if WindowsUpgrader.should_upgrade(): - self._upgrading_to_64 = True - self._singleton.unlock() - LOG.info("32bit monkey running on 64bit Windows. Upgrading.") - WindowsUpgrader.upgrade(self._opts) - return + self.upgrade_to_64_if_needed() - ControlClient.wakeup(parent=self._parent) - ControlClient.load_control_config() + ControlClient.wakeup(parent=self._parent) + ControlClient.load_control_config() - if is_windows_os(): - T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send() + if is_windows_os(): + T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send() - if not WormConfiguration.alive: - LOG.info("Marked not alive from configuration") - return + self.shutdown_by_not_alive_config() - if firewall.is_enabled(): - firewall.add_firewall_rule() + if firewall.is_enabled(): + firewall.add_firewall_rule() - monkey_tunnel = ControlClient.create_control_tunnel() - if monkey_tunnel: - monkey_tunnel.start() + monkey_tunnel = ControlClient.create_control_tunnel() + if monkey_tunnel: + monkey_tunnel.start() - StateTelem(is_done=False).send() - TunnelTelem().send() + StateTelem(is_done=False).send() + TunnelTelem().send() + LOG.debug("Starting the post-breach phase.") + self.collect_system_info_if_configured() + PostBreach().execute_all_configured() + + LOG.debug("Starting the propagation phase.") + self.shutdown_by_max_depth_reached() + + for iteration_index in range(WormConfiguration.max_iterations): + ControlClient.keepalive() + ControlClient.load_control_config() + + self._network.initialize() + + self._fingerprint = HostFinger.get_instances() + + self._exploiters = HostExploiter.get_classes() + + if not self._keep_running or not WormConfiguration.alive: + break + + machines = self._network.get_victim_machines(max_find=WormConfiguration.victims_max_find, + stop_callback=ControlClient.check_for_stop) + is_empty = True + for machine in machines: + if ControlClient.check_for_stop(): + break + + is_empty = False + for finger in self._fingerprint: + LOG.info("Trying to get OS fingerprint from %r with module %s", + machine, finger.__class__.__name__) + finger.get_host_fingerprint(machine) + + ScanTelem(machine).send() + + # skip machines that we've already exploited + if machine in self._exploited_machines: + LOG.debug("Skipping %r - already exploited", + machine) + continue + elif machine in self._fail_exploitation_machines: + if WormConfiguration.retry_failed_explotation: + LOG.debug("%r - exploitation failed before, trying again", machine) + else: + LOG.debug("Skipping %r - exploitation failed before", machine) + continue + + if monkey_tunnel: + monkey_tunnel.set_tunnel_for_host(machine) + if self._default_server: + if self._network.on_island(self._default_server): + machine.set_default_server(get_interface_to_target(machine.ip_addr) + + (':' + self._default_server_port if self._default_server_port else '')) + else: + machine.set_default_server(self._default_server) + LOG.debug("Default server for machine: %r set to %s" % (machine, machine.default_server)) + + # Order exploits according to their type + if WormConfiguration.should_exploit: + self._exploiters = sorted(self._exploiters, key=lambda exploiter_: exploiter_.EXPLOIT_TYPE.value) + host_exploited = False + for exploiter in [exploiter(machine) for exploiter in self._exploiters]: + if self.try_exploiting(machine, exploiter): + host_exploited = True + VictimHostTelem('T1210', ScanStatus.USED, machine=machine).send() + break + if not host_exploited: + self._fail_exploitation_machines.add(machine) + VictimHostTelem('T1210', ScanStatus.SCANNED, machine=machine).send() + if not self._keep_running: + break + + if (not is_empty) and (WormConfiguration.max_iterations > iteration_index + 1): + time_to_sleep = WormConfiguration.timeout_between_iterations + LOG.info("Sleeping %d seconds before next life cycle iteration", time_to_sleep) + time.sleep(time_to_sleep) + + if self._keep_running and WormConfiguration.alive: + LOG.info("Reached max iterations (%d)", WormConfiguration.max_iterations) + elif not WormConfiguration.alive: + LOG.info("Marked not alive from configuration") + + # if host was exploited, before continue to closing the tunnel ensure the exploited host had its chance to + # connect to the tunnel + if len(self._exploited_machines) > 0: + time_to_sleep = WormConfiguration.keep_tunnel_open_time + LOG.info("Sleeping %d seconds for exploited machines to connect to tunnel", time_to_sleep) + time.sleep(time_to_sleep) + + if monkey_tunnel: + monkey_tunnel.stop() + monkey_tunnel.join() + except PlannedShutdown: + LOG.info("A planned shutdown of the Monkey occurred. Logging the reason and finishing execution.") + LOG.exception("Planned shutdown, reason:") + + def shutdown_by_max_depth_reached(self): + if 0 == WormConfiguration.depth: + TraceTelem(MAX_DEPTH_REACHED_MESSAGE).send() + raise PlannedShutdown(MAX_DEPTH_REACHED_MESSAGE) + else: + LOG.debug("Running with depth: %d" % WormConfiguration.depth) + + def collect_system_info_if_configured(self): if WormConfiguration.collect_system_info: LOG.debug("Calling system info collection") system_info_collector = SystemInfoCollector() system_info = system_info_collector.get_info() SystemInfoTelem(system_info).send() - # Executes post breach actions - PostBreach().execute() + def shutdown_by_not_alive_config(self): + if not WormConfiguration.alive: + raise PlannedShutdown("Marked 'not alive' from configuration.") - if 0 == WormConfiguration.depth: - TraceTelem("Reached max depth, shutting down").send() - return - else: - LOG.debug("Running with depth: %d" % WormConfiguration.depth) - - for iteration_index in range(WormConfiguration.max_iterations): - ControlClient.keepalive() - ControlClient.load_control_config() - - self._network.initialize() - - self._fingerprint = HostFinger.get_instances() - - self._exploiters = HostExploiter.get_classes() - - if not self._keep_running or not WormConfiguration.alive: - break - - machines = self._network.get_victim_machines(max_find=WormConfiguration.victims_max_find, - stop_callback=ControlClient.check_for_stop) - is_empty = True - for machine in machines: - if ControlClient.check_for_stop(): - break - - is_empty = False - for finger in self._fingerprint: - LOG.info("Trying to get OS fingerprint from %r with module %s", - machine, finger.__class__.__name__) - finger.get_host_fingerprint(machine) - - ScanTelem(machine).send() - - # skip machines that we've already exploited - if machine in self._exploited_machines: - LOG.debug("Skipping %r - already exploited", - machine) - continue - elif machine in self._fail_exploitation_machines: - if WormConfiguration.retry_failed_explotation: - LOG.debug("%r - exploitation failed before, trying again", machine) - else: - LOG.debug("Skipping %r - exploitation failed before", machine) - continue - - if monkey_tunnel: - monkey_tunnel.set_tunnel_for_host(machine) - if self._default_server: - if self._network.on_island(self._default_server): - machine.set_default_server(get_interface_to_target(machine.ip_addr) + - (':' + self._default_server_port if self._default_server_port else '')) - else: - machine.set_default_server(self._default_server) - LOG.debug("Default server for machine: %r set to %s" % (machine, machine.default_server)) - - # Order exploits according to their type - if WormConfiguration.should_exploit: - self._exploiters = sorted(self._exploiters, key=lambda exploiter_: exploiter_.EXPLOIT_TYPE.value) - host_exploited = False - for exploiter in [exploiter(machine) for exploiter in self._exploiters]: - if self.try_exploiting(machine, exploiter): - host_exploited = True - VictimHostTelem('T1210', ScanStatus.USED, machine=machine).send() - break - if not host_exploited: - self._fail_exploitation_machines.add(machine) - VictimHostTelem('T1210', ScanStatus.SCANNED, machine=machine).send() - if not self._keep_running: - break - - if (not is_empty) and (WormConfiguration.max_iterations > iteration_index + 1): - time_to_sleep = WormConfiguration.timeout_between_iterations - LOG.info("Sleeping %d seconds before next life cycle iteration", time_to_sleep) - time.sleep(time_to_sleep) - - if self._keep_running and WormConfiguration.alive: - LOG.info("Reached max iterations (%d)", WormConfiguration.max_iterations) - elif not WormConfiguration.alive: - LOG.info("Marked not alive from configuration") - - # if host was exploited, before continue to closing the tunnel ensure the exploited host had its chance to - # connect to the tunnel - if len(self._exploited_machines) > 0: - time_to_sleep = WormConfiguration.keep_tunnel_open_time - LOG.info("Sleeping %d seconds for exploited machines to connect to tunnel", time_to_sleep) - time.sleep(time_to_sleep) - - if monkey_tunnel: - monkey_tunnel.stop() - monkey_tunnel.join() + def upgrade_to_64_if_needed(self): + if WindowsUpgrader.should_upgrade(): + self._upgrading_to_64 = True + self._singleton.unlock() + LOG.info("32bit monkey running on 64bit Windows. Upgrading.") + WindowsUpgrader.upgrade(self._opts) + raise PlannedShutdown("Finished upgrading from 32bit to 64bit.") def cleanup(self): LOG.info("Monkey cleanup started") @@ -346,9 +367,11 @@ class InfectionMonkey(object): self._default_server_port = '' def set_default_server(self): + """ + Sets the default server for the Monkey to communicate back to. + :raises PlannedShutdown if couldn't find the server. + """ if not ControlClient.find_server(default_tunnel=self._default_tunnel): - LOG.info("Monkey couldn't find server. Going down.") - return False + raise PlannedShutdown("Monkey couldn't find server with {} default tunnel.".format(self._default_tunnel)) self._default_server = WormConfiguration.current_server LOG.debug("default server set to: %s" % self._default_server) - return True diff --git a/monkey/infection_monkey/post_breach/post_breach_handler.py b/monkey/infection_monkey/post_breach/post_breach_handler.py index 7474c8ef1..d700bac62 100644 --- a/monkey/infection_monkey/post_breach/post_breach_handler.py +++ b/monkey/infection_monkey/post_breach/post_breach_handler.py @@ -20,7 +20,7 @@ class PostBreach(object): self.os_is_linux = not is_windows_os() self.pba_list = self.config_to_pba_list() - def execute(self): + def execute_all_configured(self): """ Executes all post breach actions. """