import csv

from numpy import genfromtxt
from sklearn.decomposition import PCA, NMF
from sklearn.metrics import silhouette_score


# Read all filenames from the files.csv
def get_file_names():
    with open('files.csv', 'r') as files_csv:
        reader = csv.reader(files_csv)
        return list(reader)[0]


# Utility method for converting hex strings to decimal numbers
def hex_to_dec(hexString):
    return int(hexString, base=16)


# Use numpy gen from text to read csv in one operation into numerical numpy array
def read_csv_to_numpy_array():
    data = genfromtxt('features.csv', delimiter=',')
    return data


# Perform Non-Negative Matrix Factorisation on the dataset with the specified number of components or features
def do_nmf(data, components):
    nmf = NMF(n_components=components)
    nmf.fit(data[1:])
    dataSet = nmf.transform(data[1:])
    return dataSet


# Perform Principle component Analysis on the dataset with the specified number of components or features
def do_pca(data, components):
    pca = PCA(n_components=components)
    pca.fit(data[1:])
    dataSet = pca.transform(data[1:])
    return dataSet


# Calculate the silhouette scores for the clustered dataset
def get_silhouette_score(cluster_labels, dataSet):
    return silhouette_score(dataSet, cluster_labels)

# Calculate how many files in the dataset were correctly ideentified as malicious or benign
def calculate_detection_accuracy(labels, files):
    benLabel = 0

    correct = 0

    for i in range(len(files)):
        if(files[i][:3]=="BEN"):
            benLabel=labels[i]
            break

    for i in range(len(files)):
        # If file is not benign and was not labelled as benign
        if (str(files[i])[:3] != "BEN" and labels[i] != benLabel):
            correct += 1
        # if file is benign and was labbeled as benign
        elif (str(files[i])[:3] == "BEN" and labels[i] == benLabel):
            correct += 1

    return correct / len(files) * 100

# Calculate how many files in the dataset were correctly clustered
def calculate_cluster_accuracy(labels, files):
    trojLabel = -1
    ransomLabel = -1
    adLabel = -1
    virusLabel = -1
    backDoorLabel = -1
    wormLabel = -1
    spyLabel = -1
    benLabel = -1
    malLabel = -1
    correct = -1

    # For each file type, find the the first unique label
    for i in range(len(files)):
        if (str(files[i])[:3].upper() == "TRO" and trojLabel == -1):
            if(labels[i] not in [virusLabel, ransomLabel, adLabel, backDoorLabel, wormLabel, spyLabel, benLabel, malLabel]):
                trojLabel = labels[i]
        if (str(files[i])[:3].upper() == "VIR" and virusLabel == -1):
            if(labels[i] not in [trojLabel, ransomLabel, adLabel, backDoorLabel, wormLabel, spyLabel, benLabel, malLabel]):
                virusLabel = labels[i]
        if (str(files[i])[:3].upper() == "BAC" and backDoorLabel == -1):
            if(labels[i] not in [trojLabel, ransomLabel, adLabel, virusLabel, wormLabel, spyLabel, benLabel, malLabel]):
                backDoorLabel = labels[i]
        if (str(files[i])[:3].upper() == "WOR" and wormLabel == -1):
            if(labels[i] not in [trojLabel, ransomLabel, adLabel, virusLabel, backDoorLabel, spyLabel, benLabel, malLabel]):
                wormLabel = labels[i]
        if (str(files[i])[:3].upper() == "ADW" and adLabel == -1):
            if(labels[i] not in [trojLabel, ransomLabel, backDoorLabel, virusLabel, wormLabel, spyLabel, benLabel, malLabel]):
                adLabel = labels[i]
        if (str(files[i])[:3].upper() == "SPY" and spyLabel == -1):
            if(labels[i] not in [trojLabel, ransomLabel, adLabel, virusLabel, wormLabel, backDoorLabel, benLabel, malLabel]):
                spyLabel = labels[i]
        if (str(files[i])[:3].upper() == "RAN" and ransomLabel == -1):
            if(labels[i] not in [trojLabel, backDoorLabel, adLabel, virusLabel, wormLabel, spyLabel, benLabel, malLabel]):
                ransomLabel = labels[i]
        if (str(files[i])[:3].upper() == "BEN" and benLabel == -1):
            if(labels[i] not in [trojLabel, ransomLabel, adLabel, virusLabel, wormLabel, spyLabel, backDoorLabel, malLabel]):
                benLabel = labels[i]
        if (str(files[i])[:3].upper() == "MAL" and malLabel == -1):
            if(labels[i] not in [trojLabel, ransomLabel, adLabel, virusLabel, wormLabel, spyLabel, benLabel, backDoorLabel]):
                malLabel = labels[i]

    # For each file, check if the label matches the expected label, if it does increment the correct count
    for i in range(len(files)):
        if (str(files[i])[:3].upper()  == "TRO" and labels[i] == trojLabel):
            correct += 1
        elif (str(files[i])[:3].upper()  == "VIR" and labels[i] == virusLabel):
            correct += 1
        elif (str(files[i])[:3].upper()  == "BAC" and labels[i] == backDoorLabel):
            correct += 1
        elif (str(files[i])[:3].upper()  == "WOR" and labels[i] == wormLabel):
            correct += 1
        elif (str(files[i])[:3].upper()  == "ADW" and labels[i] == adLabel):
            correct += 1
        elif (str(files[i])[:3].upper()  == "SPY" and labels[i] == spyLabel):
            correct += 1
        elif (str(files[i])[:3].upper()  == "RAN" and labels[i] == ransomLabel):
            correct += 1
        elif (str(files[i])[:3].upper()  == "BEN" and labels[i] == benLabel):
            correct += 1
        elif (str(files[i])[:3].upper()  == "MAL" and labels[i] == malLabel):
            correct += 1

    # Return percentage of dataset correctly clustered
    return correct / len(files) * 100


