Agent: Remove intialize both from monkey and dropper

Add legacy start and cleanup to the agent which
are the same code reformated in the previous commits.
Reformat start function.
This commit is contained in:
Ilija Lazoroski 2021-11-29 18:34:25 +01:00
parent 75226bdf6e
commit 72f4fc1ef6
4 changed files with 192 additions and 88 deletions

View File

@ -55,10 +55,9 @@ class MonkeyDrops(object):
"destination_path": self.opts.location,
}
def initialize(self):
logger.debug("Dropper is running with config:\n%s", pprint.pformat(self._config))
def start(self):
def legacy_start(self):
if self._config["destination_path"] is None:
logger.error("No destination path specified")
return False
@ -183,7 +182,7 @@ class MonkeyDrops(object):
if monkey_process.poll() is not None:
logger.warning("Seems like monkey died too soon")
def cleanup(self):
def legacy_cleanup(self):
logger.info("Cleaning up the dropper")
try:

View File

@ -118,16 +118,16 @@ def main():
logger.info(f"version: {get_version()}")
monkey = monkey_cls(monkey_args)
monkey.initialize()
try:
monkey.start()
monkey.legacy_start()
# monkey.start()
return True
except Exception as e:
logger.exception("Exception thrown from monkey's start function. More info: {}".format(e))
finally:
monkey.cleanup()
monkey.legacy_cleanup()
# monkey.cleanup()
if "__main__" == __name__:

View File

@ -46,6 +46,9 @@ class MockMaster(IMaster):
logger.info("Finished running system info collectors")
def _run_pbas(self):
# TODO: Create monkey_dir and revise setup in monkey.py
logger.info("Running post breach actions")
name = "AccountDiscovery"
command, result = self._puppet.run_pba(name, {})
@ -123,4 +126,5 @@ class MockMaster(IMaster):
logger.info("Terminating MockMaster")
def cleanup(self) -> None:
# TODO: Cleanup monkey_dir and send telemetry
pass

View File

