ssh tunnel monitor

gistfile1.py

#!/usr/bin/pypy
import subprocess
import shlex
import select
import socket
import time

class Proxy(object):
        def __init__(self,bind,port,remote,local):
                self.cmd = "/usr/bin/ssh -NC -D %s:%s %s" % (bind,port,remote)
                self.port = port
                self.local = local
                self.proc = None
                self.conn = None
                self.alive()

        def alive(self):
                if not self.connect():
                        self.proc = subprocess.Popen(shlex.split(self.cmd))
                        while self.proc.returncode is None:
                                time.sleep(1)
                                if self.connect():
                                        return
                        # process die
                        self.alive()

        def connect(self):
                conn = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
                conn.settimeout(None)
                try:
                        conn.connect((self.local,self.port))
                        self.conn = conn
                        Proxy.register(self)
                        print "%s %s ok" % (self.local,self.port)
                        return True
                except Exception,e:
                        return False

        def unregister(self):
                Proxy.poller.modify(self.conn.fileno(),0)
                del Proxy.fds[self.conn.fileno()]
                self.conn.close()

        @staticmethod
        def kick():
                Proxy.poller = select.epoll()
                Proxy.fds = {}

        @staticmethod
        def register(instance):
                Proxy.poller.register(instance.conn,select.EPOLLOUT)
                Proxy.fds[instance.conn.fileno()] = instance

        @staticmethod
        def monitor():
                while True:
                        for fd,event in Proxy.poller.poll():
                                instance = Proxy.fds[fd]
                                if event & select.EPOLLOUT:
                                        instance.conn.setblocking(0)
                                        Proxy.poller.modify(fd,select.EPOLLIN)
                                        continue

                                while True:
                                        if len(instance.conn.recv(1024)) == 0:
                                                instance.unregister()
                                                instance.alive()
                                                break
if __name__ == '__main__':
        Proxy.kick()
        Proxy('0.0.0.0',10087,'db1.dc','127.0.0.1')
        Proxy('0.0.0.0',10086,'db1.dc','127.0.0.1')
        Proxy.monitor()