tcp_ping.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import socket
import logging
import sys

log = logging.getLogger(__name__)


def tcp_ping(host, port, timeout=1.0, ipv6=False):
    """Test a TCP connection to a given host-port pair."""
    if ipv6:
        log.debug('using IPv6')
        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
    else:
        log.debug('using IPv4')
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    with sock:
        log.debug('setting socket timeout to %f', timeout)
        sock.settimeout(timeout)
        try:
            log.debug('connecting socket')
            sock.connect((host, int(port)))
            log.debug('shutting down socket')
            sock.shutdown(socket.SHUT_RD)
            return True
        except (ConnectionError, socket.timeout) as error:
            log.warning(error)
            return False


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('host',
                        help='host IP or domain name to ping')
    parser.add_argument('port',
                        type=int,
                        help='host port to ping')
    parser.add_argument('-t',
                        '--timeout',
                        type=float,
                        default=1.0,
                        metavar='SEC',
                        help='timeout in seconds (default: 1.0)')
    ip_grp = parser.add_mutually_exclusive_group()
    ip_grp.add_argument('-4',
                        action='store_false',
                        dest='ipv6',
                        default=False,
                        help='use IPv4 (default)')
    ip_grp.add_argument('-6',
                        action='store_true',
                        dest='ipv6',
                        help='use IPv6')
    log_group = parser.add_mutually_exclusive_group()
    log_group.add_argument(
        '-v',
        '--verbose',
        dest='verbosity',
        default=[logging.WARNING],
        action='append_const',
        const=-10,
        help='be more verbose',
    )
    log_group.add_argument(
        '-q',
        '--quiet',
        dest='verbosity',
        action='append_const',
        const=10,
        help='be more quiet',
    )
    args = parser.parse_args()
    log_level = max(logging.DEBUG, min(logging.CRITICAL, sum(args.verbosity)))
    logging.basicConfig(level=log_level)

    host = args.host
    port = args.port
    timeout = args.timeout
    ipv6 = args.ipv6

    if tcp_ping(host, port, timeout=timeout, ipv6=ipv6):
        print('Port {} is open on host {!r}'.format(port, host))
        sys.exit(0)
    else:
        print('Port {} is closed on host {!r}'.format(port, host))
        sys.exit(1)


if __name__ == '__main__':
    main()