p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Agent: Rework call of find_server in monkey.py

This commit is contained in:
Ilija Lazoroski 2022-09-07 19:53:54 +02:00 committed by Mike Salvatore
parent b89ba06fd1
commit 6ee15e22b8
1 changed files with 14 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
monkey/infection_monkey/monkey.py
View File

@ -98,9 +98,9 @@ class InfectionMonkey:
self._opts = self._get_arguments(args) self._opts = self._get_arguments(args)
# TODO: Revisit variable names # TODO: Revisit variable names
self._get_server() server = self._get_server()
self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self.server) self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server)
self._control_client = ControlClient(server_address=self.server) self._control_client = ControlClient(server_address=server)
# TODO Refactor the telemetry messengers to accept control client # TODO Refactor the telemetry messengers to accept control client
# and remove control_client_object # and remove control_client_object
@ -124,9 +124,18 @@ class InfectionMonkey:
def _get_server(self): def _get_server(self):
servers_iterator = (s for s in self._opts.servers) servers_iterator = (s for s in self._opts.servers)
self.server = find_server(servers_iterator) server = find_server(servers_iterator)
if server:
logger.debug(f"Default server set to: {server}")
else:
# TODO: Exit here
raise Exception(
f"Failed to connect to the island via any known servers: {self._opts.servers}"
)
send_relay_control_message(servers_iterator) send_relay_control_message(servers_iterator)
return server
@staticmethod @staticmethod
def _log_arguments(args): def _log_arguments(args):
arg_string = " ".join([f"{key}: {value}" for key, value in vars(args).items()]) arg_string = " ".join([f"{key}: {value}" for key, value in vars(args).items()])
@ -140,7 +149,7 @@ class InfectionMonkey:
logger.info("Agent is starting...") logger.info("Agent is starting...")
logger.info(f"Agent GUID: {GUID}") logger.info(f"Agent GUID: {GUID}")
self._connect_to_island() self._control_client.wakeup(parent=self._opts.parent)
# TODO: Reevaluate who is responsible to send this information # TODO: Reevaluate who is responsible to send this information
if is_windows_os(): if is_windows_os():
@ -158,23 +167,6 @@ class InfectionMonkey:
self._setup() self._setup()
self._master.start() self._master.start()
def _connect_to_island(self):
# Sets island's IP and port for monkey to communicate to
if self._current_server_is_set():
logger.debug(f"Default server set to: {self._control_client.server_address}")
else:
raise Exception(
f"Failed to connect to the island via " f"any known server address: {self.server}"
)
self._control_client.wakeup(parent=self._opts.parent)
def _current_server_is_set(self) -> bool:
if self.server:
return True
return False
def _setup(self): def _setup(self):
logger.debug("Starting the setup phase.") logger.debug("Starting the setup phase.")