# =============================================================================
# Data preprocessing
# Load Libraries
# =============================================================================

import os
from nltk import ngrams
import csv

# =============================================================================
# Declare variables
# =============================================================================

# Keep all the asm files under asm directory

path = "asms/"
outputPath = "csvs/"
EXT = ".csv"
opcodeDict = ["MOV","XCHG","STC","CLC","CMC","STD","CLD","STI","CLI","PUSH","PUSHF",
              "PUSHA","POP","POPF","POPA","CBW","CWD","CWDE","IN","OUT","ADD","ADC","SUB",
              "SBB","DIV","IDIV","MUL","IMUL","INC","DEC","CMP","SAL","SAR","RCL","RCR","ROL",
              "ROR","NEG","NOT","AND","OR","XOR","SHL","SHR","NOP","LEA","INT","CALL","JMP","JE",
              "JZ","JCXZ","JP","JPE","RET","JNE","JNZ","JECXZ","JNP","JPO","JA","JAE","JB","JBE","JNA",
              "JNAE","JNB","JNBE","JC","JNC","JG","JGE","JL","JLE","JNG","JNGE","JNL","JNLE","JO","JNO","JS","JN"]
sampleOpcodes = ["MOV","XCHG","STC","CLC","CMC","STD","CLD","STI","CLI","PUSH","PUSHF"]

# =============================================================================
# Clean the file, remove unwanted code
# =============================================================================

def cleanFile(path, filename):
    opcodeSequence = []
    opcodeSequenceIndex = []
    with open(path+filename) as f:
        for line in f:
            w = line.split(' ')
            opcode = w[0].upper()
            if opcode in opcodeDict:
                opcodeSequence.append(opcode) 
                index = opcodeDict.index(opcode)
                opcodeSequenceIndex.append(index)
    f.close()
    return opcodeSequence, opcodeSequenceIndex

# =============================================================================
# Create N grams for String based data - PUSH,ADD and Integer based Data PUSH = 10, etc
# =============================================================================

def createNgrams(n,text,textNumbers):
    stringNgrams = ""   # Opcodes replaced with numeric value
    stringNgrams_ = "" 
    ngram = ngrams(text, n)
    ngram_ = ngrams(textNumbers, n)
    print("Number of grams : ",n)
    for grams in ngram:
        #print(grams)
        stringNgrams = stringNgrams + ','.join(str(s) for s in grams) + ","
    for grams in ngram_:
        #print(grams)
        stringNgrams_ = stringNgrams_ + ','.join(str(s) for s in grams) + ","
        
    sequence = stringNgrams.rstrip(',')
    sequenceNumbered_ = stringNgrams_.rstrip(',')    
    return sequence,sequenceNumbered_

def generateCSV(filename, OpcodeStr):
    file = open(filename,"w") 
    file.write(OpcodeStr) 
    file.close()
    print("Generated csv file ", filename)

if __name__=="__main__":
    files = os.listdir(path)
    if '.DS_Store' in files:
        files.remove('.DS_Store')
    print(files)
    for index in range(len(files)):
        opcodes,opcodesNumered  = cleanFile(path, files[index])
        ngramsOrig, ngramsInt = createNgrams(1,opcodes,opcodesNumered)
        if not os.path.exists(outputPath):
            os.makedirs(outputPath)
        generateCSV(outputPath+"sampleStr"+str(index)+EXT,ngramsOrig)
        generateCSV(outputPath+"sampleInt"+str(index)+EXT,ngramsInt)
