networkx学习与攻击转移图可视化

f1 = open(path,'r')
list_lines = []
while True:
    line = f1.readline()
    if line:
        list_line = line[0:-1].split(',')#去掉回车符'\n'
        list_lines.append(list_line)
    else:
        break

#读取对应文件
fk = open('./killlinesort.txt','r')
dict_kill = {}
while True:
    linek = fk.readline()
    if linek:
        spilt_line = linek.split(' level')
        dict_kill[spilt_line[0]] = spilt_line[1].split(',')[1][0:-1]
    else:
        break

attack = []
for key in dict_kill:
    if dict_kill[key] not in attack:
        attack.append(dict_kill[key])
Matrix_tmp = pd.DataFrame(0, columns=attack, index=attack)

#计算转移次数
for i in list_lines:  #读取攻击过程
    for index in range(len(i)-1):
        fT = dict_kill[i[index]]
        nT = dict_kill[i[index+1]]
        if fT == nT:
            continue
        else:
            Matrix_tmp[nT][fT] = Matrix_tmp[nT][fT] + 1
Matrix_P = Matrix_tmp.div((Matrix_tmp.apply(sum,axis=1)),axis=0).fillna(0)#计算转移概率

#去掉孤立的节点
attack1 = []
for i in attack:
    row = Matrix_P.loc[:,i].values[0:]
    columns1 = Matrix_P.loc[i,:].values[0:]
    if sum(row) == 0 and sum(columns1) == 0:
        continue
    else:
        attack1.append(i)#其中是映射完并去掉孤立节点以后的攻击节点

#将转移概率矩阵转换成三元组用字典存储
dict_tup = {}
final_attack = []#最后一步攻击
for fattack in attack1:  # lines是行名也是源节点
    Max_P_Array = Matrix_P.loc[fattack].values[0:]
    if sum(Max_P_Array) == 0:
        final_attack.append(fattack)#记录下每个攻击过程的最后一步攻击，作为广度优先遍历的起始点
    col_list = Matrix_P.columns.values.tolist()
    count = 0
    for values in Max_P_Array:
        if values <= 0:#筛选边
            count += 1
            continue
        else:
            columns = col_list[count]#目的节点
            count += 1
            key = fattack + ',' + columns#结构为{‘源攻击节点，目的攻击节点’：转移概率}
            dict_tup[key] = values

#遍历有所节点
def findallnode(node_list,Matrix_P,pos_num):
    before_nodes = []#上一步的攻击节点
    row_num = 1
    pos_count = 1
    for n in node_list:
        tran_p = Matrix_P.loc[:,n].values[0:]
        count = 0
        for i in tran_p:
            if i <= 0:
                count += 1
                continue
            else:
                before = attack[count]
                count += 1
                if not before in G.nodes:
                    before_nodes.append(before)
                    G.add_node(before,pos = (pos_count*5,(pos_num*5+((-1)**row_num)*3)))
                    #pos的位置可以认为是一个坐标，同层节点之间间隔5，由于有些同层节点之间也有关系，担心重合遮挡，所以y轴坐标进行错位
                    row_num += 1
                    pos_count += 1
    pos_num += 1
    #进行迭代，直到所有的节点都被遍历，G.node的类型是list
    if (len(list(set(attack1)-set(G.nodes))) != 0):
        findallnode(before_nodes, Matrix_P, pos_num)

num = 1
#确定初始参数
for i in final_attack:
    G.add_node(i,pos = (num*10,0))
    num+=1
findallnode(final_attack,Matrix_P,pos_num=1)

#添加有向边
for key in dict_tup:
    value = round(dict_tup[key],3)
    #print(value)
    list_key = key.split(',')
    list_key.append(value)
    tup = tuple(list_key)#得到三元组
    G.add_edge(str(tup[0]),str(tup[1]),p=tup[2])

pos=nx.get_node_attributes(G,'pos')#这个函数能一字典的形式存储每个节点的坐标

#pos结果为
{'DOWNLOAD FILE': (10, 0), 'keep connect': (20, 0), 'DIRECTORY discovery': (5, 2), 'FILE OPERATION': (10, 8), 'Password attacks': (15, 2), 
'Vulnerability exploit': (20, 8), 'Initial Access': (25, 2), 'UPLOAD FILE': (30, 8), 'WEBSHELL': (35, 2), 'DDOS': (5, 7), 'Command and Control': (10, 13), 
'INFORMATION DISCLOSURE': (15, 7), 'ERROR OCCUR': (5, 12), '.': (0, 0)}

nx.draw_networkx_nodes(G, pos, node_size=50, node_color='gray')#将节点画出，各个参数可以查询关方文档，选择需要的
nx.draw_networkx_edges(G, pos, width=1.5,arrowstyle='->',arrowsize=10,alpha=0.5,)#画出边
nx.draw_networkx_edge_labels(G,pos,label_pos=0.35,font_size=5,font_weight='bold')#label_pos是转移概率在边上的位置，font_weight中bold为加粗，默认是正常

#由于pos中的坐标是以元组的形式展示，而元组不能修改，所以先转成list再进行替换
for key in pos:
    list_loc = list(pos[key])
    list_loc[1] += .7#上移的值，根据实际情况修改
    pos[key] = tuple(list_loc)
nx.draw_networkx_labels(G,pos,font_size=8,alpha=0.8)#画出攻击类型名称

#plt.show()#show()和savefig不能同时存在，不然生成的文件是白板

plt.savefig('./****.png',pad_inches=20,dpi=300)#具体参数可以上网查询