ML_classify.py

# Script to train/test SKLEARN Multi label techniques using RandomForest with D5

import os
import sys, time
import loaders_binary as ll
import loaders_multilabel as mll
from scipy.sparse import csr_matrix
import pandas as pd


def lst_df(ft_train, ft_test, labelstrain, labelstest):
	print(ft_train)
	print(len(ft_train), len(labelstrain), len(labelstest))
	print(labelstrain)
	traindf = pd.DataFrame(ft_train)
	trainlabelsdf = pd.DataFrame(labelstrain)
	featdf_train = pd.concat([traindf, trainlabelsdf], axis=1, ignore_index=True)
	print("TRAIN DF:\n", featdf_train, featdf_train.shape)
	testdf = pd.DataFrame(ft_test)
	testlabelsdf = pd.DataFrame(labelstest)
	featdf_test = pd.concat([testdf, testlabelsdf], axis=1, ignore_index=True)
	ftsindex = len(featdf_train.columns)-len(labelstrain[0])
	print("TEST DF:\n", featdf_test, featdf_test.shape)
	return featdf_train, ftsindex, featdf_test

# Split connections from PCAPs: 70-30%: train-test (D5)
def getfpaths(fourone, threetwo):
	train = []
	test = []
	for i in fourone:
		print(i)
		for j in range(0,4):
			print(j, i+"-"+str(j)+".cell")
			train += [i+"-"+str(j)+".cell"]
		test += [i+".cell"]
		print("TEST: ", i+".cell")

	for k in threetwo:
		print(k)
		for m in range(1, 4):
			print(m, k+"-"+str(m)+".cell")
			train += [k+"-"+str(m)+".cell"]
		test += [k+".cell", k+"-0.cell"]
		print("TEST: ", k+".cell", k+"-0.cell")

	print("Total PCAPs train: ",len(train), ";test: ",len(test))
	print(train, test)
	return train, test

def pcapsplit(feats, labels, fnames):
	trainconns = []
	testconns = []
	# 70-30 train-test split for PCAPs of D5 binaries (4:1/3:2)
	fourone, threetwo = mll.get_binary_splits("splitbinaries.txt")
	trainpcaps, testpcaps = getfpaths(fourone, threetwo)

	trainlabels = []
	testlabels = []
	i = 0
	for ft, fname in zip(feats, fnames):
		print(labels[i], fname)
		if fname in trainpcaps:
			trainconns += [ft]
			trainlabels += [list(labels[i])]
		else:
			assert fname in testpcaps
			testconns += [ft]
			testlabels += [list(labels[i])]
		i += 1
	print("Total train connections: ", len(trainconns))
	print("Total test connections: ", len(testconns))
	#print(testconns)
	return trainconns, trainlabels, testconns, testlabels

