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

接到一个任务，将攻击转移矩阵进行可视化，生成攻击转移概率图，便尝试用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

https://www.osgeo.cn/networkx/auto_examples/drawing/plot_weighted_graph.html#sphx-glr-auto-examples-drawing-plot-weighted-graph-py