import os
import sys
from macholib.MachO import MachO

from scipy import interp
import numpy
import matplotlib.pyplot as plt
import pandas as pd
from pandas.tools.plotting import scatter_matrix

from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import AdaBoostClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.cross_validation import StratifiedKFold
from sklearn.metrics import roc_curve, auc
import matplotlib
matplotlib.style.use('ggplot')


names = ["Nearest Neighbors", "Linear SVM", "Decision Tree", "AdaBoost", "Naive Bayes"]
classifiers = [
    KNeighborsClassifier(3),
    SVC(kernel="linear", C=0.025, probability=True),
    DecisionTreeClassifier(max_depth=5),
    AdaBoostClassifier(),
    GaussianNB()]
colors = ["b", "g", "r", "c", "m"]

def getSectionCount(macho):
    sectionCount = 0
    for load_cmd, cmd, data in macho.headers[0].commands:
        if hasattr(cmd, "segname"):
            sectionCount += len(data)
    f_sections = sectionCount
    return f_sections


def getSizeInKiloByte(fullpath):
    size = os.path.getsize(fullpath)
    f_size = size / 1024.0
    return f_size

def getData(pathToSamples, c):
    df = pd.DataFrame(columns=('f_size', 'f_sections', 'f_flags', 'c'))
    currentRow = 0

    for f in os.listdir(pathToSamples):
        if not ".json" in f:
            fullpath = os.path.join(pathToSamples, f)
            f_size = getSizeInKiloByte(fullpath)

            try:
                macho = MachO(fullpath)
                f_sections = getSectionCount(macho)
                f_flags = macho.headers[0].header.flags
            except:
                continue

            df.loc[currentRow] = [f_size, f_sections, f_flags, c]
            currentRow += 1

    return df

def prepareDatasetsForLearning(dataSet):
    X = []
    y = []

    for index, row in dataSet.iterrows():
        X.append([row.f_size, row.f_sections, row.f_flags])
        y.append(row.c)

    return X,y

def classify(X, y):
    # TODO: split learning and plotting
    X = numpy.array(X)
    y = numpy.array(y)

    X = StandardScaler().fit_transform(X)
    skf = StratifiedKFold(y, n_folds=10)

    for classifier in range(0,len(classifiers)):
        print "Training classifier %s..." % (names[classifier])

        mean_tpr = 0.0
        mean_fpr = numpy.linspace(0, 1, 100)

        for train_index, test_index in skf:
            X_train, X_test = X[train_index], X[test_index]
            y_train, y_test = y[train_index], y[test_index]
            clf = classifiers[classifier]
            probas_ = clf.fit(X_train, y_train).predict_proba(X_test)
            fpr, tpr, thresholds = roc_curve(y_test, probas_[:, 1])
            mean_tpr += interp(mean_fpr, fpr, tpr)
            mean_tpr[0] = 0.0
        mean_tpr /= len(skf)
        mean_tpr[-1] = 1.0
        mean_auc = auc(mean_fpr, mean_tpr)
        plt.plot(mean_fpr, mean_tpr, 'k--',
             label= names[classifier] + ' (AUC = %0.3f)' % mean_auc, lw=2, c=colors[classifier])

    plt.plot([0, 1], [0, 1], '--', color=(0.6, 0.6, 0.6), label='Luck')

    plt.xlim([-0.05, 1.05])
    plt.ylim([-0.05, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.legend(loc="lower right", fancybox=True, shadow=True, fontsize=10)
    plt.tight_layout()
    plt.savefig("osx_malware_goodware_roc.svg")

def plotDataset(dataSet):
    colors=['green','red']
    fig, ax = plt.subplots()
    scatter_matrix(dataSet.ix[:,:'f_flags'], alpha=0.8, figsize=(6, 6), ax=ax, marker='x', diagonal="kde", c=dataSet.c.apply(lambda x:colors[int(x)]))
    plt.tight_layout()
    fig.savefig("scatter.svg")

def main(args):
    if len(args) != 3:
        print "Usage: cluster_macho.py PATH_TO_MALWARE PATH_TO_GOODWARE"
    else:
        malicious = getData(args[1], 1)
        print "Read in %i malicious samples." % len(malicious)
        print malicious.info()
        print malicious.describe()

        benign = getData(args[2], 0)
        print "Read in %i benign samples." % len(benign)
        print benign.info()
        print benign.describe()

        dataSet = pd.concat([malicious, benign])
        X,y = prepareDatasetsForLearning(dataSet)

        classify(X, y)

        plotDataset(dataSet)

if __name__ == "__main__":
    main(sys.argv)