Agent: Address CR comments + minor changes in tcp_scanner.py

monkey/infection_monkey/network/tcp_scanner.py

@@ -46,8 +46,9 @@ def _check_tcp_ports(ip: str, ports: List[int], timeout: float = DEFAULT_TIMEOUT
     :return: List of open ports.
     """
     sockets = [socket.socket(socket.AF_INET, socket.SOCK_STREAM) for _ in range(len(ports))]
-    # CR: Don't use list comprehensions if you don't need a list
-    [s.setblocking(False) for s in sockets]
+    for s in sockets:
+        s.setblocking(False)
+
     possible_ports = []
     connected_ports_sockets = []
     try:
@@ -58,13 +59,10 @@ def _check_tcp_ports(ip: str, ports: List[int], timeout: float = DEFAULT_TIMEOUT
                 connected_ports_sockets.append((port, sock))
                 possible_ports.append((port, sock))
                 continue
-            # BUG: I don't think a socket will ever connect successfully if this error is raised.
-            #      From the documentation: "Resource temporarily unavailable... It is a nonfatal
-            #      error, **and the operation should be retried later**." (emphasis mine). If the
-            #      operation is not retried later, I don't see the point in appending this to
-            #      possible_ports.
-            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
+            if err == 10035:  # WSAEWOULDBLOCK is valid.
+                # https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-connect
+                # says, "Use the select function to determine the completion of the connection
+                # request by checking to see if the socket is writable," which is being done below.
                 possible_ports.append((port, sock))
                 continue
             if err == 115:  # EINPROGRESS     115     /* Operation now in progress */
@@ -74,15 +72,11 @@ def _check_tcp_ports(ip: str, ports: List[int], timeout: float = DEFAULT_TIMEOUT
 
         if len(possible_ports) != 0:
             timeout = int(round(timeout))  # clamp to integer, to avoid checking input
-            sockets_to_try = possible_ports[:]
-            # BUG: If any sockets were added to connected_ports_sockets on line 94, this would
-            #      remove them.
-            connected_ports_sockets = []
+            sockets_to_try = possible_ports.copy()
             while (timeout >= 0) and sockets_to_try:
                 sock_objects = [s[1] for s in sockets_to_try]
 
-                # BUG: Since timeout is 0, this could block indefinitely
-                _, writeable_sockets, _ = select.select(sock_objects, sock_objects, sock_objects, 0)
+                _, writeable_sockets, _ = select.select([], sock_objects, [], 0)
                 for s in writeable_sockets:
                     try:  # actual test
                         connected_ports_sockets.append((s.getpeername()[1], s))
@@ -97,6 +91,7 @@ def _check_tcp_ports(ip: str, ports: List[int], timeout: float = DEFAULT_TIMEOUT
                 "On host %s discovered the following ports %s"
                 % (str(ip), ",".join([str(s[0]) for s in connected_ports_sockets]))
             )
+
             banners = []
             if len(connected_ports_sockets) != 0:
                 readable_sockets, _, _ = select.select(
@@ -104,20 +99,18 @@ def _check_tcp_ports(ip: str, ports: List[int], timeout: float = DEFAULT_TIMEOUT
                 )
                 # read first BANNER_READ bytes. We ignore errors because service might not send a
                 # decodable byte string.
-                # CR: Because of how black formats this, it is difficult to parse. Refactor to be
-                # easier to read.
+                for port, sock in connected_ports_sockets:
+                    if sock in readable_sockets:
+                        banners.append(sock.recv(BANNER_READ).decode(errors="ignore"))
+                    else:
+                        banners.append("")
 
-                # TODO: Rework the return of this function. Consider using dictionary
-                banners = [
-                    sock.recv(BANNER_READ).decode(errors="ignore")
-                    if sock in readable_sockets
-                    else ""
-                    for port, sock in connected_ports_sockets
-                ]
-                pass
             # try to cleanup
-            # CR: Evaluate whether or not we should call shutdown() before close() on each socket.
-            [s[1].close() for s in possible_ports]
+            for s in possible_ports:
+                s[1].shutdown(socket.SHUT_RDWR)
+                s[1].close()
+
+            # TODO: Rework the return of this function. Consider using dictionary
             return [port for port, sock in connected_ports_sockets], banners
         else:
             return [], []