loaders_multilabel.py

# Script implementing multi label classification techniques

import sys, os
import matplotlib
import time
import pandas as pd
import ast
import numpy as np
import matplotlib.pyplot as plt
matplotlib.use("Agg")
#Classifiers
from skmultilearn.adapt import MLkNN
from sklearn.neighbors import KNeighborsClassifier as knnbase
from sklearn.multioutput import MultiOutputClassifier as MOC
from sklearn.ensemble import RandomForestClassifier as rf
from sklearn.naive_bayes import MultinomialNB as mnb
from skmultilearn.ensemble import MajorityVotingClassifier as mvc
from skmultilearn.ensemble import LabelSpacePartitioningClassifier as lspc
from skmultilearn.adapt import BRkNNaClassifier as knnA
from sklearn.multiclass import OneVsRestClassifier as OVR
from sklearn.linear_model import LogisticRegression as LR
from sklearn.naive_bayes import GaussianNB as GNB
# Multilabel techniques
from sklearn.preprocessing import MultiLabelBinarizer
from skmultilearn.problem_transform import BinaryRelevance
from skmultilearn.problem_transform import ClassifierChain as chain
from skmultilearn.problem_transform import LabelPowerset as power
# Data split types
from sklearn.model_selection import GridSearchCV as GSCV
from sklearn.model_selection import train_test_split
# Normalization and Scaling
from sklearn.preprocessing import MinMaxScaler as MMS
from sklearn.preprocessing import StandardScaler as SS
# Evaluation
from sklearn.metrics import accuracy_score, hamming_loss, precision_score, recall_score, f1_score
from sklearn.metrics import classification_report
from sklearn.metrics import multilabel_confusion_matrix as ML_matrix
from sklearn.metrics import precision_recall_fscore_support as score_multi
from sklearn.metrics import classification_report
from pickle import load, dump

# -------------------------------- HELPERS ------------------------------------------ #
def split_df(Xdata, labels, testsplit=0.3):
	Xtrain,Xtest,ytrain,ytest = train_test_split(Xdata,labels,test_size=testsplit)
	return Xtrain, Xtest, ytrain, ytest

# Rescale values to fit in a range; default: 0-1
def normalize(Xtrain, Xtest):
	scaler = MMS(feature_range=(0,1))
	Xtrainscaled = scaler.fit_transform(Xtrain)
	Xtestscaled = scaler.transform(Xtest)
	return Xtrainscaled, Xtestscaled

# Scale values such that mean = 0, std dev. = 1; Ensures robustness for new data.
def standardize(Xtrain, Xtest):
	ss = SS()
	Xtrainscaled = ss.fit_transform(Xtrain)
	Xtestscaled = ss.transform(Xtest)
	return Xtrainscaled, Xtestscaled, ss

def micro_avg(y_test_multilabel, predictions):
	precision = precision_score(y_test_multilabel, predictions, average='micro')
	recall = recall_score(y_test_multilabel, predictions, average='micro')
	f1 = f1_score(y_test_multilabel, predictions, average='micro')

	print("::Micro-average::")
	print("Precision: {:.4f}, Recall: {:.4f}, F1-measure: {:.4f}".format(precision, recall, f1))
	print("\n\n")
	return precision, recall, f1

def macro_avg(y_test_multilabel, predictions):
	precision = precision_score(y_test_multilabel, predictions, average='macro')
	recall = recall_score(y_test_multilabel, predictions, average='macro')
	f1 = f1_score(y_test_multilabel, predictions, average='macro')

	print("\nMacro-average: ")
	print("Precision: {:.4f}, Recall: {:.4f}, F1-measure: {:.4f}".format(precision, recall, f1))
	return

def per_class_dist(ytest, ypred, classorder):
	perclass = classification_report(ytest, ypred)
	print("Per class classification report: ", perclass)
	precision, recall, fscore, support = score_multi(ytest, ypred, average="micro")
	print('micro-precision: {}'.format(precision))
	print('micro-recall: {}'.format(recall))
	print('micro-fscore: {}'.format(fscore))
	print('support: {}'.format(support))
	#print(classorder)
	return

