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

Agent: Rename RELAY_CONTROL_MESSAGE -> RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST

This commit is contained in:
Shreya Malviya 2022-09-08 13:42:11 +05:30 committed by Mike Salvatore
parent cb4af415c1
commit aa1c31efb2
3 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
monkey/infection_monkey/network/relay/__init__.py
View File

@ -1,4 +1,7 @@
from .relay_connection_handler import RelayConnectionHandler, RELAY_CONTROL_MESSAGE from .relay_connection_handler import (
RelayConnectionHandler,
RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST,
)
from .relay_user_handler import RelayUser, RelayUserHandler from .relay_user_handler import RelayUser, RelayUserHandler
from .sockets_pipe import SocketsPipe from .sockets_pipe import SocketsPipe
from .tcp_connection_handler import TCPConnectionHandler from .tcp_connection_handler import TCPConnectionHandler

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

@ -4,7 +4,7 @@ from ipaddress import IPv4Address
from .relay_user_handler import RelayUserHandler from .relay_user_handler import RelayUserHandler
from .tcp_pipe_spawner import TCPPipeSpawner from .tcp_pipe_spawner import TCPPipeSpawner
RELAY_CONTROL_MESSAGE = b"infection-monkey-relay-control-message: -" RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST = b"infection-monkey-relay-control-message: -"
class RelayConnectionHandler: class RelayConnectionHandler:
@ -25,7 +25,7 @@ class RelayConnectionHandler:
control_message = sock.recv(socket.MSG_PEEK) control_message = sock.recv(socket.MSG_PEEK)
if control_message.startswith(RELAY_CONTROL_MESSAGE): if control_message.startswith(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST):
self._relay_user_handler.disconnect_user(addr) self._relay_user_handler.disconnect_user(addr)
else: else:
self._relay_user_handler.add_relay_user(addr) self._relay_user_handler.add_relay_user(addr)

4
monkey/infection_monkey/network/relay/utils.py
View File

@ -6,7 +6,7 @@ import requests
from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
from common.network.network_utils import address_to_ip_port from common.network.network_utils import address_to_ip_port
from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST
from infection_monkey.utils.threading import create_daemon_thread from infection_monkey.utils.threading import create_daemon_thread
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -57,6 +57,6 @@ def _send_remove_from_waitlist_control_message_to_relay(server: str):
try: try:
d_socket.connect((ip, int(port))) d_socket.connect((ip, int(port)))
d_socket.send(RELAY_CONTROL_MESSAGE) d_socket.send(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST)
except OSError as err: except OSError as err:
logger.error(f"Error connecting to socket {server}: {err}") logger.error(f"Error connecting to socket {server}: {err}")