def main_ml(d, malfnames, benfnames, foldtotal, maltotal, malinst, topk, hostfts, multiclass):
	# Requires 2 files: Binary index map: malwarehash->int & AVCLASS o/p: malwarehash->multi labels
	proceed = ll.check_dependencies(d)
	train_endindx = int((maltotal*70)/100)
	#assert train_endindx == 109
	if proceed:
		[avclassfile, reports] = ll.getfiles(d)
        # {shahash: avclass labels}
		classdt = ll.read_multilabels(avclassfile, reports)
        # {malwareindex: [labels, malware shahash]}
		sha_label_map = ll.sha_mapi(d, classdt)
		train_i = set()
		test_i = set()
		labels = []
		for mi, val in sha_label_map.items():
			label = val[0]
			for l in label:
				if mi <= train_endindx:
					train_i.add(l)
				else:
					test_i.add(l)
			print(mi, val[0])
			for l in label:
				if l not in labels:
					labels += [l]

		print("All classes: ", labels)
		print("Train classes covered: ", train_i)
		print("Test classes covered: ", test_i)

		# GRAPH 1: Tag distribution per malware: Frequency/count of each tag in malware dataset
		# STATS: min tag per mal, max tag per mal, avg no. of tags per mal
		mll.label_stats(sha_label_map)
		time.sleep(2)
        # dt: {fpath: [0,1,0,1,1]} based on MultilabelBinarizer
		[labeldt, classorder, mlb] = ll.label_multiclass_multilabel(malfnames, benfnames, sha_label_map)
		#print(classorder, len(classorder))
		#print(labeldt, len(labeldt))

        # 3. feature extraction
		[featdf, labeldf, fnames] = ll.extract_features(labeldt, multiclass, hostfts, top=topk, classorder=classorder)
		# Train-test PCAP split
		ft_train, labels_train, ft_test, labels_test = pcapsplit(featdf, labeldf, fnames)

		# 4. Get dataframes for classifier input
		[featdf_train, ftsindex, featdf_test] = lst_df(ft_train, ft_test, labels_train, labels_test)
		labels = None
        # Classify
		print("Features end at index: ", ftsindex)
		time.sleep(2)
		labels_i = len(list(classorder))
		ytrain = featdf_train.iloc[:,ftsindex:].to_numpy()
		ytest = featdf_test.iloc[:,ftsindex:].to_numpy()
		Xdatatrain = featdf_train.iloc[:,:ftsindex].to_numpy()
		Xdatatest = featdf_test.iloc[:,:ftsindex].to_numpy()
		# Dataframe info
		print("Features(TRAIN): \n", Xdatatrain, Xdatatrain.shape)
		print("Labels(TRAIN): \n", ytrain, ytrain.shape)
		time.sleep(2)
		print("Features(TEST): \n", Xdatatest, Xdatatest.shape)
		print("Labels(TEST): \n", ytest, ytest.shape)
		# Data normalization and standardization
		##Xtrainscaled, Xtestscaled = mll.normalize(Xtrain, Xtest)
		Xtrainscaled, Xtestscaled, ss = mll.standardize(Xdatatrain, Xdatatest)
		#print("Xtrainscaled:\n", Xtrainscaled)
		#print("Xtestscaled: \n", Xtestscaled)

		# Classifiers & ML Techniques

		# 1. Base: rf, knn
		print("Base random forest:: ")
		mll.base_rf(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb)
		#mll.base_knn(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb)

		# 2. OneVsRest: rf, knn, lr
		# RESULT: Bad performance for initial exp
		##mll.OneVsRest(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb, "knn")
		##mll.OneVsRest(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb)
		##mll.OneVsRest(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb, "rf")

		# 3. MultiOutputClassification: One classifier per target for classifiers not native to multi class classification
		##mll.multioutputLR(Xtrainscaled, ytrain, Xtestscaled, ytest)

		# ------ Problem Transformation --------
		# 3. Binary Relevance: rf, knn, NB, LR
		brr1=[]
		brr2=[]
		print("Technique: Binary Relevance::")
		# BRKNN: Classifier assigns labels assigned to >= half of the neighbors
		##mll.BRKNNA(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb)
		brr1 = mll.BR(ss, Xtrainscaled, ytrain, Xtestscaled, ytest, mlb, classorder, topk, "rf")
		#brr2 = mll.BR(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb, "knn")
		print("Technique: ClassifierChains::")
		cc1=[]
		cc2=[]
		# 4. Classifier Chains: rf, knn
		cc1 = mll.ClassifierChain(ss, Xtrainscaled, ytrain, Xtestscaled, ytest, mlb, classorder, topk, "rf")
		#cc2 = mll.ClassifierChain(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb, "knn")
		print("Technique: Label Powerset::")
		# 5. Labelpowerset: rf, knn
		lp1=[]
		lp2=[]
		lp1 = mll.labelpowerset(ss, Xtrainscaled, ytrain, Xtestscaled, ytest, mlb, classorder, topk, "rf")
		#lp2 = mll.labelpowerset(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb,"knn")
		##print("Label order: ", classorder, len(classorder))

		# 6. Adapted Knn- sklearn:  uses k-NearestNeighbors find nearest examples to a test class and uses Bayesian inference to select labels*
		##mll.adaptedknn(Xtrainscaled, ytrain, Xtestscaled, ytest)


	return [brr1,cc1,lp1]