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

Agent: Change relay control message 

'-' to 'infection-monkey-relay-control-message: -'
This commit is contained in:
Ilija Lazoroski 2022-09-01 10:35:29 +02:00
parent 94fba0fdf0
commit ecdf04a85e
2 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
monkey/infection_monkey/tcp_relay.py
View File

@ -6,6 +6,7 @@ from typing import List
from infection_monkey.transport.tcp import TcpProxy
DEFAULT_NEW_CLIENT_TIMEOUT = 3 # Wait up to 3 seconds for potential new clients to connect
RELAY_CONTROL_MESSAGE = b"infection-monkey-relay-control-message: -"
@dataclass
@ -77,7 +78,7 @@ class TCPRelay(Thread):
self._potential_users.append(RelayUser(user, time()))
def on_user_data_received(self, data: bytes, user: str) -> bool:
if data.startswith(b"-"):
if data.startswith(RELAY_CONTROL_MESSAGE):
self._disconnect_user(user)
return False
return True

4
monkey/tests/unit_tests/infection_monkey/test_tcp_relay.py
View File

@ -1,6 +1,6 @@
from threading import Thread
from monkey.infection_monkey.tcp_relay import TCPRelay
from monkey.infection_monkey.tcp_relay import RELAY_CONTROL_MESSAGE, TCPRelay
def join_or_kill_thread(thread: Thread, timeout: float):
@ -46,7 +46,7 @@ def test_user_removed_on_request():
relay = TCPRelay(9975, "0.0.0.0", 9976)
new_user = "0.0.0.1"
relay.on_user_connected(new_user)
relay.on_user_data_received(b"-", "0.0.0.1")
relay.on_user_data_received(RELAY_CONTROL_MESSAGE, "0.0.0.1")
users = relay.relay_users()
assert len(users) == 0