Example portforwarder.py

  1"""Port forwarder with graceful exit.
  2
  3Run the example as
  4
  5  python portforwarder.py :8080 gevent.org:80
  6
  7Then direct your browser to http://localhost:8080 or do "telnet localhost 8080".
  8
  9When the portforwarder receives TERM or INT signal (type Ctrl-C),
 10it closes the listening socket and waits for all existing
 11connections to finish. The existing connections will remain unaffected.
 12The program will exit once the last connection has been closed.
 13"""
 14import socket
 15import sys
 16import signal
 17import gevent
 18from gevent.server import StreamServer
 19from gevent.socket import create_connection, gethostbyname
 20
 21
 22class PortForwarder(StreamServer):
 23
 24    def __init__(self, listener, dest, **kwargs):
 25        StreamServer.__init__(self, listener, **kwargs)
 26        self.dest = dest
 27
 28    def handle(self, source, address): # pylint:disable=method-hidden
 29        log('%s:%s accepted', *address[:2])
 30        try:
 31            dest = create_connection(self.dest)
 32        except IOError as ex:
 33            log('%s:%s failed to connect to %s:%s: %s', address[0], address[1], self.dest[0], self.dest[1], ex)
 34            return
 35        forwarders = (gevent.spawn(forward, source, dest, self),
 36                      gevent.spawn(forward, dest, source, self))
 37        # if we return from this method, the stream will be closed out
 38        # from under us, so wait for our children
 39        gevent.joinall(forwarders)
 40
 41    def close(self):
 42        if self.closed:
 43            sys.exit('Multiple exit signals received - aborting.')
 44        else:
 45            log('Closing listener socket')
 46            StreamServer.close(self)
 47
 48
 49def forward(source, dest, server):
 50    try:
 51        source_address = '%s:%s' % source.getpeername()[:2]
 52        dest_address = '%s:%s' % dest.getpeername()[:2]
 53    except socket.error as e:
 54        # We could be racing signals that close the server
 55        # and hence a socket.
 56        log("Failed to get all peer names: %s", e)
 57        return
 58
 59    try:
 60        while True:
 61            try:
 62                data = source.recv(1024)
 63                log('%s->%s: %r', source_address, dest_address, data)
 64                if not data:
 65                    break
 66                dest.sendall(data)
 67            except KeyboardInterrupt:
 68                # On Windows, a Ctrl-C signal (sent by a program) usually winds
 69                # up here, not in the installed signal handler.
 70                if not server.closed:
 71                    server.close()
 72                break
 73            except socket.error:
 74                if not server.closed:
 75                    server.close()
 76                break
 77    finally:
 78        source.close()
 79        dest.close()
 80        server = None
 81
 82
 83def parse_address(address):
 84    try:
 85        hostname, port = address.rsplit(':', 1)
 86        port = int(port)
 87    except ValueError:
 88        sys.exit('Expected HOST:PORT: %r' % address)
 89    return gethostbyname(hostname), port
 90
 91
 92def main():
 93    args = sys.argv[1:]
 94    if len(args) != 2:
 95        sys.exit('Usage: %s source-address destination-address' % __file__)
 96    source = args[0]
 97    dest = parse_address(args[1])
 98    server = PortForwarder(source, dest)
 99    log('Starting port forwarder %s:%s -> %s:%s', *(server.address[:2] + dest))
100    gevent.signal_handler(signal.SIGTERM, server.close)
101    gevent.signal_handler(signal.SIGINT, server.close)
102    server.start()
103    gevent.wait()
104
105
106def log(message, *args):
107    message = message % args
108    sys.stderr.write(message + '\n')
109
110
111if __name__ == '__main__':
112    main()

Current source