monkey/chaos_monkey/transport/tcp.py

import socket
import select
from threading import Thread
from base import TransportProxyBase, update_last_serve_time
from logging import getLogger

READ_BUFFER_SIZE = 8192
DEFAULT_TIMEOUT = 30

LOG = getLogger(__name__)


class SocketsPipe(Thread):
    def __init__(self, source, dest, timeout=DEFAULT_TIMEOUT):
        Thread.__init__(self)
        self.source = source
        self.dest = dest
        self.timeout = timeout
        self._keep_connection = True
        super(SocketsPipe, self).__init__()
        self.daemon = True

    def run(self):
        sockets = [self.source, self.dest]
        while self._keep_connection:
            self._keep_connection = False
            rlist, wlist, xlist = select.select(sockets, [], sockets, self.timeout)
            if xlist:
                break
            for r in rlist:
                other = self.dest if r is self.source else self.source
                try:
                    data = r.recv(READ_BUFFER_SIZE)
                except:
                    break
                if data:
                    try:
                        other.sendall(data)
                        update_last_serve_time()
                    except:
                        break
                    self._keep_connection = True
        
        self.source.close()
        self.dest.close()


class TcpProxy(TransportProxyBase):
    
    def run(self):
        pipes = []
        l_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        l_socket.bind((self.local_host, self.local_port))
        l_socket.settimeout(DEFAULT_TIMEOUT)
        l_socket.listen(5)

        while not self._stopped:
            try:
                source, address = l_socket.accept()
            except socket.timeout:
                continue

            try:
                dest = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                dest.connect((self.dest_host, self.dest_port))
            except socket.error, ex:
                source.close()
                dest.close()
                continue

            pipe = SocketsPipe(source, dest)
            pipes.append(pipe)
            LOG.debug("piping sockets %s:%s->%s:%s", address[0], address[1], self.dest_host, self.dest_port)
            pipe.start()

        l_socket.close()
        for pipe in pipes:
            pipe.join()