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

Opportunistic waiting, make get_tcp_ports O(timeout) rather than timeout.

This commit is contained in:
Daniel Goldberg 2018-02-15 11:06:01 +02:00
parent 977e0a8769
commit 97e5b96637
1 changed files with 19 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
infection_monkey/network/tools.py
View File

@ -8,6 +8,7 @@ DEFAULT_TIMEOUT = 10
BANNER_READ = 1024 BANNER_READ = 1024
LOG = logging.getLogger(__name__) LOG = logging.getLogger(__name__)
SLEEP_BETWEEN_POLL = 0.5
def struct_unpack_tracker(data, index, fmt): def struct_unpack_tracker(data, index, fmt):
@ -126,15 +127,24 @@ def check_tcp_ports(ip, ports, timeout=DEFAULT_TIMEOUT, get_banner=False):
LOG.warning("Failed to connect to port %s, error code is %d", port, err) LOG.warning("Failed to connect to port %s, error code is %d", port, err)
if len(possible_ports) != 0: if len(possible_ports) != 0:
time.sleep(timeout) timeout = int(round(timeout)) # clamp to integer, to avoid checking input
sock_objects = [s[1] for s in possible_ports] time_left = timeout
# first filter sockets_to_try = possible_ports[:]
_, writeable_sockets, _ = select.select(sock_objects, sock_objects, sock_objects, 0) connected_ports_sockets = []
for s in writeable_sockets: while (time_left >= 0) and len(sockets_to_try):
try: # actual test sock_objects = [s[1] for s in sockets_to_try]
connected_ports_sockets.append((s.getpeername()[1], s))
except socket.error: # bad socket, select didn't filter it properly _, writeable_sockets, _ = select.select(sock_objects, sock_objects, sock_objects, 0)
pass for s in writeable_sockets:
try: # actual test
connected_ports_sockets.append((s.getpeername()[1], s))
except socket.error: # bad socket, select didn't filter it properly
pass
sockets_to_try = [s for s in sockets_to_try if s not in connected_ports_sockets]
if sockets_to_try:
time.sleep(SLEEP_BETWEEN_POLL)
timeout -= SLEEP_BETWEEN_POLL
LOG.debug( LOG.debug(
"On host %s discovered the following ports %s" % "On host %s discovered the following ports %s" %
(str(ip), ",".join([str(s[0]) for s in connected_ports_sockets]))) (str(ip), ",".join([str(s[0]) for s in connected_ports_sockets])))