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
commit 686cca1723
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 19 additions and 9 deletions

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])))