I made lots of changes to the script presented in my previous question. I was tempted to edit that one with the new code, but it would invalidate @200_success's helpful answer. It was also disallowed on a Meta thread. So I apologize for the new question, but it seemed like the right thing to do.
Again, any and all tips are appreciated! Also, this is my first time writing any docstrings, so if I'm breaking convention, let me know.
#!/usr/bin/env python3 import argparse import errno import functools import multiprocessing import os import platform import socket import time from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ProcessPoolExecutor DEFAULT_HOST = '127.0.0.1' DEFAULT_TIMEOUT = 1 DEFAULT_THREADS = 512 PORT_RANGE = range(1, 65536) def tcp_ping(host, port): """ Attempts to connect to host:port via TCP. Arguments: host: IP address or URL to hit port: Port to hit Returns: port number, if it's available; otherwise False """ try: with socket.socket() as sock: sock.connect((host, port)) print(str(port) + ' Open') return port except socket.timeout: return False except socket.error as socket_error: if socket_error.errno == errno.ECONNREFUSED: return False raise def perform_scan(host, use_threads = False): """ Perform a scan on all valid ports (1 - 65535), either by spawning threads or forking processes. Arguments: host: IP address or URL to scan use_threads: whether or not to use threads; default behaviour is to fork processes Returns: list of available ports """ if use_threads: executor = ThreadPoolExecutor(max_workers = DEFAULT_THREADS) else: executor = ProcessPoolExecutor() with executor: ping_partial = functools.partial(tcp_ping, host) return list(filter(bool, executor.map(ping_partial, PORT_RANGE))) def is_address_valid(host): """ Validate the host's existence by attempting to establish a connection to it on port 80 (HTTP). Arguments: host: IP address or URL to validate Returns: bool indicating whether the host is valid """ try: with socket.socket() as sock: sock.connect((host, 80)) return True except socket.gaierror: return False except (socket.timeout, socket.error): return True def scan_ports(host = DEFAULT_HOST, timeout = DEFAULT_TIMEOUT): """ Scan all possible ports on the specified host. If the operating system is detected as Windows, the ports will be scanned by spawning threads. Otherwise, new processes will be forked. Arguments: host: IP address or URL to scan timeout: connection timeout when testing a port """ # Set defaults if CLI arguments were passed in but not specified if host is None: host = DEFAULT_HOST if timeout is None: timeout = DEFAULT_TIMEOUT # Strip the protocol from the URL, if present if '://' in host: host = host[host.index('://') + 3:] # Set the timeout for all subsequent connections socket.setdefaulttimeout(timeout) # Validate the IP/host by testing a connection if not is_address_valid(host): print('DNS lookup for \'' + host + '\' failed.') return # Perform the scan. On Windows, thread. On all others, fork. print('Scanning ' + host + ' ...') start_time = time.clock() if os.name == 'nt': print('Running on Windows OS.') available_ports = perform_scan(host, use_threads = True) elif os.name == 'posix': print('Running on *Nix OS.') available_ports = perform_scan(host) else: print('Unidentified operating system: ' + os.name + ' [' + platform.system() + ' ' + platform.version() + ']') available_ports = perform_scan(host) end_time = time.clock() print('Done.') # Display results print() print('Time elapsed: ' + format(end_time - start_time, '.2f') + ' sec') available_ports.sort() print() print(str(len(available_ports)) + ' ports available.') print(available_ports) def main(): arg_parser = argparse.ArgumentParser() arg_parser.add_argument('-ip', '--host', help = 'IP address/host to scan') arg_parser.add_argument('-t', '--timeout', help = 'Connection timeout in seconds', type = int) args = arg_parser.parse_args() scan_ports(args.host, args.timeout) if __name__ == '__main__': main() 