def get_binary_splits(fname="splitbinaries.txt"):
	fourone = []
	threetwo = []
	with open(fname, "r") as f:
		for line in f.readlines():
			if fourone == []:
				fourone = ast.literal_eval(line.rstrip())
			else:
				threetwo = ast.literal_eval(line.rstrip())

	print("Binaries with train:test-4:1 split", fourone, "\n3:2 split", threetwo)
	return fourone, threetwo

# ---------------------------- CLASSIFIERS & ML TECHNIQUES ------------------------- #
def base_rf(Xtrain, ytrain, Xtest, ytest, mlb):
	model = rf(n_estimators= 1000, n_jobs= -1).fit(Xtrain, ytrain)
	ypred = model.predict(Xtest)
	ypredproba = model.predict_proba(Xtest)
	print("Accuracy score: ", accuracy_score(ytest, ypred))
	print("True labels: ", ytest)
	print("Predicted: ", ypred)
	hloss = hamming_loss(ytest, ypred)
	print("Hloss: ", hloss)
	micro_avg(ytest, ypred)
	#print(ypredproba)
	return

def base_knn(Xtrain, ytrain, Xtest, ytest, mlb):
	model = knnbase(n_neighbors=3, n_jobs= -1).fit(Xtrain, ytrain)
	ypred = model.predict(Xtest)
	print("Accuracy score knn: ", accuracy_score(ytest, ypred))
	print("True labels: ", ytest)
	print("Predicted: ", ypred)
	micro_avg(ytest, ypred)
	return

def OneVsRest(Xtrain, ytrain, Xtest, ytest, mlb, ctype="lr"):
	if ctype == "knn":
		print("OneVsRest KNN")
		model = OVR(knnbase(n_neighbors= 3), n_jobs=-1)
	elif ctype == "lr":
		print("OneVsRest LR")
		model = OVR(LR(class_weight= "balanced"), n_jobs=-1)
	else:
		# rf
		print("OneVsRest RF")
		model = OVR(rf(n_estimators= 1000), n_jobs=-1)

	classifier = model.fit(Xtrain, ytrain)
	ypred = classifier.predict(Xtest)
	print("True labels: \n", ytest)
	print("Predicted labels: \n ", ypred)
	score = accuracy_score(ytest, ypred)
	print("Accuracy: ", score)
	micro_avg(ytest, ypred)
	return

def adaptedknn(Xtrainscaled, ytrain, Xtestscaled, ytest):
	print("Classifier: Adapted Knn")
	scores = dict()
	for kval in range(2,15):
		print("ML KNN, k=", kval)
		classifier = MLkNN(k=kval).fit(Xtrainscaled, ytrain)
		labelstest_pred = classifier.predict(Xtestscaled)
		labeltestpred_prob = classifier.predict_proba(Xtestscaled)
		score = accuracy_score(ytest, labelstest_pred)
		hloss = hamming_loss(ytest, labelstest_pred)
		prec, rec, f1 = micro_avg(ytest, labelstest_pred)
		scores[kval] = [score, hloss, prec, rec, f1]
		print("Accuracy Adapted Knn: ", score)
		print("Hamming loss: ", hloss)
	print(scores)
	return classifier

def multioutputLR(Xtrainscaled, ytrain, Xtestscaled, ytest):
	clf = MOC(LR(class_weight= "balanced")).fit(Xtrainscaled, ytrain)
	predicted = clf.predict(Xtestscaled)
	result = clf.score(Xtestscaled, ytest)
	print("Multi Output- LR: ", result)
	print("True labels: \n", ytest)
	print("Predicted labels: \n ", predicted)
	print("Confusion Matrix\n: ", ML_matrix(ytest, predicted))
	return

