#coding:utf-8
from collections import defaultdict
import csv
import pickle
from collections import OrderedDict
#"PEInfo5_headers.csv" "PEInfo5_apis.csv" "PEInfo5_sections.csv"

def get_unigrams(filepath):
    binaryOfeachFile_dic = defaultdict(list)
    # unigramsOfeachFile_dic = defaultdict(list)
    unigramsOfeachFile_dic=OrderedDict()
    unigrams_freq = OrderedDict()
    #获取每个文件的unigrams

    csv_reader=csv.reader(open(filepath))
    for row in csv_reader:
        if row[0]:
            file_name = row[0]
        for col in row[2:]:
            if file_name not in unigramsOfeachFile_dic.keys():
                unigramsOfeachFile_dic[file_name] = []
            if col not in unigramsOfeachFile_dic[file_name] and col!= '' and col:
                col=str(col)
                if col.isupper():
                    col='.'+col.lower()
                unigramsOfeachFile_dic[file_name].append(col)
    # keys = [x[0] for x in unigramsOfeachFile_dic.items() if '.text' not in x[1]]
    # print(keys)
    print("文件数量：",len(unigramsOfeachFile_dic.keys()))

    #遍历字典统计出现频率

    for key in unigramsOfeachFile_dic:
        for item in unigramsOfeachFile_dic[key]:
            if item not in unigrams_freq.keys():
                unigrams_freq[item]=1
            else:
                unigrams_freq[item]+=1
    print("所有文件的unigram个数：",len(unigrams_freq.keys()))
    print("所有文件的unigram:",unigrams_freq.keys())
    time1=0
    for value in unigrams_freq.values():
        if value<4:
            time1=time1+1
    print("出现次数较小的unigram数",time1)
    print("总unigram出现次数分布：",sorted(unigrams_freq.values(),reverse=True))

    #输出每个文件的二值分布

    unigramsByClear=unigrams_freq.keys()
    print("整理后的unigram：",unigramsByClear)
    # for key in unigramsOfeachFile_dic:
    #     for value in unigramsByClear:
    #         if value not in unigramsOfeachFile_dic[key]:
    #             binaryOfeachFile_dic[key].append(0)
    #         else:
    #             binaryOfeachFile_dic[key].append(1)
    #
    # print(binaryOfeachFile_dic.keys())
    # print("二值分布为：",binaryOfeachFile_dic["000000000000634_Virus.Win32.Etap.exe"])
    #将二值写入文件
    # write_csv('benign-apis',binaryOfeachFile_dic)
    return unigramsOfeachFile_dic,list(unigramsByClear)

def get_binary_api(unigramsOfeachFile_dic,unigramsByClear):
    binaryOfeachFile_dic = defaultdict(list)
    for key in unigramsOfeachFile_dic:
        for value in unigramsByClear:
            if value not in unigramsOfeachFile_dic[key]:
                binaryOfeachFile_dic[key].append(0)
            else:
                binaryOfeachFile_dic[key].append(1)
    return binaryOfeachFile_dic.values()
    # write_csv('malware-apis',binaryOfeachFile_dic)
def get_binary_headers(unigram_dic,attri_dic,dict1):
    binary_dic = defaultdict(list)
    for key in unigram_dic.keys():
        for i in range(len(unigram_dic[key])):
            for j in range(len(attri_dic[dict1[i]])):
                if attri_dic[dict1[i]][j]==unigram_dic[key][i]:
                    binary_dic[key].append(1)
                else:
                    binary_dic[key].append(0)
    return binary_dic.values()
    # write_csv('benign-headers', binary_dic)
def write_csv(filename,dic):
    csvFile = open(filename+'.csv', 'w', newline='')
    writer = csv.writer(csvFile)
    for key in dic.keys():
        dic[key].insert(0,key)
        writer.writerow(dic[key])
    csvFile.close()
