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