def BRKNNA(Xtrainscaled, ytrain, Xtestscaled, ytest, mlb, k=3):
	#knnA
	print("BR Knn, k=3")
	classifier = knnA(k=3)
	classifier.fit(Xtrainscaled, ytrain)
	labelstest_pred = classifier.predict(Xtestscaled)
	score = accuracy_score(ytest, labelstest_pred)
	#score = classifier.score(Xtestscaled, ytest)
	print("Accuracy BRKnn: ", score)
	print("Hamming loss: ", hamming_loss(ytest, labelstest_pred))
	print("True labels: \n", ytest)
	print("Predicted labels: \n ", labelstest_pred)
	micro_avg(ytest, labelstest_pred)
	#macro_avg(ytest, labelstest_pred)

	params = {'k': range(1,3)}
	score = 'f1_macro'
	clf = GSCV(knnA(), params, scoring=score)
	predicted = clf.fit(Xtrainscaled, ytrain)
	print("GridSearch KnnA: Best params: ", clf.best_params_, " Best score: ", clf.best_score_)
	#print("Confusion Matrix\n: ", ML_matrix(ytest, labelstest_pred))
	return


def BR(ss, Xtrain, ytrain, Xtest, ytest, mlb, labels, top, ctype="rf"):
	print("Classifier: ", ctype)
	if ctype == "nb":
		model = GNB()
	elif ctype == "lr":
		model = LR(class_weight= "balanced")
	elif ctype == "knn":
		model = knnbase(n_neighbors=3, n_jobs= -1)
	else:
		model = rf(n_estimators= 1000, n_jobs= -1)

	br = BinaryRelevance(model).fit(Xtrain, ytrain)
	ypred = br.predict(Xtest)
	save_model(mlb, ss, br, "br_"+str(top))
	acc = accuracy_score(ytest, ypred)
	hloss = hamming_loss(ytest, ypred)
	print("Accuracy score Binary Relevance: ", acc)
	print("Hamming loss: ", hloss)
	mprec, mrecall, mf1 = micro_avg(ytest, ypred)
	#per_class_dist(ytest, ypred, labels)
	#macro_avg(ytest, ypred)
	return [acc, hloss, mprec, mrecall, mf1, ctype, model]

def ClassifierChain(ss, Xtrain, ytrain, Xtest, ytest, mlb, labels, top, ctype="rf"):
	print("Classifier: ", ctype)
	if ctype == "knn":
		model = knnbase(n_neighbors=3, n_jobs= -1)
	elif ctype == "lr":
		model = LR(class_weight= "balanced")
	else:
		model = rf(n_estimators= 1000, n_jobs= -1)

	cc = chain(model).fit(Xtrain, ytrain)
	ypred = cc.predict(Xtest)
	save_model(mlb, ss, cc, "cc_"+str(top))
	acc = accuracy_score(ytest, ypred)
	hloss = hamming_loss(ytest, ypred)
	print("Accuracy score Classifier Chains: ", acc)
	print("Hamming loss: ", hloss)
	mprec, mrecall, mf1 = micro_avg(ytest, ypred)
	#per_class_dist(ytest, ypred, labels)
	#macro_avg(ytest, ypred)
	return [acc, hloss, mprec, mrecall, mf1, ctype, model]

def labelpowerset(ss, Xtrain, ytrain, Xtest, ytest, mlb, labels, top, ctype="rf"):
	print("Classifier: ", ctype)
	if ctype == "knn":
		model = knnbase(n_neighbors=3, n_jobs= -1)
	elif ctype == "lr":
		model = LR(class_weight= "balanced")
	else:
		model = rf(n_estimators= 1000, n_jobs= -1)

	ps = power(model).fit(Xtrain, ytrain)
	ypred = ps.predict(Xtest)
	save_model(mlb, ss, ps, "lp_"+str(top))
	acc = accuracy_score(ytest, ypred)
	hloss = hamming_loss(ytest, ypred)
	print("Accuracy score LabelPowerset: ", acc)
	print("Hamming loss: ", hloss)
	mprec, mrecall, mf1 = micro_avg(ytest, ypred)
	#per_class_dist(ytest, ypred, labels)
	#macro_avg(ytest, ypred)
	return [acc, hloss, mprec, mrecall, mf1, ctype, model]


