mitmproxy/libmproxy/protocol/tcp.py

57 lines
2.4 KiB
Python

from . import ProtocolHandler
import select, socket
from cStringIO import StringIO
class TCPHandler(ProtocolHandler):
"""
TCPHandler acts as a generic TCP forwarder.
Data will be .log()ed, but not stored any further.
"""
def handle_messages(self):
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
src_str, dst_str = "client", "server"
else:
dst, src = self.c.client_conn, self.c.server_conn
dst_str, src_str = "client", "server"
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)
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" % (src_str, dst_str), ["\r\n" + data])
dst.wfile.write(data)
dst.wfile.flush()