61 lines
2.6 KiB
Python
61 lines
2.6 KiB
Python
from __future__ import absolute_import
|
|
import select, socket
|
|
from cStringIO import StringIO
|
|
from .primitives import ProtocolHandler
|
|
|
|
|
|
class TCPHandler(ProtocolHandler):
|
|
"""
|
|
TCPHandler acts as a generic TCP forwarder.
|
|
Data will be .log()ed, but not stored any further.
|
|
"""
|
|
def handle_messages(self):
|
|
self.c.establish_server_connection()
|
|
conns = [self.c.client_conn.rfile, self.c.server_conn.rfile]
|
|
while not self.c.close:
|
|
r, _, _ = select.select(conns, [], [], 10)
|
|
for rfile in r:
|
|
if self.c.client_conn.rfile == rfile:
|
|
src, dst = self.c.client_conn, self.c.server_conn
|
|
direction = "-> tcp ->"
|
|
dst_str = "%s:%s" % self.c.server_conn.address()[:2]
|
|
else:
|
|
dst, src = self.c.client_conn, self.c.server_conn
|
|
direction = "<- tcp <-"
|
|
dst_str = "client"
|
|
|
|
data = StringIO()
|
|
while range(4096):
|
|
# Do non-blocking select() to see if there is further data on in the buffer.
|
|
r, _, _ = select.select([rfile], [], [], 0)
|
|
if len(r):
|
|
d = rfile.read(1)
|
|
if d == "": # connection closed
|
|
break
|
|
data.write(d)
|
|
# OpenSSL Connections have an internal buffer that might
|
|
# contain data altough everything is read from the socket.
|
|
# Thankfully, connection.pending() returns the amount of
|
|
# bytes in this buffer, so we can read it completely at
|
|
# once.
|
|
if src.ssl_established:
|
|
data.write(rfile.read(src.connection.pending()))
|
|
else: # no data left, but not closed yet
|
|
break
|
|
data = data.getvalue()
|
|
|
|
if data == "": # no data received, rfile is closed
|
|
self.c.log("Close writing connection to %s" % dst_str, "debug")
|
|
conns.remove(rfile)
|
|
if dst.ssl_established:
|
|
dst.connection.shutdown()
|
|
else:
|
|
dst.connection.shutdown(socket.SHUT_WR)
|
|
if len(conns) == 0:
|
|
self.c.close = True
|
|
break
|
|
|
|
self.c.log("%s %s\r\n%s" % (direction, dst_str, data), "debug")
|
|
dst.wfile.write(data)
|
|
dst.wfile.flush()
|