#coding:utf-8 import networkx as nx import numpy as np import csv import random import torch import matplotlib.pyplot as plt from sklearn.model_selection import StratifiedKFold import pickle from itertools import islice lable_dict = [] for i in range(20): lable_dict.append({}) # addr0 = ["DDoS_test/"] addr0 = ['E:/DDoS/DDoS2019/0112/morefeature'] # 总地址 # addr0 = ["muti0311\\csv\\", "0112\\multi\\csv\\"] addr1 = [] num_class = 13 #addr0311=["1","2","3","4","5","6","7","LDAP","MSSQL","NetBIOS","PortMap","SYN","UDP","UDP-Lag"] #addr0311=["LDAP","MSSQL","NetBIOS","PortMap","SYN","UDP","UDP-Lag"] addr0112=["1","2","3","4","5","6","7","8","9","10","11","12","DNS","LDAP","MSSQL","NetBIOS","NTP","SNMP","SSDP","SYN","TFTP","UDP","UDP-Lag","WebDDoS"] #addr0112=["NTP","SNMP","SSDP","SYN","TFTP","UDP","UDP-Lag","WebDDoS"] #addr1.append(addr0311) addr1.append(addr0112) class S2VGraph(object): def __init__(self, g, label, node_tags=None, node_features=None): ''' g: a networkx graph label: an integer graph label node_tags: a list of integer node tags node_features: a torch float tensor, one-hot representation of the tag that is used as input to neural nets edge_mat: a torch long tensor, contain edge list, will be used to create torch sparse tensor neighbors: list of neighbors (without self-loop) ''' self.label = label self.g = g self.node_tags = node_tags self.neighbors = [] self.node_features = 0 self.edge_mat = 0 self.max_neighbor = 0 def csv_dataset(addr, ip_dst): #addr = "/home/liyuzhen/dataset/"+addr+".csv" read = csv.reader(open('%s' % addr, 'r')) len_file = len(open(addr).readlines()) print("addrr: %s: %s" % (addr, len_file)) ip_DDoS = ["172.16.0.5"] ip_DDoS.append(ip_dst) ip = {} num = 0 flows_norm = [] flows_DDoS = [] for row in islice(read, 1, None): # if(num>100000): # break num = num + 1 tmp_one_packet = [] if (row[1] == '') or (row[2] == ''): continue a = row[1] + ' ' + row[2] b = row[2] + ' ' + row[1] if (a not in ip) and (b not in ip): ip[a] = len(ip) flows_norm.append([]) try: if row[8] == '': continue else: tmp_one_packet.append(int(row[8])) # 包大小0 0 tmp_one_packet.append(row[9]) # 协议1 1 if row[5] != '': # 3 2 if int(row[5]) <= 1024: tmp_one_packet.append(int(row[5])) # print(row[7]) elif int(row[6]) <= 1024: tmp_one_packet.append(int(row[6])) # print(row[8]) else: tmp_one_packet.append(1025) # print('no port smaller than 1024:%s,%d,%s'%(addr,num,row)) else: tmp_one_packet.append(-1) if row[19] != '': tmp_one_packet.append(int(row[19])) # tcp.window_size15 3 else: tmp_one_packet.append(-1) if row[21] != '': tmp_one_packet.append(row[21]) # tcp.flags17 4 else: tmp_one_packet.append(-1) if row[23] != '': tmp_one_packet.append(row[23]) # ip.ttl19 5 else: tmp_one_packet.append(-1) except Exception as e: print(f'error in {addr}, row {num}, erro {e.args}_') continue ''' if row[3] != '': # 2 if int(row[3]) < 0: print('port smaller than 0:%s,%d,%s' % (addr, num, row)) if int(row[3]) <= 1024: tmp_one_packet.append(int(row[3])) elif int(row[4]) <= 1024: tmp_one_packet.append(int(row[4])) else: tmp_one_packet.append(1025) else: tmp_one_packet.append(-1) if row[7] != '': tmp_one_packet.append(row[7]) # frame.encap_type4 else: tmp_one_packet.append(-1) # frame.encap_type tmp_one_packet.append(float(row[10])) # 时间5 if row[10] != '': tmp_one_packet.append(float(row[10])) # http.time6 else: tmp_one_packet.append(-1) if row[11] != '': tmp_one_packet.append(int(row[11])) # icmp.len7 else: tmp_one_packet.append(-1) if row[12] != '': tmp_one_packet.append(row[12]) # icmp.type8 else: tmp_one_packet.append(-1) if row[13] != '': tmp_one_packet.append(row[13]) # irc.request9 else: tmp_one_packet.append(-1) if row[14] != '': tmp_one_packet.append(row[14]) # irc.response10 else: tmp_one_packet.append(-1) if row[15] != '': if row[15] == 0: tmp_one_packet.append(0) # tcp.ack11 else: tmp_one_packet.append(1) else: tmp_one_packet.append(-1) if row[16] != '': tmp_one_packet.append(row[16]) # tcp.ack_rtt12 else: tmp_one_packet.append(-1) if row[18] != '': tmp_one_packet.append(int(row[18])) # tcp.len14 else: tmp_one_packet.append(-1) if row[20]!='': tmp_one_packet.append(int(row[20])) # udp.length16 else: tmp_one_packet.append(-1) if row[22] != '': tmp_one_packet.append(row[22]) # ip.flags18 else: tmp_one_packet.append(-1) ''' if a in ip: tmp_one_packet[0] *= -1 flows_norm[ip[a]].append(tmp_one_packet) elif (a not in ip) and (b not in ip): print("error: in csv_dataset") elif b in ip: flows_norm[ip[b]].append(tmp_one_packet) a = ip_DDoS[0] + ' ' + ip_DDoS[1] b = ip_DDoS[1] + ' ' + ip_DDoS[0] if a in ip: flows_DDoS.append(flows_norm[ip[a]]) del flows_norm[ip[a]] elif b in ip: flows_DDoS.append(flows_norm[ip[b]]) del flows_norm[ip[b]] flows_DDoS_dict = {} len_DDoS = 0 len_norm = 0 for flows in flows_norm: len_norm += len(flows) for flows in flows_DDoS: len1 = len(flows) len_DDoS += len1 k = int(len1 / 300000) for i in range(k): flows_DDoS_dict[i] = flows[i*300000: (i+1) * 300000] if k * 300000 < len1: flows_DDoS_dict[k] = flows[(k*300000):] # print(flows_DDoS_dict) print('len_DDoS: %d' % len_DDoS) print('len_norm: %d' % len_norm) return flows_DDoS_dict, flows_norm, len_DDoS, len_norm # return flows_DDoS,flows_norm def cons_graph(all_traffic, num_node, tags, len_file, rseed): g_list = [] # print("len all_traffic: %s" %len(all_traffic)) num_graph = int(len(all_traffic) / num_node) # print("num_graph: %s" %num_graph) # print(num_graph) last_graph = len(all_traffic) - num_node*num_graph for y in range(num_graph): g = nx.Graph() node_first = 0 node_last = 0 node_tags = [] # node_features = [] a = 0 for z in range(num_node): traffic = num_node * y + z # if traffic != 0: # all_traffic[traffic][5] = float(all_traffic[traffic-1][5]) - float(all_traffic[traffic][5]) g.add_node(z) node_tags.append(all_traffic[traffic]) ''' node_tags.append([]) #print(f"len(all_traffic[traffic]): {len(all_traffic[traffic])}") for i in range(len(all_traffic[traffic])): if(all_traffic[traffic][i] not in lable_dict[i].keys() ): lable_dict[i][all_traffic[traffic][i]]=len(lable_dict[i]) #print(lable_dict) #print("label dict: %s" %lable_dict) node_tags[z].append(lable_dict[i][all_traffic[traffic][i]]) ''' # node_features.append(all_traffic[traffic][0]) if z > 0: # 构造图的边 if all_traffic[traffic][0] * all_traffic[traffic-1][0] > 0: # print(1) g.add_edge(z-1, z) # nx.draw(g,with_labels=True,pos=nx.circular_layout(g)) # plt.show() else: # print(2) g.add_edge(z, node_first) a = a+1 # plt.show() if a >= 2: g.add_edge(node_last, z-1) # nx.draw(g,with_labels=True,edge_color='b',pos=nx.circular_layout(g)) node_first=z node_last=z-1 g.add_edge(z, node_last) g_list.append(S2VGraph(g, tags, node_tags)) ''' 以下是一些平衡训练集和按比例取测试集的结果，但实际运行起来有训练集和测试集的准确率差太多的问题 ''' random.seed(rseed) random.shuffle(g_list) len_test = int(len(g_list) * 0.02) if (len_test < 30) and tags != 0: len_test = 30 if (len_test < 1) and tags == 0: len_test = 1 g_test = g_list[0: len_test] dataset_num = 20000 if (tags != 0) and (len_file >= (30 * dataset_num + len_test * 30)): len_train = len_test + int((len(all_traffic) / float(len_file)) * dataset_num) g_train = g_list[len_test: len_train] else: g_train = g_list[len_test: dataset_num + len_test] print(f"cons graph sum: {len(g_list)}") print(f"cons graph train: {len(g_train)}") print(f"cons graph test: {len(g_test)}") ''' if(tags!=0) and (len(g_list) >= (10000+len_test)): len_train=len_test+10000 g_train=g_list[len_test:len_train] else: g_train=g_list[len_test:len(g_list)] ''' del g_list return g_test, g_train #return g_list def load_data(degree_as_tag,para): ''' dataset: name of dataset test_proportion: ratio of test train split seed: random seed for random splitting of dataset ''' print('loading data') # g_list = [] all_csv_dict = {} # tag1=[0,0,0,0,0,0,0,"LDAP","MSSQL","NetBIOS","PortMap","SYN","UDP","UDP-Lag"] # tag1=[0,0,0,0,0,0,0,1,2,3,4,5,6,7] # tag1=[1,2,3,4,5,6,7] # tag2=[0,0,0,0,0,0,0,0,0,0,0,"DNS","LDAP","MSSQL","NetBIOS","NTP","SNMP","SSDP","SYN","TFTP","UDP","UDP-Lag","WebDDoS"] tag2 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] tag = [] # tag.append(tag1) tag.append(tag2) # ip_dst = ["192.168.50.4"] ip_dst = ["192.168.50.1"] # ip_dst = ["192.168.50.4","192.168.50.1"] tag_flows = {} class_flow = 13 # num_node = para['num_node'] num_node = 30 random_seeds = [2, 4, 1, 7, 9] i = 0 # root address for random_seed in random_seeds: for j in range(len(addr1[i])): # g_list = [] g_tralist = [] g_telist = [] # graph_path = "/home/liyuzhen/dataset/" + addr0[i] + '/' + 'graph_lessprotocol/'+ addr1[i][j] + '.pkl' tra_path = "E:/DDoS/DDoS2019/pkl/" + str(random_seed) + "/" + 'train/' + addr1[i][j] + '.pkl' te_path = "E:/DDoS/DDoS2019/pkl/" + str(random_seed) + "/" + 'test/' + addr1[i][j] + '.pkl' # graph = open(graph_path,"wb") train = open(tra_path, 'wb') test = open(te_path, 'wb') addr = addr0[i] + '/' + addr1[i][j] + '.csv' flows_DDoS, flows_norm, le_DDoS, le_norm = csv_dataset(addr, ip_dst[i]) # load一个文件建一次图 if tag[i][j] != 0: for k in flows_DDoS.values(): #for k in flows_DDoS: #print("flow DDoS %s :%s" % (addr1[i][j],len(k))) g_test, g_train = cons_graph(k, num_node, tag[i][j], le_DDoS, random_seed) g_tralist.extend(g_train) g_telist.extend(g_test) #g_Dlist.extend(cons_graph(k,num_node,tag[i][j],len_file)) #print("graph train DDoS %s :%s" % (addr1[i][j],len(g_tralist))) #pri`nt("graph test DDoS %s :%s" % (addr1[i][j],len(g_telist))) #print("graph DDoS %s :%s" % (addr1[i][j],len(g_Dlist))) #print("flow_DDoS:%s" %len(g_list)) else: for m in flows_norm: #print("flow norm %s :%s" % (addr1[i][j],len(m))) g_test, g_train = cons_graph(m, num_node, 0, le_norm, random_seed) g_tralist.extend(g_train) g_telist.extend(g_test) #print("graph train norm %s :%s" % (addr1[i][j],len(g_tralist))) #print("graph test norm %s :%s" % (addr1[i][j],len(g_telist))) ''' li=cons_graph(m,num_node,0,len_file) if(li != None): g_Nlist.extend(li) #print("graph normal %s :%s" % (addr1[i][j],len(g_Nlist))) #print("flow_norm:%s" %len(g_list)) ''' ''' random.seed(1) random.shuffle(g_Dlist) random.shuffle(g_Nlist) g_telist=g_Dlist[0:int(len(g_Dlist)*0.01)] g_telist.extend(g_Nlist[0:int(len(g_Nlist)*0.01)]) g_tralist=g_Nlist[int(len(g_Nlist)*0.01):] if(len(g_tralist)<=10000): g_tralist.extend(g_Dlist[int(len(g_Dlist)*0.01):]) else: g_tralist.extend(g_Dlist[int(len(g_Dlist)*0.01):int(len(g_Dlist)*0.01)+10000]) print(len(g_telist)) print(len(g_tralist)) #pickle.dump(g_list,graph) ''' print(f"{tra_path}: {len(g_tralist)}") print(f"{te_path}: {len(g_telist)}") pickle.dump(g_tralist, train) pickle.dump(g_telist, test) train.close() test.close() #print(label_dict)) def load_data2(degree_as_tag, feature, path): tagset = [] g_train = [] g_test = [] # m=[100,100,100,100,100,100,100,100,100,100,600,100,100,600] # m=[200,200,200,200,200,100,0,0,0,100,0,0,1000,1000,1000,1000,1000,1000,1000,1000,1000,1000,813,244] for i in range(len(addr0)): for j in range(len(addr1[i])): graph_path = "E:/DDoS/DDoS2019/0.01#10000pkl/" + str(path) + '/train/' + addr1[i][j] + '.pkl' graph_test = "E:/DDoS/DDoS2019/0.01#10000pkl/" + str(path) + '/test/' + addr1[i][j]+'.pkl' graph = open(graph_path, "rb") graph_test = open(graph_test, "rb") g_train.extend(pickle.load(graph)) g_test.extend(pickle.load(graph_test)) #print(type(st)) #print(st[1].node_tags) #print("每个文件的图数量：%d" %len(st)) #random.seed(1) #st=random.sample(st,m[j]) #st=random.sample(st,int(len(st)*0.1)) #g_list.extend(st) #print(len(st)) graph.close() graph_test.close() a = [] feature = feature.split(",") for i in range(len(feature)): if feature[i] != ',': a.append(int(feature[i])) print(f"feature: {feature[i]}") tagset.append({}) feature = a tagset, g_train = graph_add_features(g_train, degree_as_tag, feature, tagset) tagset, g_test = graph_add_features(g_test, degree_as_tag, feature, tagset) #g_train=one_hot(g_train,'train') #g_test=one_hot(g_test,'test') #g_train=feature_bin(g_train) #g_test=feature_bin(g_test) g_train = block_bin(g_train, tagset) g_test = block_bin(g_test, tagset) print('# maximum node tag: %d' % len(tagset)) print("# train data: %d" % len(g_train)) print('#test data: %d' % len(g_test)) return num_class, g_train, g_test def graph_add_features(g_list, degree_as_tag, feature, tagset): for g in g_list: g.neighbors = [[] for i in range(len(g.g))] for i, j in g.g.edges(): g.neighbors[i].append(j) g.neighbors[j].append(i) degree_list = [] for i in range(len(g.g)): g.neighbors[i] = g.neighbors[i] degree_list.append(len(g.neighbors[i])) g.max_neighbor = max(degree_list) #g.label = lable_dict[g.label] edges = [list(pair) for pair in g.g.edges()]#拓展边，无向边 edges.extend([[i, j] for j, i in edges]) deg_list = list(dict(g.g.degree(range(len(g.g)))).values()) g.edge_mat = torch.LongTensor(edges).transpose(0,1) if degree_as_tag: for g in g_list: g.node_tags = list(dict(g.g.degree).values()) #Extracting unique tag labels #tagset={} #tagset = set([]) for g in g_list: for j in range(len(g.node_tags)): if '0' in feature: g.node_tags[j][0]=int(g.node_tags[j][0]/100)*100 if '5' in feature: g.node_tags[j][5]=int(g.node[j][5]*100) m=[] for i in range(len(feature)): #print(f"len(g.node_tags: {len(g.node_tags[j])}") if(g.node_tags[j][feature[i]] not in tagset[i].keys()): tagset[i][g.node_tags[j][feature[i]]] = len(tagset[i]) m.append(tagset[i][g.node_tags[j][feature[i]]]) g.node_tags[j] = m #m=m+str(g.node_tags[j][int(feature[i])])+' ' #g.node_tags[j]=m #print(m) #if g.node_tags[j] not in tagset: #tagset[g.node_tags[j]]=len(tagset) #g.node_tags[j]=tagset[g.node_tags[j]] #print(tagset) #print(g.node_tags) #tagset = tagset.union(set(g.node_tags[fe])) #print(tagset) #print("union_end") #tagset = list(tagset) #tag2index = {tagset[i]:i for i in range(len(tagset))} #num_normal=0 #num_DDoS=0 return tagset, g_list def feature_bin(g_list, tagset): l = len(bin(len(tagset)).replace('0b','')) for g in g_list: g.node_features=torch.zeros(len(g.node_tags),l) for i in range(len(g.node_tags)): a=str(bin(g.node_tags[i])).replace('0b','') for j in range(len(a)): if(a[j]=='1'): g.node_features[i,l-len(a)+j]=1 #print('g.node_feature:%s,g.node_tags:%s' % (g.node_features,g.node_tags)) def block_bin(g_list, tagset): l=[] m=0 for i in range(len(tagset)): l.append(len(bin(len(tagset[i])).replace('0b',''))) m += l[i] #print(tagset) #print(m) for g in g_list: g.node_features=torch.zeros(len(g.node_tags), m) for j in range(len(g.node_tags)): for i in range(len(tagset)): a=str(bin(g.node_tags[j][i])).replace('ob','') for k in range(len(a)): if(a[k]=='1'): if(i != 0): g.node_features[j,l[i-1]+l[i]-len(a)+k]=1 else: g.node_features[j,l[i]-len(a)+k]=1 #print('g.node_feature:%s,g.node_tags:%s' % (g.node_features,g.node_tags)) return g_list def one_hot(g_list, ty, tagset): num_every={} for i in range(num_class): num_every[i]=0 for g in g_list: num_every[g.label]+=1 g.node_features = torch.zeros(len(g.node_tags), len(tagset)) g.node_features[range(len(g.node_tags)), [tagset[tag] for tag in g.node_tags]] = 1 #print("node_feature") #print("node_feature end") #num_DDoS=len(g_list)-num_normal for i in range(num_class): print("%s: class %d: %d" %(ty, i, num_every[i])) #print("num_normal:%d" %num_normal) #print("num_DDoS:%d" %num_DDoS ) #print("label_dict:%s" % lable_dict) #print('# classes: %d' %class_flow) return g_list def separate_data(graph_list, seed, fold_idx): assert 0 <= fold_idx and fold_idx < 10, "fold_idx must be from 0 to 9." skf = StratifiedKFold(n_splits=10, shuffle = True, random_state = seed) #print(f"skf: {skf}") labels = [graph.label for graph in graph_list] idx_list = [] for idx in skf.split(np.zeros(len(labels)), labels): idx_list.append(idx) #print(f"idx:{idx}") #print(idx_list) train_idx, test_idx = idx_list[fold_idx] print(f"train_num: {len(train_idx)} \n test_num: {len(test_idx)}") train_graph_list = [graph_list[i] for i in train_idx] test_graph_list = [graph_list[i] for i in test_idx] return train_graph_list, test_graph_list