# Extract our chosen features to an array which will be written to a csv later
def perform_feature_extraction(pe_file):
    section_entropy = [0] * pe_file.get_number_of_sections()
    # Calculate entropy values for each section
    for index in range(pe_file.get_number_of_sections()):
        offset = pe_file.get_section_pointer_to_raw_data(index)
        length = pe_file.get_section_size_of_raw_data(index)
        entropy = pe_file.calculate_entropy(hex_to_dec(offset), hex_to_dec(length))
        section_entropy[index] = entropy

    file_entropy = pe_file.get_file_entropy()
    code_data_ratio = pe_file.get_code_data_ratio()

#    if(code_data_ratio != 0):
#        code_data_ratio = abs(math.log(code_data_ratio)) * 8
    image_version = hex_to_dec(pe_file.get_major_image_version())
#                            if(image_version != 0):
#                                image_version = abs(math.log(image_version)) * 8

    # Initialise empty array with 18 elements as we have 18 features
    row = [0] * 18

    # Write the entropy values of the first 4 sections to the array
    for i in range(4):
        if i >= pe_file.get_number_of_sections():
            row[i] = 0
        else:
            row[i] = section_entropy[i]

    row[4] = file_entropy
    row[5] = code_data_ratio
    row[6] = image_version
    row[7] = pe_file.get_number_of_sections()

    imports = pe_file.get_imports()

    # Set flags for presence of imports
    for _import in imports:
        if _import == "KERNEL32.DLL":
            row[8] = 1
        if _import == "USER32.DLL":
            row[9] = 1
        if _import == "ADVAPI32.dll":
            row[10] = 1
        if _import == "msvcrt.dll":
            row[11] = 1
        if _import == "GDI32.dll":
            row[12] = 1
        if _import == "SHELL32.dll":
            row[13] = 1
        if _import == "ole32.dll":
            row[14] = 1
        if _import == "WS2_32.dll":
            row[15] = 1
        if _import == "SHLWAPI.dll":
            row[16] = 1
        if _import == "COMCTL32.dll":
            row[17] = 1
    return row