#!/usr/bin/env python
# coding: utf-8
# https://gist.github.com/fiorix/1878983
# http://musta.sh/2012-03-04/twisted-tcp-proxy.html

import sys

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

if len(sys.argv) != 5:
    sys.exit('usage:\n  proxy <source addr> <source port> <target addr> <target port>')

class ProxyClientProtocol(protocol.Protocol):
    def connectionMade(self):
        self.cli_queue = self.factory.cli_queue
        self.cli_queue.get().addCallback(self.serverDataReceived)

    def serverDataReceived(self, chunk):
        if chunk is False:
            self.cli_queue = None
            self.factory.continueTrying = False
            self.transport.loseConnection()
        elif self.cli_queue:
            self.transport.write(chunk)
            self.cli_queue.get().addCallback(self.serverDataReceived)
        else:
            self.factory.cli_queue.put(chunk)

    def dataReceived(self, chunk):
        self.factory.srv_queue.put(chunk)

    def connectionLost(self, why):
        if self.cli_queue:
            self.cli_queue = None


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

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

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

        factory = ProxyClientFactory(self.srv_queue, self.cli_queue)
        reactor.connectTCP(sys.argv[1], int(sys.argv[2]), factory)

    def clientDataReceived(self, chunk):
        self.transport.write(chunk)
        self.srv_queue.get().addCallback(self.clientDataReceived)

    def dataReceived(self, chunk):
        self.cli_queue.put(chunk)

    def connectionLost(self, why):
        self.cli_queue.put(False)

if __name__ == "__main__":
    factory = protocol.Factory()
    factory.protocol = ProxyServer
    reactor.listenTCP(int(sys.argv[4]), factory, interface=sys.argv[3])
    reactor.run()
