1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
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()
|
