tcp-proxy.py

import sys

from twisted.internet import defer, protocol, task, endpoints
from twisted.python import log


class ProxyClientProtocol(protocol.Protocol):
    def connectionMade(self):
        log.msg("Client: connected to peer")

    def dataReceived(self, data):
        log.msg("Client: %d bytes received from peer" % (len(data),))
        self.factory.peer.transport.write(data)

    def connectionLost(self, reason):
        log.err("Client: peer disconnected; closing server's peer", reason)
        self.factory.peer.transport.loseConnection()


class ProxyClientFactory(protocol.ClientFactory):
    protocol = ProxyClientProtocol

    def __init__(self, peer):
        self.peer = peer


class ProxyServerProtocol(protocol.Protocol):
    def connectionMade(self):
        self._queuedData = []
        self.peer = None
        (endpoints.TCP4ClientEndpoint(self.factory.reactor, '127.0.0.1', 6666)
         .connect(ProxyClientFactory(self))
         .addCallbacks(self._gotPeer, self._peerConnectionFailed))

    def _gotPeer(self, peer):
        self.peer = peer
        self.peer.transport.writeSequence(self._queuedData)
        self._queuedData = None

    def _peerConnectionFailed(self, reason):
        log.err("Server: couldn't connect to peer", reason)
        self.transport.loseConnection()

    def dataReceived(self, data):
        log.msg("Server: %d bytes received from peer" % (len(data),))
        if self.peer is None:
            self._queuedData.append(data)
        else:
            self.peer.transport.write(data)

    def connectionLost(self, reason):
        log.err("Server: peer disconnected", reason)
        if self.peer is not None:
            log.msg("Server: closing client's peer")
            self.peer.transport.loseConnection()


class ProxyServerFactory(protocol.Factory):
    protocol = ProxyServerProtocol

    def __init__(self, reactor):
        self.reactor = reactor


def main(reactor):
    log.startLogging(sys.stderr)
    endpoints.TCP4ServerEndpoint(reactor, 9999).listen(ProxyServerFactory(reactor))
    return defer.Deferred()

task.react(main, [])