Avoid global variable due to executor.map()

Question

I'm pretty new to python but after watching Mr Robot I felt inspired to make a CLI portscanner. With all the fancy console outputs you see in movies like a progress bar.

It's finished and works nicely however I am not very happy with having to write twice global host_IP in main() and scan_tcp(). I have to do this the only way I was able to get a progress bar with tqdm was by using executor.map() which as to my understanding can only take the iterator argument unlucky the executor.submit(). Is there a cleaner way to doing this?

I am also interested in any general feedback on the code or performance improvements :)

import socket
import sys
import argparse
import concurrent.futures
import pyfiglet
from tqdm import tqdm

host_IP = None


def resolve_args():
    argv = sys.argv[1:]
    parser = argparse.ArgumentParser()
    parser.add_argument("host", type=_host, help="Hostname or IP of host system to be scanned")
    parser.add_argument("-s", "--startPort", type=_port, default=0, help="Port number to start scan (0-65535)")
    parser.add_argument("-e", "--endPort", type=_port, default=65535, help="Port number to end scan (0-65535)")
    return vars(parser.parse_args(argv))


def _host(s):
    try:
        value = socket.gethostbyname(s)
    except socket.gaierror:
        raise argparse.ArgumentTypeError(f"Host '{s}' could not be resolved.")
    return value


def _port(s):
    try:
        value = int(s)
    except ValueError:
        raise argparse.ArgumentTypeError(f"Expected integer got '{s}'")
    if 0 > value > 65535:
        raise argparse.ArgumentTypeError(f"Port number must be 0-65535, got {s}")
    return value


def scan_tcp(port):
    global host_IP
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    socket.setdefaulttimeout(1)
    result = 0
    try:
        answer = sock.connect_ex((host_IP, port))
        if answer == 0:
            result = port
        sock.close()
    except socket.error:
        print(f"Could not connect to host '{host_IP}'")
        sys.exit()
    except KeyboardInterrupt:
        print("Exiting...")
        sys.exit()

    return result


def print_open_ports(results):
    results = list(filter(None, results))
    for open_port in results:
        print(f"Port {open_port} is OPEN")


def main():
    args = resolve_args()
    pyfiglet.print_figlet("PORTSCANNER", font="slant")
    global host_IP
    host_IP = args.get("host")
    ports = range(args.get("startPort"), args.get("endPort") + 1)
    print(f"Starting scan of {host_IP}...")
    with concurrent.futures.ThreadPoolExecutor() as executor:
        results = list(tqdm(executor.map(scan_tcp, ports), total=len(ports)))
    print(f"Finished scan of {host_IP}!")
    print_open_ports(results)


if __name__ == "__main__":
    main()

Answer 1

It's finished and works nicely however I am not very happy with having to write twice global host_IP in main() and scan_tcp().

I don't really understand the issue. The variable already has global scope since you assign it a None value on top of your code. Then all you have to do is pass it as a function argument, for example (preferably in lowercase notation):

def scan_tcp(host_ip, port):

But this is just variable duplication, you could downright pass args.get("host") or args.host if you add this: dest="host" in the relevant add_argument statement.

scan_tcp should probably return a boolean value and be named something else eg get_port_state. If the connection is successful you return the port number, otherwise 0 - this is not very consistent and even misleading since 0 is a port too. All you want is a yes/no answer.

Regarding the parsing of arguments you don't need this remnant of the old ways:

argv = sys.argv[1:]

This is sufficient:

parser = argparse.ArgumentParser()
parser.add_argument(...)
parser.add_argument(...)
args = parser.parse_args()

Your script will refuse to run if the host cannot be resolved. Why ? Not every IP address can be resolved to a host name. I would ditch that check.

The approach you've chosen (using socket) is quite rudimentary and is not a reliable test. Because determining what makes a port open or not is not that easy to tell reliably. For your next version you might consider nmap since it can be instrumented in Python and provides many more options. Maybe this topic could interest you: A port scanner made in python. Be aware that even nmap makes educated guesses, and I am not even touching on intrusion detection systems (IDS) or any other measures that may interfere with your scanning.

And of course scapy should definitely interest you.

I am not familiar enough with tqdm.

PS: an interesting article on using socket in Python: Socket Programming in Python (Guide)

Answer 2

Your

try:
    value = socket.gethostbyname(s)
except socket.gaierror:
    raise argparse.ArgumentTypeError(f"Host '{s}' could not be resolved.")
return value

should probably be

try:
    return socket.gethostbyname(s)
except socket.gaierror as e:
    raise argparse.ArgumentTypeError(f"Host '{s}' could not be resolved.") from e

in other words, it's usually a good idea to maintain rather than break the exception cause chain.

I would move your except KeyboardInterrupt up to a try in main, and rather than exiting, just pass. That way testers or external callers of scan_tcp can make their own decisions about how to handle keyboard interrupts.

Why build up a list of results and defer printing until the end? Why not print results as you find them?