@ -13,26 +13,23 @@ from common.version import get_version
from infection_monkey.config import WormConfiguration
from infection_monkey.control import ControlClient
from infection_monkey.exploit.HostExploiter import HostExploiter
# 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.network.firewall import app as firewall
from infection_monkey.network.HostFinger import HostFinger
from infection_monkey.network.network_scanner import NetworkScanner
from infection_monkey.network.tools import get_interface_to_target, is_running_on_island
from infection_monkey.post_breach.post_breach_handler import PostBreach
# from infection_monkey.puppet.mock_puppet import MockPuppet
from infection_monkey.puppet.mock_puppet import MockPuppet
from infection_monkey.ransomware.ransomware_payload_builder import build_ransomware_payload
from infection_monkey.system_info import SystemInfoCollector
from infection_monkey.system_singleton import SystemSingleton
from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
from infection_monkey.telemetry.attack.t1107_telem import T1107Telem
from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem
# from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import (
# LegacyTelemetryMessengerAdapter,
# )
from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import (
LegacyTelemetryMessengerAdapter,
)
from infection_monkey.telemetry.scan_telem import ScanTelem
from infection_monkey.telemetry.state_telem import StateTelem
from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
@ -46,8 +43,7 @@ from infection_monkey.utils.monkey_dir import (
remove_monkey_dir,
)
from infection_monkey.utils.monkey_log_path import get_monkey_log_path
# from infection_monkey.utils.signal_handler import register_signal_handlers
from infection_monkey.utils.signal_handler import register_signal_handlers
from infection_monkey.windows_upgrader import WindowsUpgrader
MAX_DEPTH_REACHED_MESSAGE = "Reached max depth, skipping propagation phase."
@ -58,51 +54,49 @@ logger = logging.getLogger(__name__)
class InfectionMonkey(object):
def __init__(self, args):
# self.master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
logger.info("Monkey is initializing...")
self.master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
self._keep_running = False
self._exploited_machines = set()
self._fail_exploitation_machines = set()
self._singleton = SystemSingleton()
self._parent = None
self._default_tunnel = None
self._args = args
self._network = None
self._opts = None
self._set_arguments(args)
self._parent = self._opts.parent
self._default_tunnel = self._opts.tunnel
self._default_server = self._opts.server
self._set_propagation_depth()
self._add_default_server_to_config()
self._network = NetworkScanner()
self._exploiters = None
self._fingerprint = None
self._default_server = None
self._default_server_port = None
self._opts = None
self._upgrading_to_64 = False
self._monkey_tunnel = None
self._post_breach_phase = None
def initialize(self):
logger.info("Monkey is initializing...")
self._check_for_running_monkey()
def _set_arguments(self, args):
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument("-p", "--parent")
arg_parser.add_argument("-t", "--tunnel")
arg_parser.add_argument("-s", "--server")
arg_parser.add_argument("-d", "--depth", type=int)
arg_parser.add_argument("-vp", "--vulnerable-port")
self._opts, self._args = arg_parser.parse_known_args(self._args)
self._opts, _ = arg_parser.parse_known_args(args)
self._log_arguments()
self._parent = self._opts.parent
self._default_tunnel = self._opts.tunnel
self._default_server = self._opts.server
def _log_arguments(self):
arg_string = " ".join([f"{key}: {value}" for key, value in vars(self._opts).items()])
logger.info(f"Monkey started with arguments: {arg_string}")
def _set_propagation_depth(self):
if self._opts.depth is not None:
WormConfiguration._depth_from_commandline = True
WormConfiguration.depth = self._opts.depth
logger.debug("Setting propagation depth from command line")
logger.debug(f"Set propagation depth to {WormConfiguration.depth}")
self._keep_running = True
self._network = NetworkScanner()
def _add_default_server_to_config(self):
if self._default_server:
if self._default_server not in WormConfiguration.command_servers:
logger.debug("Added default server: %s" % self._default_server)
@ -112,28 +106,38 @@ class InfectionMonkey(object):
"Default server: %s is already in command servers list" % self._default_server
)
def _check_for_running_monkey(self):
if not self._singleton.try_lock():
raise Exception("Another instance of the monkey is already running")
def _log_arguments(self):
arg_string = " ".join([f"{key}: {value}" for key, value in vars(self._opts).items()])
logger.info(f"Monkey started with arguments: {arg_string}")
def start(self):
if not self._is_another_monkey_running():
logger.info("Monkey is starting...")
self._connect_to_island()
if InfectionMonkey._is_monkey_alive_by_config():
logger.error("Monkey marked 'not alive' from configuration.")
return
if InfectionMonkey._is_upgrade_to_64_needed():
self._upgrade_to_64()
return
self._setup()
self.master.start()
else:
logger.info("Another instance of the monkey is already running")
def legacy_start(self):
if self._is_another_monkey_running():
raise Exception("Another instance of the monkey is already running")
try:
logger.info("Monkey is starting...")
# Sets the monkey up
self._setup()
# Start post breach phase
self._legacy_setup()
self._start_post_breach_async()
# Start propagation phase
self._start_propagation()
# self.master.start()
except PlannedShutdownException:
logger.info(
"A planned shutdown of the Monkey occurred. Logging the reason and finishing "
@ -141,28 +145,103 @@ class InfectionMonkey(object):
)
logger.exception("Planned shutdown, reason:")
finally:
self._teardown()
def _setup(self):
logger.debug("Starting the setup phase.")
InfectionMonkey._shutdown_by_not_alive_config()
def _connect_to_island(self):
# Sets island's IP and port for monkey to communicate to
self._set_default_server()
if not self._is_default_server_set():
raise Exception(
"Monkey couldn't find server with {} default tunnel.".format(self._default_tunnel)
)
self._set_default_port()
# Create a dir for monkey files if there isn't one
create_monkey_dir()
ControlClient.wakeup(parent=self._parent)
ControlClient.load_control_config()
if ControlClient.check_for_stop():
raise PlannedShutdownException("Monkey has been marked for shutdown.")
def _is_default_server_set(self) -> bool:
"""
Sets the default server for the Monkey to communicate back to.
:return
"""
if not ControlClient.find_server(default_tunnel=self._default_tunnel):
return False
self._default_server = WormConfiguration.current_server
logger.debug("default server set to: %s" % self._default_server)
return True
self._upgrade_to_64_if_needed()
@staticmethod
def _is_monkey_alive_by_config():
return not WormConfiguration.alive
@staticmethod
def _is_upgrade_to_64_needed():
return WindowsUpgrader.should_upgrade()
def _upgrade_to_64(self):
self._upgrading_to_64 = True
self._singleton.unlock()
logger.info("32bit monkey running on 64bit Windows. Upgrading.")
WindowsUpgrader.upgrade(self._opts)
logger.info("Finished upgrading from 32bit to 64bit.")
def _legacy_upgrade_to_64_if_needed(self):
if WindowsUpgrader.should_upgrade():
self._upgrading_to_64 = True
self._singleton.unlock()
logger.info("32bit monkey running on 64bit Windows. Upgrading.")
WindowsUpgrader.upgrade(self._opts)
raise PlannedShutdownException("Finished upgrading from 32bit to 64bit.")
def _setup(self):
logger.debug("Starting the setup phase.")
# Create a dir for monkey files if there isn't one
create_monkey_dir()
# TODO: Evaluate should we run this check
# if not ControlClient.should_monkey_run(self._opts.vulnerable_port):
# logger.error("Monkey shouldn't run on current machine "
# "(it will be exploited later with more depth).")
# return False
# Singleton should handle sending this information
if is_windows_os():
T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
if is_running_on_island():
WormConfiguration.started_on_island = True
ControlClient.report_start_on_island()
if firewall.is_enabled():
firewall.add_firewall_rule()
self._monkey_tunnel = ControlClient.create_control_tunnel()
if self._monkey_tunnel:
self._monkey_tunnel.start()
StateTelem(is_done=False, version=get_version()).send()
TunnelTelem().send()
register_signal_handlers(self.master)
return True
def _legacy_setup(self):
logger.debug("Starting the setup phase.")
self._keep_running = True
# Create a dir for monkey files if there isn't one
create_monkey_dir()
# Sets island's IP and port for monkey to communicate to
self._legacy_set_default_server()
self._set_default_port()
ControlClient.wakeup(parent=self._parent)
ControlClient.load_control_config()
InfectionMonkey._legacy_shutdown_by_not_alive_config()
self._legacy_upgrade_to_64_if_needed()
if not ControlClient.should_monkey_run(self._opts.vulnerable_port):
raise PlannedShutdownException(
@ -187,14 +266,10 @@ class InfectionMonkey(object):
StateTelem(is_done=False, version=get_version()).send()
TunnelTelem().send()
# register_signal_handlers(self.master)
def _is_another_monkey_running(self):
return not self._singleton.try_lock()
@staticmethod
def _shutdown_by_not_alive_config():
if not WormConfiguration.alive:
raise PlannedShutdownException("Marked 'not alive' from configuration.")
def _set_default_server(self):
def _legacy_set_default_server(self):
"""
Sets the default server for the Monkey to communicate back to.
:raises PlannedShutdownException if couldn't find the server.
@ -212,13 +287,10 @@ class InfectionMonkey(object):
except KeyError:
self._default_server_port = ""
def _upgrade_to_64_if_needed(self):
if WindowsUpgrader.should_upgrade():
self._upgrading_to_64 = True
self._singleton.unlock()
logger.info("32bit monkey running on 64bit Windows. Upgrading.")
WindowsUpgrader.upgrade(self._opts)
raise PlannedShutdownException("Finished upgrading from 32bit to 64bit.")
@staticmethod
def _legacy_shutdown_by_not_alive_config():
if not WormConfiguration.alive:
raise PlannedShutdownException("Marked 'not alive' from configuration.")
def _start_post_breach_async(self):
logger.debug("Starting the post-breach phase asynchronously.")
@ -424,8 +496,9 @@ class InfectionMonkey(object):
except Exception as ex:
logger.error(f"An unexpected error occurred while running the ransomware payload: {ex}")
def _teardown(self):
logger.info("Monkey teardown started")
def legacy_cleanup(self):
logger.info("Monkey cleanup started")
self._keep_running = False
if self._monkey_tunnel:
self._monkey_tunnel.stop()
self._monkey_tunnel.join()
@ -437,13 +510,6 @@ class InfectionMonkey(object):
firewall.remove_firewall_rule()
firewall.close()
# self.master.terminate()
# self.master.cleanup()
def cleanup(self):
logger.info("Monkey cleanup started")
self._keep_running = False
if self._upgrading_to_64:
InfectionMonkey._close_tunnel()
firewall.close()
@ -456,6 +522,41 @@ class InfectionMonkey(object):
InfectionMonkey._send_log()
self._singleton.unlock()
# self.master.terminate()
# self.master.cleanup()
InfectionMonkey._self_delete()
logger.info("Monkey is shutting down")
def cleanup(self):
logger.info("Monkey cleanup started")
self._keep_running = False
if self._monkey_tunnel:
self._monkey_tunnel.stop()
self._monkey_tunnel.join()
if self._post_breach_phase:
self._post_breach_phase.join()
if firewall.is_enabled():
firewall.remove_firewall_rule()
firewall.close()
if self._upgrading_to_64:
InfectionMonkey._close_tunnel()
firewall.close()
else:
StateTelem(
is_done=True, version=get_version()
).send() # Signal the server (before closing the tunnel)
InfectionMonkey._close_tunnel()
firewall.close()
InfectionMonkey._send_log()
self._singleton.unlock()
self.master.terminate()
self.master.cleanup()
InfectionMonkey._self_delete()
logger.info("Monkey is shutting down")