initmaster

1 files changed, 163 insertions, 0 deletions

diff --git a/model/DBhelper.py b/model/DBhelper.py
new file mode 100644
index 0000000..476e5c2
--- /dev/null
+++ b/model/DBhelper.py
@@ -0,0 +1,163 @@
+import mysql.connector
+import sys
+
+
+class DBhelper(object):
+    """
+    可能的问题：
+    每次进行连接会不会造成时延呢？
+    这里没有进行关闭喔
+    一条数据一条数据的更新效率可能不高
+    """
+    def __init__(self, host='localhost', port=3306, user='root',
+                 database="test"):  # 构造函数
+        self.tablename = 'flow_feature'
+        try:
+            self.conn = mysql.connector.connect(user='root', host='localhost',password='password',
+                                                database='packetin_flow',buffered=True)
+          
+            self.cursor = self.conn.cursor()
+            print('Connected', self.conn)
+            self.cursor.execute("SHOW TABLES")
+        except Exception as e:
+            print(e)
+            print('gg!')
+
+    def create_table(self):
+        print('Init table', self.tablename)
+        self.cursor.execute("DROP TABLE IF exists %s" % self.tablename)
+        self.cursor.execute("""CREATE TABLE %s (
+          `id` INT unsigned NOT NULL AUTO_INCREMENT,
+          `ipSet` varchar(128) NOT NULL,
+          `dst_port` INT NOT NULL,
+          `src_port` INT NOT NULL,
+          `flags` INT NOT NULL,
+          `window_size` INT NOT NULL,
+          `reason` INT NOT NULL,
+          `in_port` INT NOT NULL,
+          `min_len` INT NOT NULL,
+          `max_len` INT NOT NULL,
+          `means_len` FLOAT NOT NULL,
+          `pcapnum` INT NOT NULL,
+          `tagY` INT NOT NULL,
+          PRIMARY KEY (`id`)
+          ) ENGINE=InnoDB DEFAULT CHARSET=utf8""" % self.tablename)
+        return 1
+
+    def delete_table(self):
+        print('Delete table', self.tablename)
+        self.cursor.execute("DROP TABLE IF EXISTS %s" % self.tablename)
+        return 1
+
+    def execute(self, sql):
+        self.cursor.execute(sql)
+        rowcount = self.cursor.rowcount
+        return rowcount
+
+    def get_ipset_value(self, ipSet):
+        """
+        查询函数
+        :param ipSet: 传入的是一个tuple类型的ipset
+        :return: 返回对应的dict，如果有多条记录只返回一条，如果查询失败返回0
+        """
+        ipSet_str = ' '.join(ipSet)
+        sql = """select * from %s where ipSet='%s' limit 1""" % (self.tablename, ipSet_str)
+        try:
+            # 执行SQL语句
+            self.cursor.execute(sql)
+            # 使用 fetchone() 方法获取单条数据.
+            data = self.cursor.fetchone()
+        except:
+            # 发生错误时回滚
+            print('select error, now rollback...')
+            self.conn.rollback()
+            return 0
+        this_feature = list(data)
+        this_feature.pop(0)         # 把id和ipSet删了
+        this_feature.pop(0)
+        return {ipSet: this_feature}
+
+    def judgeIPset(self, ipSet):
+        # select 1 from tablename where col = 'col' limit 1
+        # select isnull((select top(1) 1 from %s where `ipSet`=%s), 0)
+        ipSet_str = ' '.join(ipSet)
+        sql = """
+            select * from %s where ipSet = '%s' limit 1
+            """ % (self.tablename, ipSet_str)
+        out = self.cursor.execute(sql)
+        out = self.cursor.rowcount              # 不知为何用execute的out不行
+        if out:
+            print('ipSet in database.')
+        else:
+            print('Not in database, now insert...')
+        return out
+
+    def insert(self, ipSet, feature_list):
+        """
+        插入数据进mysql
+        :param ipSet:
+        :param feature_list: 传入的数据是处理好pcapnum和tagY的
+        :return: 0失败 1成功
+        """
+        ipSet_str = ' '.join(ipSet)
+        sql = """
+            INSERT INTO {0} 
+            (ipSet, dst_port, src_port, flags, window_size, 
+            reason, in_port, min_len, max_len, means_len, pcapnum, tagY)
+            VALUES
+            (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+        """.format(self.tablename)
+        values = (ipSet_str,) + tuple(feature_list)
+        try:
+            self.cursor.execute(sql, values)
+            self.conn.commit()
+            print(self.cursor.rowcount, 'Insert successful', self.cursor.lastrowid)
+            # self.cursor.close()
+            # self.conn.close()
+            return self.cursor.lastrowid
+        except:
+            # Rollback in case there is any error
+            print('Insert error, now rollback...')
+            self.conn.rollback()
+            return 0
+
+    def update(self, feature_list, ipSet):
+        """
+        数据库中本身存在，更新预测结果
+        :param feature_list: 传入一个list，已经更新好的喔
+        :param ipSet
+        :return:
+        """
+        ipSet_str = ' '.join(ipSet)
+        sql = """
+            UPDATE {0} SET
+            `ipSet` = %s, `dst_port` = %s, `src_port` = %s, `flags` = %s, `window_size` = %s,
+            `reason` = %s, `in_port` = %s, `min_len` = %s, `max_len` = %s, `means_len` = %s,
+            `pcapnum` = %s, `tagY` = %s
+            where ipSet='{1}'
+            """.format(self.tablename, ipSet_str)
+        values = (ipSet_str,) + tuple(feature_list)
+        try:
+            self.cursor.execute(sql, values)
+            self.conn.commit()
+            return 1
+        except:
+            # Rollback in case there is any error
+            print('更新失败，要哭了')
+            self.conn.rollback()
+            return 0
+
+
+if __name__ == '__main__':
+    """
+    需要先打开mysql服务，并确保有database名为packetin_flow
+    在执行packetinpredict的时候，需要先建表
+    mydb.create_table()
+    """
+    mydb = DBhelper()
+    mydb.create_table()
+    out = mydb.execute("""
+    select * from flow_feature where ipSet = '188.42.254.65 10.0.2.107' limit 1
+    """)
+    print('rowcount: ', mydb.cursor.rowcount)
+    # mydb.delete_table()