210 lines
9.3 KiB
Python
210 lines
9.3 KiB
Python
from collections import defaultdict
|
|
from tqdm import tqdm
|
|
import pandas as pd
|
|
import os
|
|
import csv
|
|
import argparse
|
|
import statistics
|
|
import plotly.express as px
|
|
|
|
###################################################################################################
|
|
## Program shall take two csv files of different classes - benign and malware
|
|
## It will compute ngrams for each of the classes seperately and find the delta frequencies
|
|
## of each computed ngram. delta_frequencies = (class1 - class2)
|
|
###################################################################################################
|
|
|
|
#--------------------------------------------------------------------------------------------------
|
|
## 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])
|
|
# Selects high frequency ngram until the mean value
|
|
# Returns both complete and filtered dictionary of ngrams
|
|
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"):
|
|
opcodes = corpusList[idx]
|
|
for item in generate_N_grams(opcodes, N):
|
|
#compute frequency of all unique ngrams
|
|
if len(opcodes) == 0:
|
|
continue
|
|
ngramDictionary[item] += 1
|
|
#compute ngram file count
|
|
for item in ngramDictionary:
|
|
ngramFileCount[item] += 1
|
|
|
|
filteredNgramDictionary = defaultdict(int)
|
|
#Filter those ngrams which meet percent of Total files criteria
|
|
filterCnt = round(int((percent * total)/ 100), 0)
|
|
for item in ngramFileCount:
|
|
if ngramFileCount[item] >= filterCnt:
|
|
#Add to filtered dictionary the item which meets file count criteria
|
|
filteredNgramDictionary[item] = ngramDictionary[item]
|
|
|
|
#Filter ngram with a minimum frequency
|
|
if (filterFreq):
|
|
for item in ngramDictionary:
|
|
if ngramDictionary[item] < filterFreq and item in filteredNgramDictionary:
|
|
#Remove the item which below the frequency threshold
|
|
filteredNgramDictionary.pop(item)
|
|
|
|
#print(f"Total ngrams:{len(ngramDictionary.items())} => filtered: {len(filteredNgramDictionary.items())}\n")
|
|
return [ngramDictionary, filteredNgramDictionary]
|
|
|
|
#--------------------------------------------------------------------------------------------------
|
|
# Calculate a normalization factor for frequency values of class1 and class2
|
|
# For class which are high in frequency due their sample size, a normalization may required to be
|
|
# factored for correctly resizgin the frequencies of the small class set.
|
|
# input list of frequencies of class1 and class 2
|
|
def normalization_factor(class1, class2):
|
|
mean1 = statistics.mean(class1)
|
|
mean2 = statistics.mean(class2)
|
|
return mean1/mean2
|
|
|
|
#--------------------------------------------------------------------------------------------------
|
|
# Write the data into the given csv file handle
|
|
def WriteCSV (file, csvFields, dataDictionary):
|
|
writer = csv.DictWriter(file, fieldnames=csvFields)
|
|
writer.writeheader()
|
|
writer.writerows(dataDictionary)
|
|
|
|
#--------------------------------------------------------------------------------------------------
|
|
# 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
|
|
|
|
|
|
# Get user arguments
|
|
malware_csvfile = os.path.join('./out/output_malware.csv')
|
|
benign_csvfile = os.path.join('./out/output_benign.csv')
|
|
maxgrams = 3
|
|
|
|
# Error check and exit if not a file
|
|
if not (os.path.isfile(malware_csvfile) and os.path.isfile(benign_csvfile)):
|
|
print (f"Path should be csv file!")
|
|
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
|
|
frequencyFilter = 20 ## select ngrams with frequency greater than this value
|
|
|
|
malwareNgram = defaultdict(int) ## full list of ngrams in malware corpus
|
|
benignNgram = defaultdict(int) ## full list of ngrams in benign corpus
|
|
filteredMalwareNgram = defaultdict(int) ## filtered list of ngrams from malware corpus
|
|
filteredBenignNgram = defaultdict(int) ## filtered list of ngrams from benign corpus
|
|
|
|
## 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 ...")
|
|
malwareNgram.clear()
|
|
filteredMalwareNgram.clear()
|
|
benignNgram.clear()
|
|
filteredBenignNgram.clear()
|
|
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)
|
|
|
|
[benignNgram, filteredBenignNgram] = filter_N_grams(benignDF['corpus'].values, idx+1,
|
|
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.
|
|
for item in mergedList:
|
|
key = item #get the ngram only
|
|
if key in filteredBenignNgram:
|
|
if key in filteredMalwareNgram:
|
|
filteredMergedNgram[key] = filteredBenignNgram[key] - filteredMalwareNgram[key]
|
|
elif item in malwareNgram:
|
|
filteredMergedNgram[key] = filteredBenignNgram[key] - malwareNgram[key]
|
|
else:
|
|
filteredMergedNgram[key] = filteredBenignNgram[key]
|
|
elif key in filteredMalwareNgram:
|
|
if key in benignNgram:
|
|
filteredMergedNgram[key] = benignNgram[key] - filteredMalwareNgram[key]
|
|
else:
|
|
filteredMergedNgram[key] = filteredMalwareNgram[key]
|
|
|
|
print(f"Merged: {idx+1}gramCnt={len(filteredMergedNgram.keys())}")
|
|
## get a sorted list of merged ngrams with relative frequencies
|
|
sortedMergedNgramList = sorted(filteredMergedNgram.items(), key=lambda x: x[1])
|
|
|
|
#Plot a scatter graph -
|
|
# y values as relative frequency benign-malware
|
|
# x values as max frequency of a ngram max(malware, benign)
|
|
# 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)
|
|
|
|
#write the final ngrams into a file for feature selection
|
|
ngramDictList = []
|
|
for item in sortedMergedNgramList:
|
|
dictItem = {}
|
|
key = item[0]
|
|
dictItem['ngram'] = key
|
|
dictItem['count'] = max(filteredMalwareNgram[key], filteredBenignNgram[key])
|
|
ngramDictList.append(dictItem)
|
|
|
|
csvfields = ['ngram', 'count']
|
|
csvname = str(idx+1) + "gram.csv"
|
|
try:
|
|
csvfile = open(csvname, 'w')
|
|
except Exception as err:
|
|
print(f"Error: writing csvfile {err}")
|
|
WriteCSV(csvfile, csvfields, ngramDictList)
|
|
csvfile.close()
|
|
|