diff options
Diffstat (limited to 'model/DBhelper.py')
| -rw-r--r-- | model/DBhelper.py | 163 |
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() |