# -*- coding: utf-8 -*-
"""
PROJECT: Fighting malware using machine learning.

This file contains all the functions needed get a list of opcodes from the 
files provided in the data set from the Microsoft malware competition and 
save these opcode lists to a pkl file for later use in training.
@author: Jacques-Antoine Portal
"""
import os
import pickle
import progressbar
from pyparsing import Word, hexnums, WordEnd, alphas, Optional, alphanums



def readAsmFile(inputFolder = 'train/', outputFolder = 'OpcodeList/', opcodeDic = {}):
    '''
    this functions reads the .asm files from the input folder and dumps 
    an array of all the opcodes found in that file as a pkl file in the 
    output folder.
    '''
    bar = progressbar.ProgressBar()
    emptyFiles = []
    for filename in bar(os.listdir(inputFolder)):
        if filename.endswith(".asm"): 
            asmFile = open(os.path.join(inputFolder, filename), "r", encoding='ISO-8859-1')
            source = asmFile.readlines()
            asmFile.close()
            fileOpcodes, opcodeDic = parseAsm(source, opcodeDic)
            if len(fileOpcodes) < 100:
                emptyFiles += [filename]
            output = open(os.path.join(outputFolder, (filename + ".pkl")), 'wb')
            pickle.dump(fileOpcodes, output, -1)
            output.close()
    return emptyFiles, opcodeDic

#these are constants used to parse the strings found in the file.
hex_integer = Word(hexnums) + WordEnd()
alphanumsPlusPoint = alphanums + ".:"
hex_int_plus = Word(hexnums) + "+" + WordEnd() # use WordEnd to avoid parsing leading a-f of non-hex numbers as a hex
LINE = (Word(alphanumsPlusPoint) + (hex_integer*(1,))("instructions") + Optional(hex_int_plus) + Word(alphas)("opcode"))


def isItAnOpcode(op):
    '''
    This fucntion returns a boolean if the string provided as an argument qualifies
    as an opcode (no capital letter, not data (dd, db, dw)) and size between 2 and 5.
    '''
    return (op == op.lower() and op not in "dd db dw" and len(op) > 1 and len(op) < 6)


def parseAsm(source, opcodeDic, line = LINE):
    '''
    This parsing function goes through the source code of an .asm file and 
    returns an ordered list of all the opcodes it found.
    '''
    fileOpcodes = []
    for source_line in source:
        try:
            result = line.parseString(source_line)
            if "opcode" in result:
                keys = opcodeDic.keys()
                op = result.opcode
                if isItAnOpcode(op):
                    if op not in keys:
                        opcodeDic[op] = len(keys)
                    fileOpcodes += [opcodeDic[op]]#, result.instructions.asList(zyydirectory = os.fsencode(directory_in_str)
        except:
            continue
    return fileOpcodes, opcodeDic

def getOpcodeList(fileName, opcodeDic, line = LINE):
    '''
    This function extracts the opcodes from one file and returns the list of 
    opocdes as an integer list as well as the opcode dictionary.
    '''
    asmFile = open(fileName, "r", encoding='ISO-8859-1')
    source = asmFile.readlines()
    asmFile.close()
    fileOpcodes, opcodeDic = parseAsm(source, opcodeDic)
    return fileOpcodes, opcodeDic

def getIntegerListOfOpcode(inputFolder = 'OpcodeList/', outputFolder = 'IntegerList/'):
    '''
    This function transforms lists of opcodes found in the input folder and 
    outputs files containing lists of integers in the output folder.
    It also creates a dirrectory of all opcodes and returns it for later use
    in opcode analysis.
    '''
    bar = progressbar.ProgressBar()
    OPCODES = {}
    i = 0
    for filename in bar(os.listdir(inputFolder)):
        file = open(os.path.join(inputFolder, filename), 'rb')
        opcodes = pickle.load(file)
        file.close()
        integers = []
        for op in opcodes:
            if op not in OPCODES.keys():
                OPCODES[op] = i
                i+=1
            integers += [OPCODES[op]]
        output = open(os.path.join(outputFolder, (filename)), 'wb')
        pickle.dump(integers, output, -1)
        output.close()
    output = open('OPCODES.pkl', 'wb')
    pickle.dump(OPCODES, output, -1)
    output.close()
    return OPCODES
        

#This is an example of how to use these functions
#OPCODES = getIntegerListOfOpcode()
#asmFile = open(os.path.join("train/", "4kuxDJlwWHMi5KX3SVQq.asm"), "r", encoding='ISO-8859-1')
#source = asmFile.readlines()
#fileOpcodes = parseAsm(source)
#asmFile.close()
#print(fileOpcodes)
#emptyFiles, OPCODES = readAsmFile()

#print(emptyFiles)
#print("DONE")