import json
import copy
import sys
import csv
import os
import os
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
# sys.path.append('....')
from lib import get_allpath as method_get_all_path
from lib import cal_metrics as method_cal_metrics
# import MABAC as method_MABAC
from lib.Rank.mabac import method_MABAC
def get_path(topo_path):
with open(topo_path, 'r') as f:
Topo = json.load(f)
f.close()
start = Topo['start']
end = Topo['goal'][0]['host']
nodes = Topo['hosts']
edges = Topo['edges']
pathsnum = Topo['nums']
newnodes = []
newedges = []
tempedges={}
for node in nodes:
newnodes.append(node['host_name'])
for edge in edges:
path = [edge['source'], edge['target']]
newedges.append(path)
return start, end, newnodes, newedges,pathsnum
def get_path1(start, end, nodes, edges):
edgesDict = {}
Numedges = {}
edgesNum = {}
for edge in edges:
if edgesDict.get(edge[0]) == None:
edgesDict[edge[0]] = [edge[1]]
else:
edgesDict[edge[0]].append(edge[1])
edgesDict[end] = []
myPath = method_get_all_path.GetPath(nodes, edgesDict, start, end)
Path = myPath.getPath() # 存放全部攻击路径
edgePath = myPath.tran(Path)
return edgePath
def all_paths_conversion(all_paths, topo_path, start):
path_id = [] # 每条边用编号代替
path_node = [] # 记录攻击路径上的点
with open(topo_path, 'r') as f:
Topo = json.load(f)
f.close()
edge_ID_dict = {}
for edge in Topo['edges']:
dict = {edge['source'] + ',' + edge['target']: edge['edge_name']}
edge_ID_dict.update(dict)
# print(edge_ID_dict)
for path in all_paths:
temp_id = []
temp_path = [start]
for source, target in path:
temp_id.append(edge_ID_dict[source + ',' + target])
temp_path.append(target)
path_id.append(temp_id)
path_node.append(temp_path)
return path_id, path_node
def get_rank(m1, m2, m3, m4, weight, visit, edgesnum, topo_path, data_path, out,flag):
# # 从右往左pop
# for i in range(edgesnum - 1, -1, -1):
# if visit[i] == 1:
# m1.pop(i)
# m2.pop(i)
# m3.pop(i)
# m4.pop(i)
# # m10.pop(i)
B_or_C = ['B', 'C', 'C', 'C', 'B', 'B', 'B', 'B', 'C', 'B'] # B表示Benefit,C表示Cost型
with open(topo_path, 'r') as f:
edges = json.load(f)['edges']
f.close()
m5_TFN = []
m6_TFN = []
m7_TFN = []
m8 = []
m9_TFN = []
for edge in edges:
if edge['edge_name'] == 'E1':
m5_TFN.append((0,0,0))
m6_TFN.append((0,0,0))
m7_TFN.append((0,0,0))
m8.append(0)
m9_TFN.append((0,0,0))
continue
# print(edge['val'])
# 机密性C(b) 模糊值
m5_TFN.append(edge['val']['C'])
# 完整性I(b) 模糊值
m6_TFN.append(edge['val']['I'])
# 可用性A(b） 模糊值
m7_TFN.append(edge['val']['A'])
# exploit(b)
m8.append(edge['val']['ES'])
# 设备影响值（b）模糊值
m9_TFN.append(edge['val']['TFN'])
# print(edge['edge_name'], edge['val']['C'], len(m5_TFN))
cias=[copy.deepcopy(m5_TFN[1:]),copy.deepcopy(m6_TFN[1:]),copy.deepcopy(m7_TFN[1:]),copy.deepcopy(m9_TFN[1:])]
m5 = method_MABAC.TFN_exact_value(m5_TFN)
m6 = method_MABAC.TFN_exact_value(m6_TFN)
m7 = method_MABAC.TFN_exact_value(m7_TFN)
m9 = method_MABAC.TFN_exact_value(m9_TFN)
for i in range(edgesnum - 1, -1, -1):
if visit[i] == 1:
m1.pop(i)
m2.pop(i)
m3.pop(i)
m4.pop(i)
m5.pop(i)
m6.pop(i)
m7.pop(i)
m8.pop(i)
m9.pop(i)
metrics_10 = [m1] + [m2] + [m3] + [m4] + [m5] + [m6] + [m7] + [m8] + [m9]
all_metrics_final = []
g = []
A_cias=[];n_cias=[]
for i in range(len(metrics_10)):
temp = metrics_10[i]
if B_or_C[i] == 'B':
# benefit
n_m = method_MABAC.benefit_normalize(copy.deepcopy(temp))
else:
# cost
n_m = method_MABAC.cost_normalize(copy.deepcopy(temp))
A_m = method_MABAC.B_to_A_weight(copy.deepcopy(n_m), weight[i])
if i in [4,5,6,8]:
n_cias.append(n_m)
A_cias.append(A_m)
g_i = method_MABAC.get_g(copy.deepcopy(A_m), visit)
final_m = method_MABAC.refresh_A_with_g(copy.deepcopy(A_m), g_i)
all_metrics_final.append(final_m)
if(flag):
out={}
out.update({'CIA_Service': [cias,n_cias,A_cias]})
with open("result/tmp/CIA_Service.json", 'w') as f:
json.dump(out, f)
f.close()
# print(all_metrics_final)
data = list(zip(*all_metrics_final))
# 保存数据
with open(data_path, 'w', newline='') as csvfile:
writer = csv.writer(csvfile)
writer.writerows(data)
temp_sum = []
temp_sum_exp = []
temp_c3 = []
for i in range(len(data)):
temp_sum.append(sum(data[i]))
temp_sum_exp.append(sum(data[i]) - data[i][2])
temp_c3.append(data[i][2])
edge_scores = [-100] * edgesnum
visit_num = 0
# print(len(data), len(visit))
for i in range(len(visit)):
# print(i-visit_num)
if visit[i] == 1:
visit_num += 1
continue
else:
edge_scores[i] = sum(data[i - visit_num])
edgeID_score = []
for i in range(len(edge_scores)):
#edgeID_score.append(['E' + str(i + 1), edge_scores[i]])
edgeID_score.append(['E' + str(i + 1), edge_scores[i], help[i]])
ranking=Single_round_sorting_results(edgeID_score,len(edgeID_score))
best_score = max(edge_scores)
rank_1_edge = edge_scores.index(best_score)
# rank1_edge这个是索引,加入索引为3，说明最重要的边是E4
return rank_1_edge,ranking
def refresh_path(best_edge, all_paths, topo_path):
# 删除all_paths中包含best_edge的边
with open(topo_path, 'r') as f:
Topo = json.load(f)
f.close()
edge_ID_dict = {}
for edge in Topo['edges']:
dict = {edge['source'] + ',' + edge['target']: edge['edge_name']}
edge_ID_dict.update(dict)
for key in edge_ID_dict.keys():
if edge_ID_dict[key] == best_edge:
# print(key)
delete_edge = (key.split(',')[0], key.split(',')[1])
# print(delete_edge)
for i in range(len(all_paths) - 1, -1, -1):
if delete_edge in all_paths[i]:
all_paths.pop(i)
#print(edge_ID_dict)
return all_paths
def get_pathnum(topo):
with open(topo_path, 'r') as f:
topo = json.load(f)
edges = {}
for edge in topo["edges"]:
key = (edge["source"], edge["target"])
edges[key] = edge
pathsnum = caculatepathsnum(edges)
return pathsnum
def Single_round_sorting_results(edgeID_score, edgesnum):
edgeID_score.sort(key=lambda x: -x[1])
# edgeID_score.sort(key=lambda x: (-x[1], x[2]))
ranking = [0] * edgesnum
result=[]
for i, res in enumerate(edgeID_score):
result.append(res)
ranking[int(res[0][1:]) - 1] = i + 1
# print('ranking=', ranking)
out.update({'ranking': result})
return ranking
if __name__ == '__main__':
topo_path = sys.argv[1]; weight_path = sys.argv[2]; data_path = sys.argv[3]; output_path = sys.argv[4]; metrics_path = sys.argv[5]
with open(weight_path, 'r') as f:
weight = json.load(f)
f.close()
weight = weight['weight']
with open('input/MainInput.json', 'r') as f:
input = json.loads(f.read())
f.close()
method=input['rrisk_cal']['method']
start, end, nodes, edges,nums= get_path(topo_path)
all_paths = get_path1(start, end, nodes, edges) # all_paths 记录了攻击路径
cutset = []; visit = [0] * nums[2]; visit[0] = 1; best_edge = 0;help = [0] * nums[2];help[0] = 1
out = {};tmp={}
rankings=[]
with open(metrics_path, 'r') as f:
metrics = list(csv.reader(f))
m1 = [float(row[1]) for row in metrics[1:]]
m2 = [float(row[2]) for row in metrics[1:]]
m3 = [float(row[3]) for row in metrics[1:]]
m4 = [float(row[4]) for row in metrics[1:]]
# m10 = [float(row[5]) for row in metrics[1:]]
f.close()
best_edge_index,ranking = get_rank(m1, m2, m3, m4, weight, visit, nums[2], topo_path, data_path, out,True)
best_edge = 'E' + str(best_edge_index + 1)
cutset.append(best_edge)
rankings.append(ranking[1:])
visit[best_edge_index] = 1
help[best_edge_index] = max(help) + 1
# 更新攻击路径
all_paths = refresh_path(best_edge, copy.deepcopy(all_paths), topo_path)
# print('best_edge', best_edge)
# print('===========================')
with open(output_path, 'w') as f:
json.dump(out, f)
f.close()
if(method=='d'):
while (len(all_paths) > 0):
path_id, path_node = all_paths_conversion(all_paths, topo_path, start)
m1 = method_cal_metrics.cal_AttackPath_Count(path_id,nums[2])
m2 = method_cal_metrics.cal_Shortest_AP_Length(path_id,nums[2])
m3 = method_cal_metrics.cal_AttackPath_Length_Average(path_id,nums[2])
m4 = method_cal_metrics.cal_shortest_distance_to_S0(path_id,nums[2])
best_edge_index,ranking = get_rank(m1, m2, m3, m4, weight, visit, nums[2], topo_path, data_path, out,False)
best_edge = 'E' + str(best_edge_index + 1)
cutset.append(best_edge)
rankings.append(ranking[1:])
visit[best_edge_index] = 1
help[best_edge_index] = max(help) + 1
# 更新攻击路径
all_paths = refresh_path(best_edge, copy.deepcopy(all_paths), topo_path)
# print('best_edge', best_edge)
# print('===========================')
tmp.update({'ranking':rankings})
tmp.update({'delete_edges':cutset})
with open("result/tmp/ranking.json", 'w') as f:
json.dump(tmp, f)
f.close()
如何在别的程序中调用这个排序算法程序，需要传入的参数以及程序输出的排序结果在哪里
超级会员免费看
订阅专栏 解锁全文
扫一扫
29
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
