12
\$\begingroup\$

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()
\$\endgroup\$
  • 12
    \$\begingroup\$ Asking a new question is the right thing to do! \$\endgroup\$ – 200_success Apr 11 '14 at 6:27
4
\$\begingroup\$

Trying to connect to TCP port 80 in order to check whether the DNS lookup succeeds is overkill. You should just call socket.gethostbyname().

\$\endgroup\$
3
\$\begingroup\$

I notice a few little things that I would adjust just from a housekeeping perspective.

  1. You import multiprocessing but I don't see where it is used.. I would remove if it is not needed.
  2. some small pep8 stuff - Keyword args should not have spaces around the equal, 2 blank lines between each function, and more than 80 line length on the argpase.add_arg near the bottom.
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.