# -*- coding: utf-8 -*-
## instructions frequency

from multiprocessing import Pool
import os
import csv
import re
#全局参数配置
labelNum = 3
fileBasePath = 'E:\\kiggle\\kiggleData\\'
trainLabelFilePath = r'E:\kiggle\kiggleData\\trainLabels.csv'
paths = ['train','test']

#指令集
instr_set = set(['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','ja','jae','jb','jbe','jna',
	'jnae','jnb','jnbe','jc','jnc','ret','jne','jnz','jecxz',
	'jnp','jpo','jg','jge','jl','jle','jng','jnge','jnl','jnle',
	'jo','jno','js','jns'])


def consolidate(fpath,instr_set = instr_set):
    global fileBasePath
    path = fileBasePath + fpath
    Files = os.listdir(path)
    asmFiles = [i for i in Files if '.asm' in i]
    consolidatedFile = 'featfile\\'+ fpath + '_instrucion_freq.csv' #windows文件路径命名方式
    import time
    start = time.time()
    with open(consolidatedFile, 'w',newline='') as f:
        fieldnames = ['Id'] + list(instr_set)
        writer = csv.DictWriter(f, fieldnames=fieldnames)  # write header
        writer.writeheader()
        opcode_seq = []
        for t, fname in enumerate(asmFiles):
            consolidation = dict(zip(instr_set, [0] * len(instr_set)))
            consolidation['Id'] = fname[:fname.find('.asm')]
            with open(path + '\\' + fname, 'r',errors='ignore') as f: #windows文件路径命名方式
                for line in f:
                    if 'text' in line:
                        if ';'  in line and 'call' not in line: #针对asm文件修改+
                            continue
                        row = line.lower().strip().split('  ')[1:]
                        if row:
                            tmp_list = [x.strip() for x in row if x != '']
                            if tmp_list[0] == 'call' or (len(tmp_list) == 2 and tmp_list[0] in consolidation):
                                consolidation[tmp_list[0]] += 1

                writer.writerow(consolidation)
    end = time.time()
    print('inst freq time up per:',(end-start)/len(asmFiles))
        # if (t+1)%100 == 0:
        #	print str(t+1) + 'files loaded for ' + path

if __name__ == '__main__':
    if not os.path.exists('featfile'):
        os.makedirs('featfile')
    # test function
    #consolidate('kaggle_dataSample')
    p = Pool(2)
    p.map(consolidate, paths)
    print("DONE instruction count!")