def get_headers(filepath,type):
    dict1={} #保存属性列名
    unigram_dic=defaultdict(list)
    binary_dic=defaultdict(list)
    attri_dic = defaultdict(list)
    with open(filepath) as fin:
        line=fin.readline()
        tokens=line.strip().split(',')
        if type=='malware':
            for i in range(len(tokens)):
                tokens[i]=tokens[i][7:]
        attri_dic.fromkeys(tokens[1:])
        for i in range(len(tokens[1:])):
            dict1[i]=tokens[i+1]
        for row in fin.readlines():
            row=row.strip().split(',')
            file_name=row[0]
            unigram_dic[file_name]=row[1:]
            row=row[1:]
            for i in range(len(row)):
                if row[i] not in attri_dic[dict1[i]] and row[i]!='' and row[i]:
                    attri_dic[dict1[i]].append(row[i])
    fin.close()
    list1=[]
    for i in range(len(dict1.keys())):
        list1+=attri_dic[dict1[i]]
    for key in unigram_dic.keys():
        for i in range(len(unigram_dic[key])):
            for j in range(len(attri_dic[dict1[i]])):
                if attri_dic[dict1[i]][j]==unigram_dic[key][i]:
                    binary_dic[key].append(1)
                else:
                    binary_dic[key].append(0)
    # write_csv('malware-headers',binary_dic)
    return unigram_dic,attri_dic,dict1

    # col=[row[2] for row in csv_reader]



#获取训练样本header信息并写入文件
# unigram_dic,attri_dic,dict1=get_headers('./mixed_samples/malware data/malware-PEInfo5_headers.csv','malware')
# unigram_dic1,attri_dic1,dict11=get_headers('./mixed_samples/benign data/benign-PEInfo5_headers.csv','benign')
# for key in attri_dic.keys():
#     for data in attri_dic[key]:
#         if data not in attri_dic1[key]:
#             attri_dic1[key].append(data)
# f=open('headerOfall','wb')
# f2=open('headerColName','wb')
# pickle.dump(attri_dic1,f,True)
# pickle.dump(dict11,f2,True)
# f.close()
# f2.close()
# print(unigram_dic1)
# get_binary_headers(unigram_dic1,attri_dic1,dict11)

#获取训练样本api信息
# unigramsOfeachFile_dic,unigramsByClear=get_unigrams('./mixed_samples/malware data/malware-PEInfo5_apis.csv')
# unigramsOfeachFile_dic1,unigramsByClear1=get_unigrams('./mixed_samples/benign data/benign-PEInfo5_apis.csv')
#
# for data in unigramsByClear:
#     if data not in unigramsByClear1:
#         unigramsByClear1.append(data)
# get_binary_api(unigramsOfeachFile_dic,unigramsByClear1)
# #保存合并的unigrams
# print(unigramsByClear1)
# f=open('apiOfall','wb')
# pickle.dump(unigramsByClear1,f,True)
# f.close()



#测试生成单个文件的api的0,1序列
# unigramsOfeachFile_dic,unigramsByClear=get_unigrams('./mixed_samples/malware data/malware-PEInfo5_apis.csv')
# f=open('apiOfall','rb')
# apiOfall=pickle.load(f)
# print(unigramsOfeachFile_dic['1.exe'])
# f.close()
# get_binary_api(unigramsOfeachFile_dic,apiOfall)

#测试生成单个文件的header的特征序列
# unigram_dic1,attri_dic1,dict11=get_headers('2.csv','benign')
# f=open('headerOfall','rb')
# headerOfall=pickle.load(f)
# print(headerOfall.keys())
# f.close()
# get_binary_headers(unigram_dic1,headerOfall,dict11)



# for data in get_unigrams('./mixed_samples/malware data/malware-PEInfo5_apis.csv'):
#     if data not in api_unigram:
#         api_unigram.append(data)
# print(len(api_unigram))
# api_dic=get_unigrams('./mixed_samples/benign data/benign-PEInfo5_apis.csv')

# header_dic=get_unigrams('./mixed_samples/benign data/benign-PEInfo5_sections.csv').keys()
# print(api_dic)
# print(header_dic)
# print("api length",len(api_dic))
# print("headers length:",len(header_dic))


#求并集
# ret_list1 = list(set(api_dic).union(set(header_dic)))

#求交集
# ret_list = list((set(api_dic).union(set(header_dic)))^(set(api_dic)^set(header_dic)))

#求差集
# print(api_dic-header_dic)
# print(header_dic-api_dic)
# section_dic=get_unigrams('./mixed_samples/benign data/PEInfo5_sections.csv').keys()
# print(header_dic-api_dic)
# print([item for item in header_dic if item not in api_dic])
# print(len(get_unigrams('./mixed_samples/benign data/PEInfo5_apis.csv').keys()))
