p17860254
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: readability and style changes in monkey.py: refactored back from raising exceptions to logging and returning, not storing part of island config options as separate parameters, etc.

This commit is contained in:
VakarisZ 2021-12-01 18:13:27 +02:00
parent 81e61dcea5
commit ad6b309523
1 changed files with 28 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
monkey/infection_monkey/monkey.py
View File

@ -23,7 +23,6 @@ from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter im
from infection_monkey.telemetry.state_telem import StateTelem
from infection_monkey.telemetry.tunnel_telem import TunnelTelem
from infection_monkey.utils.environment import is_windows_os
from infection_monkey.utils.exceptions.planned_shutdown_error import PlannedShutdownError
from infection_monkey.utils.monkey_dir import get_monkey_dir_path, 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
@ -32,24 +31,18 @@ from infection_monkey.windows_upgrader import WindowsUpgrader
logger = logging.getLogger(__name__)
class PlannedShutdownError(Exception):
# Raise when we deliberately want to shut down the agent
pass
class InfectionMonkey:
def __init__(self, args):
logger.info("Monkey is initializing...")
self._master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
self._singleton = SystemSingleton()
self._opts = self._get_arguments(args)
self._parent = self._opts.parent
self._default_tunnel = self._opts.tunnel
self._default_server = self._opts.server
# TODO Used in propagation phase to set the default server for the victim
self._default_server_port = None
self._set_propagation_depth()
self._add_default_server_to_config()
self._monkey_tunnel = None
# TODO used in propogation phase
self._monkey_inbound_tunnel = None
@staticmethod
def _get_arguments(args):
@ -87,7 +80,8 @@ class InfectionMonkey:
def start(self):
if self._is_another_monkey_running():
raise PlannedShutdownError("Another instance of the monkey is already running.")
logger.info("Another instance of the monkey is already running")
return
logger.info("Monkey is starting...")
@ -98,11 +92,13 @@ class InfectionMonkey:
T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
if InfectionMonkey._is_monkey_alive_by_config():
raise PlannedShutdownError("Monkey marked 'not alive' from configuration.")
logger.info("Monkey marked 'not alive' from configuration.")
return
if InfectionMonkey._is_upgrade_to_64_needed():
self._upgrade_to_64()
raise PlannedShutdownError("32 bit Agent can't run on 64 bit system.")
logger.info("32 bit Agent can't run on 64 bit system.")
return
self._setup()
self._master.start()
@ -111,11 +107,13 @@ class InfectionMonkey:
# Sets island's IP and port for monkey to communicate to
if not self._is_default_server_set():
raise Exception(
"Monkey couldn't find server with {} default tunnel.".format(self._default_tunnel)
"Monkey couldn't find server with {} default tunnel.".format(
self._opts._default_tunnel
)
)
self._set_default_port()
ControlClient.wakeup(parent=self._parent)
ControlClient.wakeup(parent=self._opts._parent)
ControlClient.load_control_config()
def _is_default_server_set(self) -> bool:
@ -123,7 +121,7 @@ class InfectionMonkey:
Sets the default server for the Monkey to communicate back to.
:return
"""
if not ControlClient.find_server(default_tunnel=self._default_tunnel):
if not ControlClient.find_server(default_tunnel=self._opts._default_tunnel):
return False
self._default_server = WormConfiguration.current_server
logger.debug("default server set to: %s" % self._default_server)
@ -146,21 +144,26 @@ class InfectionMonkey:
def _setup(self):
logger.debug("Starting the setup phase.")
self._should_run_check_for_performance()
if self._should_exit_for_performance():
logger.info(
"Monkey shouldn't run on current machine to improve perfomance"
"(it will be exploited later with more depth)."
)
return
if firewall.is_enabled():
firewall.add_firewall_rule()
self._monkey_tunnel = ControlClient.create_control_tunnel()
if self._monkey_tunnel:
self._monkey_tunnel.start()
self._monkey_inbound_tunnel = ControlClient.create_control_tunnel()
if self._monkey_inbound_tunnel:
self._monkey_inbound_tunnel.start()
StateTelem(is_done=False, version=get_version()).send()
TunnelTelem().send()
register_signal_handlers(self._master)
def _should_run_check_for_performance(self):
def _should_exit_for_performance(self):
"""
This method implements propagation performance enhancing algorithm that
kicks in if the run was started from the Island.
@ -170,11 +173,7 @@ class InfectionMonkey:
WormConfiguration.started_on_island = True
ControlClient.report_start_on_island()
if not ControlClient.should_monkey_run(self._opts.vulnerable_port):
raise PlannedShutdownError(
"Monkey shouldn't run on current machine to improve perfomance"
"(it will be exploited later with more depth)."
)
return not ControlClient.should_monkey_run(self._opts.vulnerable_port)
def _is_another_monkey_running(self):
return not self._singleton.try_lock()
@ -195,9 +194,9 @@ class InfectionMonkey:
if self._master:
self._master.cleanup()
if self._monkey_tunnel:
self._monkey_tunnel.stop()
self._monkey_tunnel.join()
if self._monkey_inbound_tunnel:
self._monkey_inbound_tunnel.stop()
self._monkey_inbound_tunnel.join()
if firewall.is_enabled():
firewall.remove_firewall_rule()