From b8cfa343aeeb22f2fd0ddfa20e3b97f010f63de4 Mon Sep 17 00:00:00 2001 From: JonnyWong16 Date: Sat, 23 Nov 2019 19:01:14 -0800 Subject: [PATCH] Add portend-2.6 --- lib/portend.py | 222 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 222 insertions(+) create mode 100644 lib/portend.py diff --git a/lib/portend.py b/lib/portend.py new file mode 100644 index 00000000..5aefa725 --- /dev/null +++ b/lib/portend.py @@ -0,0 +1,222 @@ +# -*- coding: utf-8 -*- + +""" +A simple library for managing the availability of ports. +""" + +from __future__ import print_function, division + +import time +import socket +import argparse +import sys +import itertools +import contextlib +import platform + +try: + from collections import abc +except ImportError: + import collections as abc + +from tempora import timing + + +def client_host(server_host): + """Return the host on which a client can connect to the given listener.""" + if server_host == '0.0.0.0': + # 0.0.0.0 is INADDR_ANY, which should answer on localhost. + return '127.0.0.1' + if server_host in ('::', '::0', '::0.0.0.0'): + # :: is IN6ADDR_ANY, which should answer on localhost. + # ::0 and ::0.0.0.0 are non-canonical but common + # ways to write IN6ADDR_ANY. + return '::1' + return server_host + + +class Checker(object): + def __init__(self, timeout=1.0): + self.timeout = timeout + + def assert_free(self, host, port=None): + """ + Assert that the given addr is free + in that all attempts to connect fail within the timeout + or raise a PortNotFree exception. + + >>> free_port = find_available_local_port() + + >>> Checker().assert_free('localhost', free_port) + >>> Checker().assert_free('127.0.0.1', free_port) + >>> Checker().assert_free('::1', free_port) + + Also accepts an addr tuple + + >>> addr = '::1', free_port, 0, 0 + >>> Checker().assert_free(addr) + + Host might refer to a server bind address like '::', which + should use localhost to perform the check. + + >>> Checker().assert_free('::', free_port) + """ + if port is None and isinstance(host, abc.Sequence): + host, port = host[:2] + if platform.system() == 'Windows': + host = client_host(host) + info = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM) + list(itertools.starmap(self._connect, info)) + + def _connect(self, af, socktype, proto, canonname, sa): + s = socket.socket(af, socktype, proto) + # fail fast with a small timeout + s.settimeout(self.timeout) + + with contextlib.closing(s): + try: + s.connect(sa) + except socket.error: + return + + # the connect succeeded, so the port isn't free + port, host = sa[:2] + tmpl = "Port {port} is in use on {host}." + raise PortNotFree(tmpl.format(**locals())) + + +class Timeout(IOError): + pass + + +class PortNotFree(IOError): + pass + + +def free(host, port, timeout=float('Inf')): + """ + Wait for the specified port to become free (dropping or rejecting + requests). Return when the port is free or raise a Timeout if timeout has + elapsed. + + Timeout may be specified in seconds or as a timedelta. + If timeout is None or ∞, the routine will run indefinitely. + + >>> free('localhost', find_available_local_port()) + """ + if not host: + raise ValueError("Host values of '' or None are not allowed.") + + timer = timing.Timer(timeout) + + while True: + try: + # Expect a free port, so use a small timeout + Checker(timeout=0.1).assert_free(host, port) + return + except PortNotFree: + if timer.expired(): + raise Timeout("Port {port} not free on {host}.".format(**locals())) + # Politely wait. + time.sleep(0.1) + + +wait_for_free_port = free + + +def occupied(host, port, timeout=float('Inf')): + """ + Wait for the specified port to become occupied (accepting requests). + Return when the port is occupied or raise a Timeout if timeout has + elapsed. + + Timeout may be specified in seconds or as a timedelta. + If timeout is None or ∞, the routine will run indefinitely. + + >>> occupied('localhost', find_available_local_port(), .1) + Traceback (most recent call last): + ... + Timeout: Port ... not bound on localhost. + """ + if not host: + raise ValueError("Host values of '' or None are not allowed.") + + timer = timing.Timer(timeout) + + while True: + try: + Checker(timeout=0.5).assert_free(host, port) + if timer.expired(): + raise Timeout("Port {port} not bound on {host}.".format(**locals())) + # Politely wait + time.sleep(0.1) + except PortNotFree: + # port is occupied + return + + +wait_for_occupied_port = occupied + + +def find_available_local_port(): + """ + Find a free port on localhost. + + >>> 0 < find_available_local_port() < 65536 + True + """ + infos = socket.getaddrinfo(None, 0, socket.AF_UNSPEC, socket.SOCK_STREAM) + family, proto, _, _, addr = next(iter(infos)) + sock = socket.socket(family, proto) + sock.bind(addr) + addr, port = sock.getsockname()[:2] + sock.close() + return port + + +class HostPort(str): + """ + A simple representation of a host/port pair as a string + + >>> hp = HostPort('localhost:32768') + + >>> hp.host + 'localhost' + + >>> hp.port + 32768 + + >>> len(hp) + 15 + """ + + @property + def host(self): + host, sep, port = self.partition(':') + return host + + @property + def port(self): + host, sep, port = self.partition(':') + return int(port) + + +def _main(): + parser = argparse.ArgumentParser() + + def global_lookup(key): + return globals()[key] + + parser.add_argument('target', metavar='host:port', type=HostPort) + parser.add_argument('func', metavar='state', type=global_lookup) + parser.add_argument('-t', '--timeout', default=None, type=float) + args = parser.parse_args() + try: + args.func(args.target.host, args.target.port, timeout=args.timeout) + except Timeout as timeout: + print(timeout, file=sys.stderr) + raise SystemExit(1) + + +if __name__ == '__main__': + _main()