# Stats & Graph: Label distribution
def label_stats(sha_label_map):
	#print(sha_label_map)
	freq_labels = dict()
	taglens = []
	avglen = 0
	for mali, taglst in sha_label_map.items():
		print(mali, taglst)
		tags = taglst[0]
		taglen = len(tags)
		taglens += [taglen]
		avglen += taglen

		for tag in tags:
			if tag not in freq_labels:
				freq_labels[tag] = 1
			else:
				freq_labels[tag] += 1

	print("Label distribution: ", freq_labels)
	# Stats: min tag per mal, max tag per mal, avg no. of tags per mal
	taglens.sort()
	print("Min tag length: ", taglens[0])
	print("Max tag length: ", taglens[-1])
	print("Avg no. of tags/mal: ", avglen/len(sha_label_map), len(sha_label_map))

	return

def save_model(mlb, scaler, model, naming, dir="output/multi/"):
	curwd = os.getcwd()
	if not os.path.exists(dir):
		os.system("mkdir "+curwd+"/output")
		os.system("mkdir "+curwd+"/"+dir)
	dump(model, open(dir+naming+'_model.pkl', 'wb'))
	# save the scaler
	dump(scaler, open(dir+naming+'_scaler.pkl', 'wb'))
	dump(mlb, open(dir+naming+'_binarizer.pkl', 'wb'))
	return


# ------------------------------- MODEL TESTING ------------------------------------- #
# Xtest: dataframe with features and malware label
def test_model(Xtest_label, mname, ztestlabels, mpath="output/multi/"):
	mfullpath = mpath+mname+"_model.pkl"
	spath = mpath+mname+"_scaler.pkl"
	binpath = mpath+mname+"_binarizer.pkl"
	if not os.path.exists(mfullpath):
		print("Error! Pre trained multilabel model not found on path: 'output/multi/'. Train model using 'classify_topk.py' using options-D5/-D5_host, multiclass mode:1, --train")
		return
	# load the model
	model = load(open(mfullpath, 'rb'))
	# load the scaler
	scaler = load(open(spath, 'rb'))
	# load label binarizer
	mlb = load(open(binpath, 'rb'))
	# Drop labels 0/1 and add ztest as labels
	newtestdf = Xtest_label.iloc[:,:-1].copy()
	truelabels = mlb.transform(ztestlabels).tolist()[0]
	print("All class labels (trained on): ", list(mlb.classes_))
	print("ACTUAL labels for zeroday binary: ", ztestlabels, " encoded: ",truelabels)
	time.sleep(2)
	labellst = [truelabels]*42
	Xtest = Xtest_label.iloc[:,:-1].copy()
	labels = np.array(labellst)
	X_test_scaled = scaler.transform(Xtest)
	# Make predictions on the test df
	ypred = model.predict(X_test_scaled)
	mlb = load(open(binpath, 'rb'))
	pred_tags = mlb.inverse_transform(ypred)
	print("\nPredicted tags (per test instance): ", pred_tags)
	#print("For following encoded labels: ", labels)
	hloss = hamming_loss(labels, ypred)
	print("\nHloss for test: ", hloss, mname)
	micro_avg(labels, ypred)
	time.sleep(2)
	return

def test_models(Xtest, topk, ztestlabels = [["grayware", "worm", "ransomware", "downloader"]]):
	if topk == 3:
		mnames = ["br_3", "cc_3", "lp_3"] # models trained using topk=3
	else:
		mnames = ["br_1", "cc_1", "lp_1"] # models trained using topk=1

	for mname in mnames:
		if "br" in mname:
			print("Testing Multilabel technique: Binary Relevance")
		elif "cc" in mname:
			print("Testing Multilabel technique: Classifier Chains")
		else:
			print("Testing Multilabel technique: Label Powerset")
		test_model(Xtest, mname, ztestlabels)
	return