add generic TCP handler with SSL support, move StateObject into netlib

1 files changed, 57 insertions, 0 deletions

diff --git a/libmproxy/protocol/tcp.py b/libmproxy/protocol/tcp.py
new file mode 100644
index 00000000..9fb2e01b
--- /dev/null
+++ b/libmproxy/protocol/tcp.py
@@ -0,0 +1,57 @@
+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()
\ No newline at end of file