Move tcp scanner to use new check_tcp_pors

chaos_monkey/network/tcp_scanner.py

@@ -1,9 +1,7 @@
-import time
 from random import shuffle
 
-from model.host import VictimHost
 from network import HostScanner, HostFinger
-from network.tools import check_tcp_port
+from network.tools import check_tcp_ports
 
 __author__ = 'itamar'
 
@@ -18,29 +16,25 @@ class TcpScanner(HostScanner, HostFinger):
         return self.get_host_fingerprint(host, True)
 
     def get_host_fingerprint(self, host, only_one_port=False):
-        assert isinstance(host, VictimHost)
+        """
+        Scans a target host to see if it's alive using the tcp_target_ports specified in the configuration.
+        :param host: VictimHost structure
+        :param only_one_port: Currently unused.
+        :return: T/F if there is at least one open port. In addition, the host object is updated to mark those services as alive.
+        """
 
-        count = 0
         # maybe hide under really bad detection systems
         target_ports = self._config.tcp_target_ports[:]
         shuffle(target_ports)
 
-        for target_port in target_ports:
-
-            is_open, banner = check_tcp_port(host.ip_addr,
-                                             target_port,
-                                             self._config.tcp_scan_timeout / 1000.0,
-                                             self._config.tcp_scan_get_banner)
-
-            if is_open:
-                count += 1
+        ports, banners = check_tcp_ports(host.ip_addr, target_ports, self._config.tcp_scan_timeout / 1000.0)
+        if len(ports) != 0:
+            for target_port, banner in zip(ports, banners):
                 service = 'tcp-' + str(target_port)
                 host.services[service] = {}
                 if banner:
                     host.services[service]['banner'] = banner
                 if only_one_port:
                     break
-            else:
-                time.sleep(self._config.tcp_scan_interval / 1000.0)
 
-        return count != 0
+        return len(ports) != 0

chaos_monkey/network/tools.py

@@ -2,6 +2,7 @@ import logging
 import select
 import socket
 import struct
+import time
 
 DEFAULT_TIMEOUT = 10
 BANNER_READ = 1024
@@ -36,6 +37,14 @@ def struct_unpack_tracker_string(data, index):
 
 
 def check_tcp_port(ip, port, timeout=DEFAULT_TIMEOUT, get_banner=False):
+    """
+    Checks if a given TCP port is open
+    :param ip: Target IP
+    :param port: Target Port
+    :param timeout: Timeout for socket connection
+    :param get_banner:  if true, pulls first BANNER_READ bytes from the socket.
+    :return: Tuple, T/F + banner if requested.
+    """
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     sock.settimeout(timeout)
 
@@ -54,7 +63,7 @@ def check_tcp_port(ip, port, timeout=DEFAULT_TIMEOUT, get_banner=False):
             read_ready, _, _ = select.select([sock], [], [], timeout)
             if len(read_ready) > 0:
                 banner = sock.recv(BANNER_READ)
-    except:
+    except socket.error:
         pass
 
     sock.close()
@@ -62,6 +71,13 @@ def check_tcp_port(ip, port, timeout=DEFAULT_TIMEOUT, get_banner=False):
 
 
 def check_udp_port(ip, port, timeout=DEFAULT_TIMEOUT):
+    """
+    Checks if a given UDP port is open by checking if it replies to an empty message
+    :param ip:  Target IP
+    :param port: Target port
+    :param timeout: Timeout to wait
+    :return: Tuple, T/F + banner
+    """
     sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
     sock.settimeout(timeout)
 
@@ -77,3 +93,53 @@ def check_udp_port(ip, port, timeout=DEFAULT_TIMEOUT):
     sock.close()
 
     return is_open, data
+
+
+def check_tcp_ports(ip, ports, timeout=DEFAULT_TIMEOUT, get_banner=False):
+    """
+    Checks whether any of the given ports are open on a target IP.
+    :param ip:  IP of host to attack
+    :param ports: List of ports to attack. Must not be empty.
+    :param timeout: Amount of time to wait for connection
+    :param get_banner: T/F if to get first packets from server
+    :return: list of open ports. If get_banner=True, then a matching list of banners.
+    """
+    sockets = [socket.socket(socket.AF_INET, socket.SOCK_STREAM) for _ in range(len(ports))]
+    [s.setblocking(0) for s in sockets]
+    good_ports = []
+    try:
+        for sock, port in zip(sockets, ports):
+            LOG.debug("Connecting to port %d" % port)
+            err = sock.connect_ex((ip, port))
+            if err == 0:
+                good_ports.append((port, sock))
+            if err == 10035:  # WSAEWOULDBLOCK is valid, see https://msdn.microsoft.com/en-us/library/windows/desktop/ms740668%28v=vs.85%29.aspx?f=255&MSPPError=-2147217396
+                good_ports.append((port, sock))
+
+        if len(good_ports) != 0:
+            time.sleep(timeout)
+            read_sockets, write_sockets, _ = \
+                select.select(
+                    [s[1] for s in good_ports],
+                    [s[1] for s in good_ports],
+                    [s[1] for s in good_ports],
+                    0)  # no timeout because we've already slept
+            connected_ports_sockets = [x for x in good_ports if x[1] in write_sockets]
+            LOG.debug(
+                "On host %s discovered the following ports %s" %
+                (str(ip), ",".join([str(x) for x in connected_ports_sockets])))
+            banners = []
+            if get_banner:
+                # read first X bytes
+                banners = [sock.recv(BANNER_READ) if sock in read_sockets else ""
+                           for port, sock in connected_ports_sockets]
+                pass
+            # try to cleanup
+            [s[1].close() for s in good_ports]
+            return [port for port, sock in connected_ports_sockets], banners
+        else:
+            return [], []
+
+    except socket.error as exc:
+        LOG.warning("Exception when checking ports on host %s, Exception: %s", str(ip), exc)
+        return [], []