import pandas as pd

from code.tools import *


def extractDomain(filename):
    domains = set()
    ips = set()
    inf = open(filename,encoding='utf-8')
    reader = csv.reader(inf)
    headings = next(reader)
    for data in reader:
        url = data[1]
        if '/' in url:
            vals = url.split('/')
            fulldomain = vals[2].split(':')[0]
        else:
            fulldomain = url
        if checkip(fulldomain):
            ips.add(fulldomain)
        else:
            domains.add(fulldomain)
    return domains,ips


def loadDomains():
    domains,IPs = extractDomain(fileMap['web'])
    registeredDomain= set()
    for key in ['fqdn','region']:
        newDomain,newIP = extractDomain(fileMap[key])
        domains = domains | newDomain
        registeredDomain = registeredDomain | newDomain    
    return domains,registeredDomain

def recordList(data,outname):
    outf = open(outMap[outname],'w+')
    for line in data:
        outf.write(line+'\n')

def makeDomainList():
    domains,registeredDomain = loadDomains()
    print(len(domains),len(registeredDomain),len(domains-registeredDomain))
    recordList(registeredDomain,'registered_domain')
    recordList(domains,'all_domain')

def ip2domain():
    ipMap = {}
    filename = fileMap['fqdn']
    inf = open(filename,encoding='utf-8')
    reader = csv.reader(inf)
    headings = next(reader)
    for data in reader:
        domain = data[1]
        ip = data[2]
        if ip not in ipMap:
            ipMap[ip] =  set()
        ipMap[ip].add(domain)
    for ip in ipMap:
        if len(ipMap[ip])>1:
            print(ipMap[ip])
            
def noRegisterDomain():
    noRegisterDomain = set()
    domains,registeredDomain = loadDomains()
    for domain in domains:
        if domain not in registeredDomain and len(domain.split('.'))==2:
            noRegisterDomain.add(domain)
        
    inf = open(fileMap['region'],encoding='utf-8')
    reader = csv.reader(inf)
    headings = next(reader)
    for data in reader:
        if len(data[2])<1 and not checkip(data[1]):
            domain = data[1].lower()
            noRegisterDomain.add(domain)
    print(len(noRegisterDomain))
    return noRegisterDomain


def staticFeature(domain):
    vals = domain.split('.')
    charNum =0
    numberNum = 0
    allLen = 0
    specialChar = 0
    for char in domain:
        if char == '.':
            continue        
        allLen+=1
        if char>='0' and  char<='9':
            numberNum+=1
        elif char>='a' and char<='z':
            charNum+=1
        else:
            specialChar+=1
    return {   
    'domain':domain,
    'enddomain':vals[-1],
    'block':len(vals),
    'charLen':charNum,
    'numLen':numberNum,
    'specialChar':specialChar    
    } 

def clusterFromDomains(domains):
    check = []
    for domain in domains:
        features = staticFeature(domain)
        check.append(features)
    oneHotFeature(check,'enddomain')
    train = pd.DataFrame(check)    
    train.to_csv(outMap['domainFeature'],index =0)


def DBSCANCluster():
    domain2group = {}
    X = pd.read_csv(outMap['domainFeature'])
    domains = X['domain']
    del(X['domain'])
    # from sklearn.cluster import KMeans
    # y_pred = KMeans(n_clusters=100, random_state=9).fit_predict(X)
    from sklearn.cluster import DBSCAN
    y_pred = DBSCAN().fit_predict(X)
    group2domain = {}
    outf = open(outMap['dga2level'],'w+')
    outf.write('domain,family\n')
    for i in range(len(y_pred)):
        domain2group[domains[i]] = y_pred[i]
        if y_pred[i] not in group2domain:
            group2domain[y_pred[i]] = []
        group2domain[y_pred[i]].append(domains[i])
    for group in group2domain:
        if group <0:
            continue
        for domain in group2domain[group]:
            outf.write('%s,%d\n'%(domain,group))

    return domain2group


def expendDomain():
    secondRegistered = set()
    sha2domain,domain2sha = readDomainSet()
    domain2group = {}
    index = 10000
    for domain in domain2sha:
        domain2group[domain] = index
        index+=1
    domain2group.update(DBSCANCluster())
    clusterSet(domain2group,sha2domain)




def dga2graph():
    sha2domain,domain2sha = readDomainSet()
    outf = open('graph.txt','a+')
    inf = open(outMap['dga2level'],encoding='utf-8')
    reader = csv.reader(inf)    
    headings = next(reader)
    whiteSet = loadWhiteMap()
    for data in reader:
        domain = data[0]
        group = data[1]
        secondDomain = '.'.join(domain.split('.')[-2:])
        if domain in whiteSet or secondDomain in whiteSet:
            print(domain,group)
            continue
        for sha in domain2sha[domain]:
            outf.write('%s,domain2level_%s\n'%(sha,group))




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)
    print(secondDomainSet)
    return whiteSet


def filterDomains():
    domains,registeredDomain = loadDomains()
    unregistered = domains - registeredDomain
    whiteSet = loadWhiteMap()

    whiteSet = filterDomainSet(whiteSet)
    domains = filterDomainSet(domains)

    index = 0
    whiteOut = open(outMap['whiteDomain'],'w+')
    blackOut = open(outMap['blackDomain'],'w+')
    whiteSecond = set()
    for domain in domains:
        vals = domain.split('.')
        if len(vals)<=2:
            secondDomain = domain
        else:
            secondDomain = '.'.join(vals[1:])
        if domain in whiteSet or secondDomain in whiteSet:
            whiteSecond.add(secondDomain)
            whiteOut.write("%s,%s\n"%(domain,secondDomain))
        else:
            blackOut.write("%s,%s\n"%(domain,secondDomain))
    # print(whiteSecond)


def loadDangerDomain():
    result = set()
    for line in open(outMap['blackDomain']).readlines():
        domain = line.split(',')[0]
        result.add(domain)
    return result


       
if __name__=='__main__':
    # makeDomainList()
    # ip2domain()
    # clusterFromDomains(noRegisterDomain())
    # clusterFromDomains(loadDangerDomain())
    # clusterFromDomains(extractDomain(fileMap['web'])[0])
    # DBSCANCluster()
    # expendDomain()
    # checkFamily()
    # dga2graph()
    filterDomains()
    # checkBlackDomain()
