使用Python的twisted和socket模块实现端口的负载分发

很简单的需求，自己写个类似iptables那样的dnat端口转发，简单实现像lvs那样的nat模式的端口的负载分发，当然性能堪忧哈~

这个例子是 监听 本地ip的9999端口，然后映射到另一个端口上，也可以利用random参数，进行多个端口的轮训，当然他的算法和性能简单，不能和lvs 相比了。

映射的版本：

执行方法：python th.py 8888 10.10.10.63 80

#!/usr/bin/env python
#python th.py 8888 10.10.10.63 80
import sys, socket, time, threading
loglock = threading.Lock()
def log(msg):
loglock.acquire()
try:
print '[%s]: \n%s\n' % (time.ctime(), msg.strip())
sys.stdout.flush()
finally:
loglock.release()
class PipeThread(threading.Thread):
def __init__(self, source, target):
threading.Thread.__init__(self)
self.source = source
self.target = target
def run(self):
while True:
try:
data = self.source.recv(1024)
log(data)
if not data: break
self.target.send(data)
except:
break
log('PipeThread done')
class Forwarding(threading.Thread):
def __init__(self, port, targethost, targetport):
threading.Thread.__init__(self)
self.targethost = targethost
self.targetport = targetport
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.sock.bind(('0.0.0.0', port))
self.sock.listen(10)
def run(self):
while True:
client_fd, client_addr = self.sock.accept()
target_fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
target_fd.connect((self.targethost, self.targetport))
log('new connect')
# two direct pipe
PipeThread(target_fd, client_fd).start()
PipeThread(client_fd, target_fd).start()
if __name__ == '__main__':
print 'Starting'
import sys
try:
port = int(sys.argv[1])
targethost = sys.argv[2]
try: targetport = int(sys.argv[3])
except IndexError: targetport = port
except (ValueError, IndexError):
print 'Usage: %s port targethost [targetport]' % sys.argv[0]
sys.exit(1)
#sys.stdout = open('forwaring.log', 'w')
Forwarding(port, targethost, targetport).start()

附带负载的版本：

#!/usr/bin/env python
import sys, socket, time, threading
loglock = threading.Lock()
def log(msg):
loglock.acquire()
try:
print '[%s]: \n%s\n' % (time.ctime(), msg.strip())
sys.stdout.flush()
finally:
loglock.release()
class PipeThread(threading.Thread):
def __init__(self, source, target):
threading.Thread.__init__(self)
self.source = source
self.target = target
def run(self):
while True:
try:
data = self.source.recv(1024)
log(data)
if not data: break
self.target.send(data)
except:
break
log('PipeThread done')
class Forwarding(threading.Thread):
def __init__(self, port, targethost, targetport):
number = random.randint(1,2)
numbers = '%d' %number
portlist={'1':80,'2':81}
host='10.10.10.63'
threading.Thread.__init__(self)
self.targethost = targethost
self.targetport = portlist[numbers]
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.sock.bind(('0.0.0.0', port))
self.sock.listen(10)
def run(self):
while True:
client_fd, client_addr = self.sock.accept()
target_fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
target_fd.connect((self.targethost, portlist[numbers]))
log('new connect')
# two direct pipe
PipeThread(target_fd, client_fd).start()
PipeThread(client_fd, target_fd).start()
if __name__ == '__main__':
print 'Starting'
import sys
import random
number = random.randint(1,2)
numbers = '%d' %number
#sys.stdout = open('forwaring.log', 'w')
portlist={'1':80,'2':81}
host='10.10.10.63'
Forwarding(9999,host, portlist[numbers]).start()

对于性能要求高的，可以试试twisted端口转发的代码。

调优方向  I/O模型，比如阻塞、非阻塞和反应式（select,poll,WaitForMultipleObject)

现在用gevent比较多，一是gevent性能会更好一些，而且用同步的方式来实现异步(使用greenlet)，twisted中defer 和
callback 会把逻辑打散。
当初从twisted转到gevent最主要的原因是在twisted上很难实现多进程。
gevent不好的地方就是它的mongkey_patch 有时候会带来一些奇怪的问题。

个人感觉，twisted的优点是程序的各种设计，要是针对于咱们这个映射需求的话，推荐使用gevent网络并发框架。  这个也是我较常用的。

from twisted.internet.protocol import Protocol,ClientCreator
from twisted.internet import reactor
from twisted.protocols.basic import LineReceiver
from twisted.internet.protocol import Factory,ClientFactory
class Transfer(Protocol):
def __init__(self):
pass
def connectionMade(self):
c = ClientCreator(reactor,Clienttransfer)
c.connectTCP("10.10.10.63",80).addCallback(self.set_protocol)
self.transport.pauseProducing()
def set_protocol(self,p):
self.server = p
p.set_protocol(self)
def dataReceived(self,data):
self.server.transport.write(data)
def connectionLost(self,reason):
self.transport.loseConnection()
self.server.transport.loseConnection()
class Clienttransfer(Protocol):
def __init__(self):
pass
def set_protocol(self,p):
self.server = p
self.server.transport.resumeProducing()
pass
def dataReceived(self,data):
self.server.transport.write(data)
pass
factory = Factory()
factory.protocol = Transfer
reactor.listenTCP(9999,factory)
reactor.run()

完成 ！