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
                    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)
                    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))
                dst.wfile.write(data)
                dst.wfile.flush()