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

Merge pull request #153 from guardicore/feature/async_scan 

Opportunistic waiting, make get_tcp_ports O(timeout)
This commit is contained in:
itaymmguardicore 2018-08-30 15:00:33 +03:00 committed by GitHub
parent 62f6d7748b 97e5b96637
commit 686cca1723
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
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
LOG = logging.getLogger(__name__)
SLEEP_BETWEEN_POLL = 0.5
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)
if len(possible_ports) != 0:
time.sleep(timeout)
sock_objects = [s[1] for s in possible_ports]
# first filter
_, writeable_sockets, _ = select.select(sock_objects, sock_objects, sock_objects, 0)
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
timeout = int(round(timeout)) # clamp to integer, to avoid checking input
time_left = timeout
sockets_to_try = possible_ports[:]
connected_ports_sockets = []
while (time_left >= 0) and len(sockets_to_try):
sock_objects = [s[1] for s in sockets_to_try]
_, writeable_sockets, _ = select.select(sock_objects, sock_objects, sock_objects, 0)
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(
"On host %s discovered the following ports %s" %
(str(ip), ",".join([str(s[0]) for s in connected_ports_sockets])))