Twisted 编写的TCP转向连接程序

阅读更多

机房对外的系统安装的是Linux，内网有一台Windows机器，想找工作把3389端口转到Linux服务器上，每次临时启用。

安装了iptables，竟然提示缺少模块，配置也比较繁琐，所以就放弃了。

想到以前用过的twisted，于是自己写了一个小工具：

from twisted.internet.protocol import Protocol, ClientFactory, ServerFactory
from twisted.internet import reactor
import sys

class ForwardServer(Protocol):
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.data = ""
        self._connected = False

    def dataReceived(self, data):
        #print "Received %d bytes from client\n" % len(data)
        self.data += data
        #print "%d bytes in buffer" % len(self.data)
        if self._connected and (len(self.data) > 0):
            self.connector.transport.write(self.data)
            #print "Sent %d bytes to server" % len(self.data)
            self.data = ""

    def connectionMade(self):
        self.connector = reactor.connectTCP(self.host, self.port, ForwardClientFactory(self))
        print "Client connected"

    def setConnected(self, flag):
        if flag:
            self.onConnected()
        else:
            self.transport.loseConnection()
        self._connected = flag

    def onConnected(self):
        if len(self.data) > 0:
            self.connector.transport.write(self.data)
            self.data = ""

    def connectionLost(self, reason):
        self.connector.transport.loseConnection()
        self._connected = False
        print "Client disonnected"

class ForwardClient(Protocol):
    def __init__(self, forward):
        self.forward = forward

    def dataReceived(self, data):
        #print "Received %d bytes from server\n" % len(data)
        self.forward.transport.write(data)

    def connectionMade(self):
        print "Connected to server"
        self.forward.setConnected(True)

    def connectionLost(self, reason):
        self.forward.setConnected(False)
        print "Disconnected from server"

class ForwardServerFactory(ServerFactory):
    def __init__(self, host, port):
        self.host = host
        self.port = port

    def buildProtocol(self, addr):
        return ForwardServer(self.host, self.port)

class ForwardClientFactory(ClientFactory):
    def __init__(self, forward):
        self.forward = forward

    def buildProtocol(self, addr):
        return ForwardClient(self.forward)

    def clientConnectionFailed(self, connector, reason):
        self.forward.transport.loseConnection()


if __name__ == "__main__":
    if len(sys.argv) != 4:
        print "USAGE: %s   " % sys.argv[0]
        sys.exit(1)
    host, port, listen_port = sys.argv[1], int(sys.argv[2]), int(sys.argv[3])
    server_factory = ForwardServerFactory(host, port)
    reactor.listenTCP(listen_port, server_factory)
    reactor.run()

运行：
python port_forward.py 192.168.0.10 3389 3389

就可以在本机上监听3389端口，有客户连接时就建立起到服务器的连接，并在2台机器间转发数据。

twisted是reactor架构，可以支持多个侦听器和连接器，也可以接受多个客户同时访问。作些修改，就可以改成简单的负载均衡服务器。twisted易用性、效率都是很不错的。