twgs-proxy/tcp-proxy.py

#!/usr/bin/env python3

import sys

from twisted.internet import defer
from twisted.internet import protocol
from twisted.internet import reactor
from twisted.python import log

HOST = '127.0.0.1'
PORT = 2002
LISTEN_PORT = 9999
LISTEN_ON = '127.0.0.1'


class ProxyClientProtocol(protocol.Protocol):
    def connectionMade(self):
        log.msg("Client: connected to peer")
        self.cli_queue = self.factory.cli_queue
        self.cli_queue.get().addCallback(self.serverDataReceived)

    def serverDataReceived(self, chunk):
        # TODO: Line processing, and line cleaning (remove ANSI color codes)
        if chunk is False:
            self.cli_queue = None
            log.msg("Client: disconnecting from peer")
            self.factory.continueTrying = False
            self.transport.loseConnection()
        elif b"$" == chunk:
            self.factory.svr_queue.put(b"HELLO.\r\n")
            self.cli_queue.get().addCallback(self.serverDataReceived)           
        elif self.cli_queue:
            log.msg("Client: writing %d bytes to peer" % len(chunk))
            log.msg(">>", chunk)
            self.transport.write(chunk)
            self.cli_queue.get().addCallback(self.serverDataReceived)
        else:
            self.factory.cli_queue.put(chunk)

    def dataReceived(self, chunk):
        log.msg("Client: %d bytes received from peer" % len(chunk))
        log.msg("<<", chunk)
        self.factory.svr_queue.put(chunk)

    def connectionLost(self, why):
        log.msg("connectionLost because: %s" % why)
        # if self.cli_queue:
        #     self.cli_queue = None
        #     log.msg("Client: peer disconnect unexpectedly")
        self.factory.svr_queue.put(False)
        self.cli_queue = None
        self.transport.loseConnection()

# class ProxyClientFactory(protocol.ReconnectingClientFactory):
class ProxyClientFactory(protocol.ClientFactory):    
    maxDelay = 10
    # continueTrying = True
    # continueTrying = False
    protocol = ProxyClientProtocol

    def __init__(self, svr_queue, cli_queue):
        self.svr_queue = svr_queue
        self.cli_queue = cli_queue

class ProxyServer(protocol.Protocol):
    def connectionMade(self):
        self.svr_queue = defer.DeferredQueue()
        self.cli_queue = defer.DeferredQueue()
        self.svr_queue.get().addCallback(self.clientDataReceived)

        factory = ProxyClientFactory(self.svr_queue, self.cli_queue)
        reactor.connectTCP(HOST, PORT, factory)

    def clientDataReceived(self, chunk):
        if chunk is False:
            self.transport.loseConnection()
        else:
            log.msg("Server: writing %d bytes to original client" % len(chunk))
            self.transport.write(chunk)
            self.svr_queue.get().addCallback(self.clientDataReceived)

    def dataReceived(self, chunk):
        log.msg("Server: %d bytes received" % len(chunk))
        self.cli_queue.put(chunk)

    def connectionLost(self, why):
        log.msg("lost connection %s" % why)
        self.cli_queue.put(False)



if __name__ == "__main__":
    log.startLogging(sys.stdout)
    factory = protocol.Factory()
    factory.protocol = ProxyServer
    reactor.listenTCP(LISTEN_PORT, factory, interface=LISTEN_ON)
    reactor.run()