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

Agent: Improve readability of _check_tcp_ports()

This commit is contained in:
Mike Salvatore 2022-02-10 10:17:29 -05:00 committed by Ilija Lazoroski
parent 2ae77ce897
commit 21ede3e341
1 changed files with 19 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
monkey/infection_monkey/network/tcp_scanner.py
View File

@ -53,8 +53,8 @@ def _check_tcp_ports(
for s in sockets: for s in sockets:
s.setblocking(False) s.setblocking(False)
possible_ports = [] possible_ports = set()
connected_ports_sockets = [] connected_ports = set()
open_ports = {} open_ports = {}
try: try:
@ -64,19 +64,17 @@ def _check_tcp_ports(
for sock, port in zip(sockets, ports_to_scan): for sock, port in zip(sockets, ports_to_scan):
err = sock.connect_ex((ip, port)) err = sock.connect_ex((ip, port))
if err == 0: # immediate connect if err == 0: # immediate connect
connected_ports_sockets.append((port, sock)) connected_ports.add((port, sock))
possible_ports.append((port, sock)) possible_ports.add((port, sock))
continue elif err == 10035: # WSAEWOULDBLOCK is valid.
if err == 10035: # WSAEWOULDBLOCK is valid.
# https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-connect # 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 # 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. # request by checking to see if the socket is writable," which is being done below.
possible_ports.append((port, sock)) possible_ports.add((port, sock))
continue elif err == 115: # EINPROGRESS 115 /* Operation now in progress */
if err == 115: # EINPROGRESS 115 /* Operation now in progress */ possible_ports.add((port, sock))
possible_ports.append((port, sock)) else:
continue logger.warning("Failed to connect to port %s, error code is %d", port, err)
logger.warning("Failed to connect to port %s, error code is %d", port, err)
if len(possible_ports) != 0: if len(possible_ports) != 0:
sockets_to_try = possible_ports.copy() sockets_to_try = possible_ports.copy()
@ -90,24 +88,25 @@ def _check_tcp_ports(
_, writeable_sockets, _ = select.select([], sock_objects, [], timer.time_remaining) _, writeable_sockets, _ = select.select([], sock_objects, [], timer.time_remaining)
for s in writeable_sockets: for s in writeable_sockets:
try: # actual test try: # actual test
connected_ports_sockets.append((s.getpeername()[1], s)) connected_ports.add((s.getpeername()[1], s))
except socket.error: # bad socket, select didn't filter it properly except socket.error: # bad socket, select didn't filter it properly
pass pass
sockets_to_try = [s for s in sockets_to_try if s not in connected_ports_sockets]
sockets_to_try = sockets_to_try - connected_ports
logger.debug( logger.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]))
) )
open_ports = {port: "" for port, _ in connected_ports_sockets} open_ports = {port: "" for port, _ in connected_ports}
if len(connected_ports_sockets) != 0: if len(connected_ports) != 0:
readable_sockets, _, _ = select.select( readable_sockets, _, _ = select.select(
[s[1] for s in connected_ports_sockets], [], [], timer.time_remaining [s[1] for s in connected_ports], [], [], timer.time_remaining
) )
# read first BANNER_READ bytes. We ignore errors because service might not send a # read first BANNER_READ bytes. We ignore errors because service might not send a
# decodable byte string. # decodable byte string.
for port, sock in connected_ports_sockets: for port, sock in connected_ports:
if sock in readable_sockets: if sock in readable_sockets:
open_ports[port] = sock.recv(BANNER_READ).decode(errors="ignore") open_ports[port] = sock.recv(BANNER_READ).decode(errors="ignore")
else: else:
@ -116,7 +115,7 @@ def _check_tcp_ports(
except socket.error as exc: except socket.error as exc:
logger.warning("Exception when checking ports on host %s, Exception: %s", str(ip), exc) logger.warning("Exception when checking ports on host %s, Exception: %s", str(ip), exc)
_clean_up_sockets(possible_ports, connected_ports_sockets) _clean_up_sockets(possible_ports, connected_ports)
return open_ports return open_ports