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()