线程池版本
This commit is contained in:
parent
52b0cf6db3
commit
3f4bde2989
18
OpcodeGet.py
18
OpcodeGet.py
@ -6,12 +6,14 @@ from tqdm import tqdm
|
||||
import r2pipe
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def Opcode_to_csv(opcode_list, file_type):
|
||||
logger.info("*======================start write==================================*")
|
||||
csv_write(f'output_{file_type}.csv', opcode_list)
|
||||
logger.info(f"done {done_file_num} files")
|
||||
logger.info("*=================write to csv success==============================*")
|
||||
|
||||
|
||||
def csv_write(file_name, data: list):
|
||||
"""write data to csv"""
|
||||
df = pd.DataFrame(data)
|
||||
@ -19,6 +21,8 @@ def csv_write(file_name, data: list):
|
||||
for i in range(0, len(df), chunksize):
|
||||
df.iloc[i:i + chunksize].to_csv(f'./out/{file_name}', mode='a', header=False, index=False)
|
||||
return True
|
||||
|
||||
|
||||
def extract_opcode(disasm_text):
|
||||
"""
|
||||
从反汇编文本中提取操作码和操作数
|
||||
@ -34,6 +38,7 @@ def extract_opcode(disasm_text):
|
||||
return opcode
|
||||
return ""
|
||||
|
||||
|
||||
def get_graph_r2pipe(r2pipe_open, file_type):
|
||||
# 获取基础块内的操作码序列
|
||||
opcode_Sequence = []
|
||||
@ -66,16 +71,15 @@ def get_graph_r2pipe(r2pipe_open, file_type):
|
||||
if op["type"] == "invalid":
|
||||
continue
|
||||
block_opcode_Sequence.append(extract_opcode(op["opcode"]))
|
||||
opcode_Sequence.append([file_type, file_type, len(block_opcode_Sequence), ' '.join(block_opcode_Sequence)])
|
||||
opcode_Sequence.append(
|
||||
[file_type, file_type, len(block_opcode_Sequence), ' '.join(block_opcode_Sequence)])
|
||||
except:
|
||||
print("Error: get function list failed")
|
||||
return opcode_Sequence
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
logger = setup_logger('logger', 'log/opcode_benign.log')
|
||||
logger = setup_logger('logger', './log/opcode_benign.log')
|
||||
file_type = 'benign'
|
||||
file_path = os.path.join('/mnt/d/bishe/dataset/train_benign')
|
||||
file_list = os.listdir(file_path)[:10000]
|
||||
@ -94,12 +98,6 @@ if __name__ == '__main__':
|
||||
else:
|
||||
csv_write(f'output_{file_type}.csv', done_list)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# node_list = []
|
||||
# edge_list = []
|
||||
# temp_edge_list = []
|
||||
|
||||
110
ngram.py
110
ngram.py
@ -6,6 +6,11 @@ import csv
|
||||
import argparse
|
||||
import statistics
|
||||
import plotly.express as px
|
||||
import concurrent.futures
|
||||
from functools import partial
|
||||
import logging
|
||||
import contextlib
|
||||
|
||||
|
||||
###################################################################################################
|
||||
## Program shall take two csv files of different classes - benign and malware
|
||||
@ -16,12 +21,12 @@ import plotly.express as px
|
||||
# --------------------------------------------------------------------------------------------------
|
||||
## Generate ngrams given the corpus and factor n
|
||||
def generate_N_grams(corpus, n=1):
|
||||
|
||||
words = [word for word in corpus.split(" ")]
|
||||
temp = zip(*[words[i:] for i in range(0, n)])
|
||||
ngram = [' '.join(n) for n in temp]
|
||||
return ngram
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------
|
||||
## Creates ngrams for the corpus List for given N and Filters it based on following criteria
|
||||
# file count >= percent of Total corpus len (pecent in [1..100])
|
||||
@ -31,8 +36,10 @@ def filter_N_grams (corpusList, N, percent, filterFreq=0):
|
||||
total = len(corpusList)
|
||||
ngramDictionary = defaultdict(int)
|
||||
ngramFileCount = defaultdict(int)
|
||||
for idx in tqdm(range(0, total), ncols=100, desc="Computing ngrams"):
|
||||
for idx in range(0, total):
|
||||
opcodes = corpusList[idx]
|
||||
if type(opcodes) is not str:
|
||||
continue
|
||||
for item in generate_N_grams(opcodes, N):
|
||||
# compute frequency of all unique ngrams
|
||||
if len(opcodes) == 0:
|
||||
@ -58,7 +65,8 @@ def filter_N_grams (corpusList, N, percent, filterFreq=0):
|
||||
filteredNgramDictionary.pop(item)
|
||||
|
||||
# print(f"Total ngrams:{len(ngramDictionary.items())} => filtered: {len(filteredNgramDictionary.items())}\n")
|
||||
return [ngramDictionary, filteredNgramDictionary]
|
||||
return ngramDictionary, filteredNgramDictionary
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------
|
||||
# Calculate a normalization factor for frequency values of class1 and class2
|
||||
@ -70,6 +78,7 @@ def normalization_factor(class1, class2):
|
||||
mean2 = statistics.mean(class2)
|
||||
return mean1 / mean2
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------
|
||||
# Write the data into the given csv file handle
|
||||
def WriteCSV(file, csvFields, dataDictionary):
|
||||
@ -77,18 +86,43 @@ def WriteCSV (file, csvFields, dataDictionary):
|
||||
writer.writeheader()
|
||||
writer.writerows(dataDictionary)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
def process_csv_file(csvfile, ngram_type, file_percent_filter, frequency_filter):
|
||||
"""处理CSV文件并并行计算n-gram"""
|
||||
print(f"start load csv file:{os.path.basename(csvfile)}")
|
||||
dataframe = pd.read_csv(csvfile, encoding="utf8")
|
||||
print(f"end load")
|
||||
ngram_list = defaultdict(int)
|
||||
filtered_ngram_list = defaultdict(int)
|
||||
process_bar = tqdm(total=len(dataframe['corpus'].values), desc=f'Computing {ngram_type}-gram on files')
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor: # 调整线程池大小
|
||||
future_to_args = {
|
||||
executor.submit(filter_N_grams, dataframe['corpus'].values[start: start + 10000],
|
||||
idx + 1, file_percent_filter, frequency_filter): start for start in
|
||||
range(0, len(dataframe['corpus'].values), 10000)
|
||||
}
|
||||
for future in concurrent.futures.as_completed(future_to_args):
|
||||
try:
|
||||
sub_ngram_list, sub_filtered_ngram_list = future.result()
|
||||
for i in [sub_ngram_list, ngram_list]:
|
||||
for key, value in i.items():
|
||||
ngram_list[key] += value
|
||||
for i in [sub_filtered_ngram_list, filtered_ngram_list]:
|
||||
for key, value in i.items():
|
||||
filtered_ngram_list[key] += value
|
||||
process_bar.update(10000) # 手动更新进度条
|
||||
except Exception as exc:
|
||||
logging.error(f"Error processing {idx + 1}-gram: {exc}")
|
||||
return ngram_list, filtered_ngram_list
|
||||
|
||||
# --------------------------------------------------------------------------------------------------
|
||||
# Execution starts here
|
||||
# Add command line arguments
|
||||
# CSV header: class, sub-class, size, corpus
|
||||
parser = argparse.ArgumentParser(description="ngram analysis on a given corpus csv file.")
|
||||
parser.add_argument('malware_csvfile', help='path to the malware corpus csv file')
|
||||
parser.add_argument('benign_csvfile', help='path to the benign corpus csv file')
|
||||
parser.add_argument('ngram', help='ngram to compute, higher value will be compute intensive')
|
||||
|
||||
# Execute the parse_args() method
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Get user arguments
|
||||
malware_csvfile = os.path.join('./out/output_malware.csv')
|
||||
benign_csvfile = os.path.join('./out/output_benign.csv')
|
||||
@ -100,11 +134,6 @@ if not (os.path.isfile(malware_csvfile) and os.path.isfile(benign_csvfile)):
|
||||
exit(1)
|
||||
|
||||
# Read the csv file using pandas into data frame
|
||||
try:
|
||||
malwareDF = pd.read_csv(malware_csvfile, encoding = "utf8")
|
||||
benignDF = pd.read_csv(benign_csvfile, encoding="utf8")
|
||||
except Exception as error:
|
||||
print(error)
|
||||
|
||||
# Build a frequency list for ngrams
|
||||
filePercentFilter = 80 ## select ngrams present in x% of files
|
||||
@ -118,10 +147,14 @@ filteredBenignNgram = defaultdict(int) ## filtered list of ngrams from benign
|
||||
## common list ngrams from both malware and benign corpus with relative frequency (benignFreq - malwareFreq)
|
||||
filteredMergedNgram = defaultdict(int)
|
||||
|
||||
|
||||
# run for only the maxgram provided, change lower value to 0 to run for all values [1..N]
|
||||
for idx in range(maxgrams - 1, maxgrams):
|
||||
print(f"Computing {idx + 1}gram on files ...")
|
||||
print(f"CPU core {os.cpu_count()} on use")
|
||||
malwareNgram = []
|
||||
filteredMalwareNgram = []
|
||||
benignNgram = []
|
||||
filteredBenignNgram = []
|
||||
malwareNgram.clear()
|
||||
filteredMalwareNgram.clear()
|
||||
benignNgram.clear()
|
||||
@ -129,19 +162,13 @@ for idx in range(maxgrams-1, maxgrams):
|
||||
filteredMergedNgram.clear()
|
||||
|
||||
# opcodes decoded from pe file in sequence is stored as corpus in the csv
|
||||
[malwareNgram, filteredMalwareNgram] = filter_N_grams(malwareDF['corpus'].values, idx+1,
|
||||
filePercentFilter, frequencyFilter)
|
||||
malwareNgram, filteredMalwareNgram = process_csv_file(malware_csvfile, 'malware', filePercentFilter, frequencyFilter)
|
||||
|
||||
[benignNgram, filteredBenignNgram] = filter_N_grams(benignDF['corpus'].values, idx+1,
|
||||
filePercentFilter, frequencyFilter)
|
||||
benignNgram, filteredBenignNgram = process_csv_file(benign_csvfile, 'benign', filePercentFilter, frequencyFilter)
|
||||
|
||||
# creates a sorted list of ngram tuples with their frequency for 1 .. maxgram
|
||||
print(f"Malware: {idx+1}gramCnt={len(malwareNgram.items())}, filterenCnt={len(filteredMalwareNgram.items())}")
|
||||
print(f"Benign: {idx+1}gramCnt={len(benignNgram.items())}, filterenCnt={len(filteredBenignNgram.items())}")
|
||||
|
||||
## Make a intersection of filtered list between malware and benign ngrams
|
||||
mergedList = list(set().union(filteredMalwareNgram.keys(), filteredBenignNgram.keys()))
|
||||
|
||||
## Now find the relative frequency b/w benign and malware files. = benign - malware
|
||||
## write this for cases where ngrams only present in one of the clases malware or benign
|
||||
## for reusability in case a union of classes is taken.
|
||||
@ -170,24 +197,23 @@ for idx in range(maxgrams-1, maxgrams):
|
||||
# color labels as 'a' + frequency % 26
|
||||
# size as frequency/max * 100
|
||||
# hover name is ngram name
|
||||
titlestr = str(idx+1) + "gram: Total samples(" + str(len(sortedMergedNgramList)) + ")"
|
||||
htmlfile = str (idx+1) +"gram.html"
|
||||
hovername = [item[0] for item in sortedMergedNgramList]
|
||||
yval = [item[1] for item in sortedMergedNgramList]
|
||||
xval = []
|
||||
for key in hovername:
|
||||
xval.append(max(filteredMalwareNgram[key], filteredBenignNgram[key]))
|
||||
colors = [chr(ord('a')+ (value %26)) for value in xval]
|
||||
maxval = max(xval)
|
||||
sizeval = [(int((val/maxval)*100)+1) for val in xval]
|
||||
|
||||
fig = px.scatter(title=titlestr, y=yval, x=xval, color=colors,
|
||||
size=sizeval, hover_name=hovername, log_x=True,
|
||||
labels = {
|
||||
"x": "Absolute Frequency",
|
||||
"y": "Relative Frequency"})
|
||||
fig.show()
|
||||
fig.write_html(htmlfile)
|
||||
# titlestr = str(idx + 1) + "gram: Total samples(" + str(len(sortedMergedNgramList)) + ")"
|
||||
# htmlfile = str(idx + 1) + "gram.html"
|
||||
# hovername = [item[0] for item in sortedMergedNgramList]
|
||||
# yval = [item[1]/1e10 for item in sortedMergedNgramList]
|
||||
# xval = []
|
||||
# for key in hovername:
|
||||
# xval.append(max(filteredMalwareNgram[key], filteredBenignNgram[key]))
|
||||
# colors = [chr(ord('a') + (value % 26)) for value in xval]
|
||||
# maxval = max(xval)
|
||||
# sizeval = [(int((val / maxval) * 100) + 1) for val in xval]
|
||||
#
|
||||
# fig = px.scatter(title=titlestr, y=yval, x=xval, color=colors,
|
||||
# size=sizeval, hover_name=hovername, log_x=True,
|
||||
# labels={
|
||||
# "x": "Absolute Frequency",
|
||||
# "y": "Relative Frequency"})
|
||||
# fig.write_html(htmlfile)
|
||||
|
||||
# write the final ngrams into a file for feature selection
|
||||
ngramDictList = []
|
||||
@ -199,11 +225,11 @@ for idx in range(maxgrams-1, maxgrams):
|
||||
ngramDictList.append(dictItem)
|
||||
|
||||
csvfields = ['ngram', 'count']
|
||||
csvname = str(idx+1) + "gram.csv"
|
||||
csvname = "./out/"+str(idx + 1) + "gram.csv"
|
||||
print("*======================start write csv=======================================*")
|
||||
try:
|
||||
csvfile = open(csvname, 'w')
|
||||
except Exception as err:
|
||||
print(f"Error: writing csvfile {err}")
|
||||
WriteCSV(csvfile, csvfields, ngramDictList)
|
||||
csvfile.close()
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user