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

Agent: Address addr type TODO in RelayConnectionHandler

This commit is contained in:
Mike Salvatore 2022-09-06 14:54:13 -04:00
parent 1cd544e227
commit 76fe0be990
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
monkey/infection_monkey/network/relay/relay_connection_handler.py
View File

@ -13,10 +13,13 @@ class RelayConnectionHandler:
self._relay_user_handler = relay_user_handler self._relay_user_handler = relay_user_handler
def handle_new_connection(self, sock: socket.socket): def handle_new_connection(self, sock: socket.socket):
addr, _ = sock.getpeername()
addr = IPv4Address(addr)
control_message = sock.recv(socket.MSG_PEEK) control_message = sock.recv(socket.MSG_PEEK)
addr, _ = sock.getpeername() # TODO check the type of the addr object
if control_message.startswith(RELAY_CONTROL_MESSAGE): if control_message.startswith(RELAY_CONTROL_MESSAGE):
self._relay_user_handler.disconnect_user(IPv4Address(addr)) self._relay_user_handler.disconnect_user(addr)
else: else:
self._relay_user_handler.add_relay_user(IPv4Address(addr)) self._relay_user_handler.add_relay_user(addr)
self._pipe_spawner.spawn_pipe(sock) self._pipe_spawner.spawn_pipe(sock)