Frist Version

Server/shadowsocks/db_transfer.py

@@ -0,0 +1,132 @@
+#!/usr/bin/python
+# -*- coding: UTF-8 -*-
+
+import logging
+import cymysql
+import time
+import sys
+from server_pool import ServerPool
+import Config
+
+class DbTransfer(object):
+
+    instance = None
+
+    def __init__(self):
+        self.last_get_transfer = {}
+
+    @staticmethod
+    def get_instance():
+        if DbTransfer.instance is None:
+            DbTransfer.instance = DbTransfer()
+        return DbTransfer.instance
+
+    def push_db_all_user(self):
+        #更新用户流量到数据库
+        last_transfer = self.last_get_transfer
+        curr_transfer = ServerPool.get_instance().get_servers_transfer()
+        #上次和本次的增量
+        dt_transfer = {}
+        for id in curr_transfer.keys():
+            if id in last_transfer:
+                if last_transfer[id][0] == curr_transfer[id][0] and last_transfer[id][1] == curr_transfer[id][1]:
+                    continue
+                elif curr_transfer[id][0] == 0 and curr_transfer[id][1] == 0:
+                    continue
+                elif last_transfer[id][0] <= curr_transfer[id][0] and \
+                last_transfer[id][1] <= curr_transfer[id][1]:
+                    dt_transfer[id] = [curr_transfer[id][0] - last_transfer[id][0],
+                                       curr_transfer[id][1] - last_transfer[id][1]]
+                else:
+                    dt_transfer[id] = [curr_transfer[id][0], curr_transfer[id][1]]
+            else:
+                if curr_transfer[id][0] == 0 and curr_transfer[id][1] == 0:
+                    continue
+                dt_transfer[id] = [curr_transfer[id][0], curr_transfer[id][1]]
+
+        self.last_get_transfer = curr_transfer
+        query_head = 'UPDATE user'
+        query_sub_when = ''
+        query_sub_when2 = ''
+        query_sub_in = None
+        last_time = time.time()
+        for id in dt_transfer.keys():
+            query_sub_when += ' WHEN %s THEN u+%s' % (id, dt_transfer[id][0])
+            query_sub_when2 += ' WHEN %s THEN d+%s' % (id, dt_transfer[id][1])
+            if query_sub_in is not None:
+                query_sub_in += ',%s' % id
+            else:
+                query_sub_in = '%s' % id
+        if query_sub_when == '':
+            return
+        query_sql = query_head + ' SET u = CASE port' + query_sub_when + \
+                    ' END, d = CASE port' + query_sub_when2 + \
+                    ' END, t = ' + str(int(last_time)) + \
+                    ' WHERE port IN (%s)' % query_sub_in
+        #print query_sql
+        conn = cymysql.connect(host=Config.MYSQL_HOST, port=Config.MYSQL_PORT, user=Config.MYSQL_USER,
+                               passwd=Config.MYSQL_PASS, db=Config.MYSQL_DB, charset='utf8')
+        cur = conn.cursor()
+        cur.execute(query_sql)
+        cur.close()
+        conn.commit()
+        conn.close()
+
+    @staticmethod
+    def pull_db_all_user():
+        #数据库所有用户信息
+        conn = cymysql.connect(host=Config.MYSQL_HOST, port=Config.MYSQL_PORT, user=Config.MYSQL_USER,
+                               passwd=Config.MYSQL_PASS, db=Config.MYSQL_DB, charset='utf8')
+        cur = conn.cursor()
+        cur.execute("SELECT port, u, d, transfer_enable, passwd, switch, enable FROM user")
+        rows = []
+        for r in cur.fetchall():
+            rows.append(list(r))
+        cur.close()
+        conn.close()
+        return rows
+
+    @staticmethod
+    def del_server_out_of_bound_safe(rows):
+    #停止超流量的服务
+    #启动没超流量的服务
+    #修改下面的逻辑要小心包含跨线程访问
+        for row in rows:
+            if ServerPool.get_instance().server_is_run(row[0]) is True:
+                if row[5] == 0 or row[6] == 0:
+                    #stop disable or switch off user
+                    logging.info('db stop server at port [%s] reason: disable' % (row[0]))
+                    ServerPool.get_instance().del_server(row[0])
+                elif row[1] + row[2] >= row[3]:
+                    #stop out bandwidth user
+                    logging.info('db stop server at port [%s] reason: out bandwidth' % (row[0]))
+                    ServerPool.get_instance().del_server(row[0])
+                if ServerPool.get_instance().tcp_servers_pool[row[0]]._config['password'] != row[4]:
+                    #password changed
+                    logging.info('db stop server at port [%s] reason: password changed' % (row[0]))
+                    ServerPool.get_instance().del_server(row[0]) 
+            else:
+                if row[5] == 1 and row[6] == 1 and row[1] + row[2] < row[3]:
+                    logging.info('db start server at port [%s] pass [%s]' % (row[0], row[4]))
+                    ServerPool.get_instance().new_server(row[0], row[4])
+
+    @staticmethod
+    def thread_db():
+        import socket
+        import time
+        timeout = 60
+        socket.setdefaulttimeout(timeout)
+        while True:
+            #logging.warn('db loop')
+            try:
+                DbTransfer.get_instance().push_db_all_user()
+                rows = DbTransfer.get_instance().pull_db_all_user()
+                DbTransfer.del_server_out_of_bound_safe(rows)
+            except Exception as e:
+                logging.warn('db thread except:%s' % e)
+            finally:
+                time.sleep(15)
+
+
+#SQLData.pull_db_all_user()
+#print DbTransfer.get_instance().test()