import re
import os
import csv
dirname = '/home/malware/trace2_test'
outdir = './data'

def filterDomainSet(domains):
    result = set()
    for domain in domains:
        if domain.startswith('www.'):
            result.add(domain[4:])
        else:
            result.add(domain)
    return result

def loadWhiteMap():
    whiteSet = set()
    secondDomainSet = set()
    for line in open('prank.top.1m.20180322').readlines():
        domain = line.split('\t')[0]
        whiteSet.add(domain)
    filtered = filterDomainSet(whiteSet)
    for domain in filtered:
        vals = domain.split('.')
        secondDomain = '.'.join(vals[1:])
        secondDomainSet.add(secondDomain)
    return whiteSet

fileMap = {
    'fqdn':os.path.join(dirname,'fqdn_ip.csv'),
    'email':os.path.join(dirname,'trojan_email.csv'),
    'phone':os.path.join(dirname,'trojan_phone.csv'),
    'web':os.path.join(dirname,'trojan_web.csv'),
    'region':os.path.join(dirname,'fqdn_reginfo.csv'),
}
outMap = {
    'all_domain':os.path.join(outdir,'domains.txt'),
    'registered_domain':os.path.join(outdir,'registered.txt'),
    'domainFeature':os.path.join(outdir,'static_domain.txt'),
    'dga2level':os.path.join(outdir,'dga2level.txt'),
    'graph':os.path.join(outdir,'graph.txt'),
    'blackDomain':os.path.join(outdir,'blackdomain.txt'),
    'whiteDomain':os.path.join(outdir,'whiteDomain.txt'),



}



def readDomainSet():
    sha2domain = {}
    domain2sha = {}
    for key in ['web','fqdn','region']:
        filename = fileMap[key]
        inf = open(filename,encoding='utf-8')
        reader = csv.reader(inf)
        headings = next(reader)
        for data in reader:
            sha = data[0]
            if sha not in sha2domain:
                sha2domain[sha] = set()
            url = data[1]
            if '/' in url:
                vals = url.split('/')
                fulldomain = vals[2].split(':')[0]
            else:
                fulldomain = url
            if fulldomain.startswith('www.'):
                fulldomain = fulldomain[4:]
            if not checkip(fulldomain):
                if fulldomain not in domain2sha:
                    domain2sha[fulldomain] = set()
                sha2domain[sha].add(fulldomain)
                domain2sha[fulldomain].add(sha)
    return sha2domain,domain2sha


def domain2key(domain):
    vals = domain.split('.')
    for val in vals[::-1]:
        if val not in  {'edu','cn','com','net'}:
            return val
def checkip(ip):  
    p = re.compile('^((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)$')  
    if p.match(ip):  
        return True  
    else:  
        return False  

def oneHotFeature(dataList,featureName):
    keys = {}
    for data in dataList:
        key = featureName+data[featureName]
        if key not in keys:
            keys[key] = 0
    for data in dataList:
        data.update(keys)
        key = featureName+data[featureName]
        data[key] = 1
        del(data[featureName])

def clusterSet(vecrot2Group,node2vector):
    def rootInGroupTree(groupTree,node):
        if node not in groupTree:
            groupTree[node]=node
        while groupTree[node]!=node:
            node=groupTree[node]
        return node
    groupTree = {}
    for node,vectors in node2vector.items():
        group = None
        vectors = list(vectors)
        for vector in vectors:
            if vector not in vecrot2Group:
                continue
            nodeGroup = rootInGroupTree(groupTree,vecrot2Group[vector])
            if group is None:
                group = nodeGroup
            else:
                groupTree[rootInGroupTree(groupTree,vecrot2Group[vector])] = group
    for node  in groupTree:
        groupTree[node] = rootInGroupTree(groupTree,node)

    for vector in vecrot2Group:
        vecrot2Group[vector] = rootInGroupTree(groupTree,vecrot2Group[vector])

def rootInGroupTree(groupTree,node):
    if node not in groupTree:
        groupTree[node]=node
    while groupTree[node]!=node:
        node=groupTree[node]
    return node


# input : vec2group sha: groupID  , node2vector:  key : [sha list]
def clusterSet(vecrot2Group,node2vector):
    groupTree = {}
    for node,vectors in node2vector.items():
        group = None
        vectors = list(vectors)
        for vector in vectors:
            if vector not in vecrot2Group:
                continue
            nodeGroup = rootInGroupTree(groupTree,vecrot2Group[vector])
            if group is None:
                group = nodeGroup
            else:
                groupTree[rootInGroupTree(groupTree,vecrot2Group[vector])] = group
    for node  in groupTree:
        groupTree[node] = rootInGroupTree(groupTree,node)

    for vector in vecrot2Group:
        vecrot2Group[vector] = rootInGroupTree(groupTree,vecrot2Group[vector])
