relay.py

"""Port forwarder with graceful exit.
Run the example as
  python portforwarder.py :8080 gevent.org:80
Then direct your browser to http://localhost:8080 or do "telnet localhost 8080".
When the portforwarder receives TERM or INT signal (type Ctrl-C),
it closes the listening socket and waits for all existing
connections to finish. The existing connections will remain unaffected.
The program will exit once the last connection has been closed.
"""
import sys
import signal
import gevent
import socket
from gevent.server import StreamServer
from gevent.socket import create_connection, gethostbyname
import logging
import logging.config

LOGGING = {
    'version': 1,
    'formatters': {
        'standard': {
            'format': '%(asctime)s|%(levelname)s|%(process)d|%(filename)s:%(lineno)d|%(module)s.%(funcName)s|%(message)s',
                    },
                   },
    'handlers': {
        'main': {
            'class': 'logging.handlers.RotatingFileHandler',
            'filename': 'log',
            'maxBytes': 1024*1024*50,
            'formatter': 'standard',
            'backupCount': 10,
        }
    },
    'loggers': {
        'main': {
            'level': 'INFO',
            'handlers': ['main']
        },
    },
}

logging.config.dictConfig(LOGGING)
logger = logging.getLogger('main')

class PortForwarder(StreamServer):

    def __init__(self, listener, dest, **kwargs):
        StreamServer.__init__(self, listener, **kwargs)
        self.dest = dest

    def handle(self, source, address):
        logger.info('%s:%s accepted', *address[:2])
        try:
            dest = create_connection(self.dest)
        except IOError as ex:
            logger.error('%s:%s failed to connect to %s:%s: %s', address[0], address[1], self.dest[0], self.dest[1], ex)
            return

        forwarders = (gevent.spawn(forward, source, dest),
                      gevent.spawn(forward, dest, source))
        # if we return from this method, the stream will be closed out
        # from under us, so wait for our children
        gevent.joinall(forwarders)


    def close(self):
        if self.closed:
            sys.exit('Multiple exit signals received - aborting.')
        else:
            logger.info('Closing listener socket')
            StreamServer.close(self)


def forward(source, dest):
    #source_address = '%s:%s' % source.getpeername()[:2]
    #dest_address = '%s:%s' % dest.getpeername()[:2]

    try:
        while True:
            data = source.recv(4096)
            #logger.debug('%s->%s: %r', source_address, dest_address, data)
            if not data:
                break
            dest.sendall(data)
    except Exception as e:
        pass
    finally:
        source.close()
        dest.close()

def parse_address(address):
    try:
        hostname, port = address.rsplit(':', 1)
        port = int(port)
    except ValueError:
        sys.exit('Expected HOST:PORT: %r' % address)
    return gethostbyname(hostname), port


def main():
    args = sys.argv[1:]
    if len(args) != 2:
        sys.exit('Usage: %s source-address destination-address' % __file__)
    source = args[0]
    dest = parse_address(args[1])
    server = PortForwarder(source, dest)
    logger.info('Starting port forwarder %s:%s -> %s:%s', *(server.address[:2] + dest))
    gevent.signal(signal.SIGTERM, server.close)
    gevent.signal(signal.SIGINT, server.close)
    server.start()
    gevent.wait()



if __name__ == '__main__':
    main()