247 lines
12 KiB
Python
247 lines
12 KiB
Python
|
import angr
|
|||
|
|
import csv
|
|||
|
|
from angrutils import plot_cfg, hook0, set_plot_style
|
|||
|
|
import bingraphvis
|
|||
|
|
import networkx as nx
|
|||
|
|
import os
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
import sys
|
|||
|
|
# sys.path.append(r'../ASM2VEC_plus_scripts/')
|
|||
|
|
# from func2vec import func2vec,load_model
|
|||
|
|
# from node_feature import asm2vec_plus
|
|||
|
|
import lief
|
|||
|
|
from detect_pe_packer import detect_pack_res
|
|||
|
|
# exit()
|
|||
|
|
from node_feature import *
|
|||
|
|
|
|||
|
|
def my_round(vec_list):
|
|||
|
|
for i in range(len(vec_list)):
|
|||
|
|
vec_list[i] = round(vec_list[i], 5)
|
|||
|
|
return vec_list
|
|||
|
|
|
|||
|
|
def get_node_feature(hex_asm="558bec83ec085756bf0bb80000ff15d4804000",node_feature_method="n_gram"):
|
|||
|
|
if node_feature_method =="asm2vec_plus_16":
|
|||
|
|
return asm2vec_plus_16(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_plus_32":
|
|||
|
|
return asm2vec_plus_32(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_plus_64":
|
|||
|
|
return asm2vec_plus_64(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_plus_128":
|
|||
|
|
return asm2vec_plus_128(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_plus_256":
|
|||
|
|
return asm2vec_plus_256(hex_asm=hex_asm).tolist()
|
|||
|
|
|
|||
|
|
elif node_feature_method =="asm2vec_base_16":
|
|||
|
|
return asm2vec_base_16(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_base_32":
|
|||
|
|
return asm2vec_base_32(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_base_64":
|
|||
|
|
return asm2vec_base_64(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_base_128":
|
|||
|
|
return asm2vec_base_128(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_base_256":
|
|||
|
|
return asm2vec_base_256(hex_asm=hex_asm).tolist()
|
|||
|
|
|
|||
|
|
elif node_feature_method == "asm2vec_s_base_16":
|
|||
|
|
return asm2vec_s_base_16(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method == "asm2vec_s_base_32":
|
|||
|
|
return asm2vec_s_base_32(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method == "asm2vec_s_base_64":
|
|||
|
|
return asm2vec_s_base_64(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method == "asm2vec_s_base_128":
|
|||
|
|
return asm2vec_s_base_128(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method == "asm2vec_s_base_256":
|
|||
|
|
return asm2vec_s_base_256(hex_asm=hex_asm).tolist()
|
|||
|
|
|
|||
|
|
elif node_feature_method == "asm2vec_s368_base_16":
|
|||
|
|
return asm2vec_s368_base_16(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method == "asm2vec_s368_base_32":
|
|||
|
|
return asm2vec_s368_base_32(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method == "asm2vec_s368_base_64":
|
|||
|
|
return asm2vec_s368_base_64(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method == "asm2vec_s368_base_128":
|
|||
|
|
return asm2vec_s368_base_128(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method == "asm2vec_s368_base_256":
|
|||
|
|
return asm2vec_s368_base_256(hex_asm=hex_asm).tolist()
|
|||
|
|
|
|||
|
|
elif node_feature_method =="malconv":
|
|||
|
|
return malconv(hex_asm=hex_asm)
|
|||
|
|
elif node_feature_method =="n_gram":
|
|||
|
|
return n_gram(hex_asm=hex_asm)
|
|||
|
|
elif node_feature_method =="word_frequency":
|
|||
|
|
return word_frequency(hex_asm=hex_asm)
|
|||
|
|
elif node_feature_method =="asm2vec_plus_small":
|
|||
|
|
return asm2vec_plus_small(hex_asm=hex_asm).tolist()
|
|||
|
|
elif node_feature_method =="asm2vec_base_small":
|
|||
|
|
return asm2vec_base_small(hex_asm=hex_asm).tolist()
|
|||
|
|
|
|||
|
|
asm2vec_model_path="../ASM2VEC_plus_scripts/asm2vec_checkpoints/model_100.pt"
|
|||
|
|
|
|||
|
|
# print(bengin_list)
|
|||
|
|
|
|||
|
|
# jmp_family=["jmp","call","ret","retf",
|
|||
|
|
# "ja","jnbe","jae","jnb","jb","jane","jbe","jna",
|
|||
|
|
# "jg","jnle","jge","jnl","jl","jnge","jle","jng",
|
|||
|
|
# "je","jz","jne","jnz","jc","jnc","jno","jnp","jpo",
|
|||
|
|
# "jns","jo","jp","jpe","js"
|
|||
|
|
# "loop","loope","loopz","loopne","loopnz","jcxz","jecxz"
|
|||
|
|
# ]
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_new_section_addr(bin_parse):
|
|||
|
|
entry_point = bin_parse.optional_header.addressof_entrypoint
|
|||
|
|
# 找到入口点所在的section名字
|
|||
|
|
entryname = bin_parse.section_from_rva(entry_point).name
|
|||
|
|
for section in bin_parse.sections:
|
|||
|
|
if section.name == entryname:
|
|||
|
|
text_characteristics = section.characteristics
|
|||
|
|
virtual_address = section.virtual_address
|
|||
|
|
virtual_size = section.size
|
|||
|
|
virtual_offset = section.offset
|
|||
|
|
# 创建新段，并设置偏移位置
|
|||
|
|
new_section = lief.PE.Section("test")
|
|||
|
|
if virtual_size % 0x1000 == 0:
|
|||
|
|
mod_num = int(virtual_size / 0x1000)
|
|||
|
|
else:
|
|||
|
|
mod_num = int(virtual_size / 0x1000) + 1
|
|||
|
|
new_section.virtual_address = virtual_address + mod_num * 0x1000
|
|||
|
|
# print(hex(new_section.virtual_address))
|
|||
|
|
# exit()
|
|||
|
|
# print(hex(new_section.virtual_address + bin_parse.imagebase))
|
|||
|
|
|
|||
|
|
new_section.offset = virtual_offset + virtual_size
|
|||
|
|
return new_section.offset
|
|||
|
|
|
|||
|
|
|
|||
|
|
def cfg_extract(file_list,data_dir="../data/malware",CFG_dir="../CFG_data/malware",csv_save_path="../CFG_data/malware_msg.csv",header = ['malware_name','nodes_num','edgs_num']):
|
|||
|
|
|
|||
|
|
csv_data=[]
|
|||
|
|
#载入asm2vec的模型
|
|||
|
|
# asm2vec_model=load_model(asm2vec_model_path)
|
|||
|
|
for i in tqdm(range(len(file_list))):
|
|||
|
|
# for i in tqdm(range(0,50)):
|
|||
|
|
file_item=file_list[i]
|
|||
|
|
|
|||
|
|
#剔除加壳程序
|
|||
|
|
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
if detect_pack_res(os.path.join(data_dir, file_item)) == True:
|
|||
|
|
continue
|
|||
|
|
bin_parse = lief.PE.parse(os.path.join(data_dir,file_item))
|
|||
|
|
p = angr.Project(os.path.join(data_dir,file_item), load_options={'auto_load_libs': False})
|
|||
|
|
cfg = p.analyses.CFGFast(show_progressbar=True, normalize=False, resolve_indirect_jumps=False,
|
|||
|
|
force_smart_scan=False, symbols=False, data_references=False)
|
|||
|
|
# new_section_addr = get_new_section_addr(bin_parse)
|
|||
|
|
except:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
|
|||
|
|
G = cfg.graph
|
|||
|
|
|
|||
|
|
#如果反汇编有出错，则当前这个文件舍弃
|
|||
|
|
flag=1
|
|||
|
|
#为每个节点设置相关信息
|
|||
|
|
insert_point_count=0
|
|||
|
|
for e,dict in G.nodes.items():
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
asm_hex=e.block.bytes.hex().replace("0x","")
|
|||
|
|
#不能正确生成vector的函数不要
|
|||
|
|
# dict["feature"] = str(vec)
|
|||
|
|
#判断是否有插入点
|
|||
|
|
# dict["has_insertPoint"] = False
|
|||
|
|
#该basic_block的bytes指令
|
|||
|
|
|
|||
|
|
dict["asm2vec_plus_16"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_plus_16")))
|
|||
|
|
except Exception as ea:
|
|||
|
|
flag = 0
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
dict["asm2vec_plus_32"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_plus_32")))
|
|||
|
|
dict["asm2vec_plus_64"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_plus_64")))
|
|||
|
|
dict["asm2vec_plus_128"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_plus_128")))
|
|||
|
|
dict["asm2vec_plus_256"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_plus_256")))
|
|||
|
|
|
|||
|
|
dict["asm2vec_base_16"] = str(my_round(get_node_feature(hex_asm=asm_hex,node_feature_method="asm2vec_base_16")))
|
|||
|
|
dict["asm2vec_base_32"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_base_32")))
|
|||
|
|
dict["asm2vec_base_64"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_base_64")))
|
|||
|
|
dict["asm2vec_base_128"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_base_128")))
|
|||
|
|
dict["asm2vec_base_256"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_base_256")))
|
|||
|
|
|
|||
|
|
dict["asm2vec_s_base_16"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s_base_16")))
|
|||
|
|
dict["asm2vec_s_base_32"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s_base_32")))
|
|||
|
|
dict["asm2vec_s_base_64"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s_base_64")))
|
|||
|
|
dict["asm2vec_s_base_128"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s_base_128")))
|
|||
|
|
dict["asm2vec_s_base_256"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s_base_256")))
|
|||
|
|
|
|||
|
|
dict["asm2vec_s368_base_16"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s368_base_16")))
|
|||
|
|
dict["asm2vec_s368_base_32"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s368_base_32")))
|
|||
|
|
dict["asm2vec_s368_base_64"] = str( my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s368_base_64")))
|
|||
|
|
dict["asm2vec_s368_base_128"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s368_base_128")))
|
|||
|
|
dict["asm2vec_s368_base_256"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="asm2vec_s368_base_256")))
|
|||
|
|
|
|||
|
|
dict["n_gram"] = str(my_round(get_node_feature(hex_asm=asm_hex, node_feature_method="n_gram")))
|
|||
|
|
dict["malconv"]= str(my_round(get_node_feature(hex_asm=asm_hex,node_feature_method="malconv")))
|
|||
|
|
dict["word_frequency"]= str(my_round(get_node_feature(hex_asm=asm_hex,node_feature_method="word_frequency")))
|
|||
|
|
# dict["bytes"]=asm_hex
|
|||
|
|
|
|||
|
|
#用于记录当前basicblock插入func后，该func的地址值，初始值是原始程序的新建section的地址
|
|||
|
|
|
|||
|
|
|
|||
|
|
if flag==1:
|
|||
|
|
print("成功写入",os.path.join(CFG_dir,file_item)+".gexf")
|
|||
|
|
nx.write_gexf(G, os.path.join(CFG_dir,file_item)+".gexf")
|
|||
|
|
#写入节点数量与边数量信息
|
|||
|
|
nodes_num=len(cfg.graph.nodes())
|
|||
|
|
edgs_num=len(cfg.graph.edges())
|
|||
|
|
csv_data.append([file_item+".gexf",str(nodes_num),str(edgs_num)])
|
|||
|
|
|
|||
|
|
with open(csv_save_path,'w',encoding='utf-8',newline='') as fp:
|
|||
|
|
writer =csv.writer(fp)
|
|||
|
|
writer.writerow(header)
|
|||
|
|
writer.writerows(csv_data)
|
|||
|
|
|
|||
|
|
#用于直接读取gexf信息保存的代码
|
|||
|
|
def write_csv(CFG_dir = "../CFG_data/malware",csv_save_path="../CFG_data/malware_msg.csv",header = ['malware_name','nodes_num','edgs_num']):
|
|||
|
|
file_list=os.listdir(CFG_dir)
|
|||
|
|
# benign_list=os.listdir(bengin_CFG_dir)
|
|||
|
|
csv_data = []
|
|||
|
|
|
|||
|
|
for item_name in tqdm(file_list):
|
|||
|
|
file=os.path.join(CFG_dir,item_name)
|
|||
|
|
G=nx.read_gexf(file)
|
|||
|
|
nodes_num = len(G.nodes())
|
|||
|
|
edgs_num = len(G.edges())
|
|||
|
|
csv_data.append([item_name, str(nodes_num), str(edgs_num)])
|
|||
|
|
|
|||
|
|
with open(csv_save_path,'w',encoding='utf-8',newline='') as fp:
|
|||
|
|
writer =csv.writer(fp)
|
|||
|
|
writer.writerow(header)
|
|||
|
|
writer.writerows(csv_data)
|
|||
|
|
print("成功结束")
|
|||
|
|
# pass
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
malware_data_dir = "../data/malware"
|
|||
|
|
bengin_data_dir = "../data/benign"
|
|||
|
|
# pack_data_dir= "../data/pack"
|
|||
|
|
|
|||
|
|
malware_CFG_dir = "../cfg_data_with_feature/malware"
|
|||
|
|
bengin_CFG_dir = "../cfg_data_with_feature/benign"
|
|||
|
|
# pack_CFG_dir = "../cfg_data/pack"
|
|||
|
|
|
|||
|
|
malware_list = os.listdir(malware_data_dir)
|
|||
|
|
bengin_list = os.listdir(bengin_data_dir)
|
|||
|
|
# pack_list = os.listdir(pack_data_dir)
|
|||
|
|
# print(pack_list)
|
|||
|
|
# exit()
|
|||
|
|
#
|
|||
|
|
# cfg_extract(pack_list, data_dir=pack_data_dir, CFG_dir=pack_CFG_dir,
|
|||
|
|
# csv_save_path="../cfg_data/pack_msg.csv", header=['bengin_name', 'nodes_num', 'edgs_num'])
|
|||
|
|
|
|||
|
|
# cfg_extract(malware_list,data_dir=malware_data_dir,CFG_dir=malware_CFG_dir,csv_save_path="../cfg_data_with_feature/malware_msg1.csv",header = ['malware_name','nodes_num','edgs_num'])
|
|||
|
|
cfg_extract(bengin_list, data_dir=bengin_data_dir ,CFG_dir=bengin_CFG_dir ,csv_save_path="../cfg_data_with_feature/benign_msg.csv", header = ['bengin_name', 'nodes_num', 'edgs_num'])
|
|||
|
|
# write_csv(CFG_dir = "../CFG_data/malware_asm2vec_base",csv_save_path="../CFG_data/malware_msg_asm2vec_base.csv",header = ['malware_name','nodes_num','edgs_num','insert_point_count'])
|
|||
|
|
# write_csv(CFG_dir="../CFG_data/benign_asm2vec_base", csv_save_path="../CFG_data/benign_msg_asm2vec_base.csv",header=['benign_name', 'nodes_num', 'edgs_num','insert_point_count'])
|