接到一個任務，將攻擊轉移矩陣進行視覺化，生成攻擊轉移概率圖，便嘗試用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