python twisted man-in-the-middle implementation

Question

What I need is a sort of man-in-the-middle implementation: I need a server who receives connections from clients (binary data with different lengths) and forwards the stream to a server it connects to (acting as a client), and then sends the data back from the server it is connected to, to the clients.

It actually works standing between the clients and the servers, and passing the data they exchange (which is a stream, so it continuously get from one side and sends to the other one).
The server is static, so it is always the same, and its address can even be hardcoded; however when a client drops the connection, this server must also drop the connection to the "real" server.

I've been looking around, but couldn't find a solution or an example for such a simple problem.

The code I've made works actually, but I have not yet managed to find how to put a reference into the server part that says "this is your assigned client", or into the client that says "this is your server". Here's my code:

#!/usr/bin/env python

from twisted.internet import protocol, reactor
from twisted.protocols import basic

client = None
server = None

class ServerProtocol(protocol.Protocol):
    def connectionMade(self):
        global server
        factory = protocol.ClientFactory()
        factory.protocol = ClientProtocol
        server = self
        reactor.connectTCP('localhost', 1324, factory)

    def dataReceived(self, data):
        global client
        client.transport.write(data)

class ClientProtocol(protocol.Protocol):
    def connectionMade(self):
        global client
        # Here's the instance of the client
        client = self

    def dataReceived(self, data):
        global server
        server.transport.write(data)

def main():
    import sys
    from twisted.python import log
    log.startLogging(sys.stdout)
    factory = protocol.ServerFactory()
    factory.protocol = ServerProtocol

    # Here's the instance of the server
    server = ServerProtocol

    reactor.listenTCP(2593, factory)
    reactor.run()

if __name__ == '__main__':
    main()

Now, the point is that the instance can't be contained into the global objects, and should be put inside the two classes: how?

Answer 1

I've managed to solve the issue by myself and, for future references (or to help anybody else who had this problem), here's the code I used to solve it.

I think both my solution and the one kindly given by jedwards work; now I just have to study his own a little more to be sure that what I've done is correct: this is my first application using the Twisted framework and studying somebody else's solution is the way to learn something new! :)

#!/usr/bin/env python

from twisted.internet import protocol, reactor
from twisted.protocols import basic

class ServerProtocol(protocol.Protocol):
    def __init__(self):
        self.buffer = None
        self.client = None

    def connectionMade(self):
        factory = protocol.ClientFactory()
        factory.protocol = ClientProtocol
        factory.server = self

        reactor.connectTCP('gameserver16.gamesnet.it', 2593, factory)

    def dataReceived(self, data):
        if (self.client != None):
            self.client.write(data)
        else:
            self.buffer = data

    def write(self, data):
        self.transport.write(data)
        print 'Server: ' + data.encode('hex')

class ClientProtocol(protocol.Protocol):
    def connectionMade(self):
        self.factory.server.client = self
        self.write(self.factory.server.buffer)
        self.factory.server.buffer = ''

    def dataReceived(self, data):
        self.factory.server.write(data)

    def write(self, data):
        self.transport.write(data)
        print 'Client: ' + data.encode('hex')

def main():
    import sys
    from twisted.python import log

    log.startLogging(sys.stdout)

    factory = protocol.ServerFactory()
    factory.protocol = ServerProtocol

    reactor.listenTCP(2593, factory)
    reactor.run()

if __name__ == '__main__':
    main()

Answer 2

Consider this approach

#!/usr/bin/env python

import sys
from twisted.internet import reactor
from twisted.internet.protocol import ServerFactory, ClientFactory, Protocol
from twisted.protocols import basic
from twisted.python import log

LISTEN_PORT = 2593
SERVER_PORT = 1234


class ServerProtocol(Protocol):
    def connectionMade(self):
        reactor.connectTCP('localhost', SERVER_PORT, MyClientFactory(self))

    def dataReceived(self, data):
        self.clientProtocol.transport.write(data)

class ClientProtocol(Protocol):
    def connectionMade(self):
        # Pass ServerProtocol a ref. to ClientProtocol
        self.serverProtocol.clientProtocol = self;  

    def dataReceived(self, data):
        self.serverProtocol.transport.write(data)

class MyServerFactory(ServerFactory):
    protocol = ServerProtocol
    def buildProtocol(self, addr):
        # Create ServerProtocol
        p = ServerFactory.buildProtocol(self, addr)
        return p

class MyClientFactory(ClientFactory):
    protocol = ClientProtocol
    def __init__(self, serverProtocol_):
        self.serverProtocol = serverProtocol_

    def buildProtocol(self, addr):
        # Create ClientProtocol
        p = ClientFactory.buildProtocol(self,addr)
        # Pass ClientProtocol a ref. to ServerProtocol
        p.serverProtocol = self.serverProtocol
        return p

def main():
    log.startLogging(sys.stdout)

    reactor.listenTCP(LISTEN_PORT, MyServerFactory())
    reactor.run()

if __name__ == '__main__':
    main()

The ServerProtcol instance, passes a reference of itself to the MyClientFactory constructor, which then sets tells the ClientProtcol what ServerProtocol instance it's associated with.

Similarly, when the ClientProtocol connection is established, it uses it's reference to the ServerProtocol to tell the ServerProtocol what ClientProtocol to use.

Note: There's no error checking in this code, so you may encounter errors regarding NoneType if things go wrong (for example, if the real server isn't listening).

The important lines are:

reactor.connectTCP('localhost', SERVER_PORT, MyClientFactory(self))
#...
def __init__(self, serverProtocol_):
    self.serverProtocol = serverProtocol_

Here, you pass a reference to the ServerProtocol to the MyClientFactory constructor. It stores this reference in a member variable. You do this so that that when the client factory creates a ClientProtocol, it can pass the reference on:

# Pass ClientProtocol a ref. to ServerProtocol
p.serverProtocol = self.serverProtocol

Then, once the connection is made from your script to the real server, the reverse happens. The ClientProtocol gives the ServerProtocol a reference to itself:

# Pass ServerProtocol a ref. to ClientProtocol
self.serverProtocol.clientProtocol = self;

Finally, both protocols use the stored references of each other to send data when it is received:

def dataReceived(self, data):
    self.clientProtocol.transport.write(data)
#...
def dataReceived(self, data):
    self.serverProtocol.transport.write(data)