Hide navigation sidebar
Hide table of contents sidebar
Back to top
Toggle table of contents sidebar

# 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