接到一个任务，将攻击转移矩阵进行可视化，生成攻击转移概率图，便尝试用python实现一下。
查阅资料，看大家都在用networkx和matplotlib进行可视化，便边学边做，记录一下学习笔记。
任务：将手里了多个攻击过程，如图1所示，生成为攻击转移模型，如图2所示
图1：
图2：
由于直接画图节点过多，于是对节点进行了映射。
过程
首先读取文件并以列表的形式表示：
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)#具体参数可以上网查询
-----------------------------------------------------------------------------------------------------------------------------------------
参考资料
https://zhuanlan.zhihu.com/p/36700425