import json
from code.domain import readDomains,clusterDomains,group2file
from code.secondDomain import saveDGA
import pandas as pd
from code.tools import oneHotFeature,checkip
from sklearn.cluster import DBSCAN
from publicsuffix import PublicSuffixList
from sklearn import preprocessing
import csv
import codecs
psl_file = codecs.open('suffix.dat', encoding='utf8')
psl = PublicSuffixList(psl_file)

def makeGraphFile():
    readDomains()
    clusterDomains(json.load(open('data/noRegisterBlackTwoLevel.txt')), 'data/featureCluster.txt')
    group2file()
    saveDGA()

def addToSetMap(data,key,val):
    if key not in data:
        data[key] = set()
    data[key].add(val)

def readDGADomains():
    domainMap = {}
    filenames = [
        # '/home/malware/geek/data/keyDomainMap.txt',
        # '/home/malware/geek/data/featureCluster.txt',
        '/home/OpenCode/geekai/geek/code/secondCluster.txt'

    ]
    for filename in filenames:
        print(filename)
        data = json.load(open(filename))
        for key in data:
            for domain in data[key]:
                addToSetMap(domainMap,domain,key)


    return list(domainMap.keys())

def staticFeature(domain,suffixDomainMap,meanLenMap):
    vals = domain.split('.')
    block = len(vals)
    suffix = psl.get_public_suffix(domain)
    charNum =0
    numberNum = 0
    allLen = 0
    specialChar = 0
    for char in suffix:
        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':block,
    # 'allLen':allLen,
    'charLen':charNum/(allLen+0.0),
    'numLen':numberNum/(allLen+0.0),
    'specialChar':specialChar/(allLen+0.0),
    'subDomainNum':1 if suffixDomainMap[suffix]>100 else 0,
    'meanDomain':meanLenMap[suffix]
    }

def clusterDomains():
    domains = readDGADomains()
    suffixMap = {}
    sumLenMap = {}
    meanLenMap = {}
    for domain in domains:
        suffix = psl.get_public_suffix(domain)
        if suffix not in suffixMap:
            suffixMap[suffix] = 0
            sumLenMap[suffix] = 0
        suffixMap[suffix] += 1
        sumLenMap[suffix] += len(domain) - len(suffix)
    for suffix in suffixMap:
        meanLenMap[suffix] = sumLenMap[suffix]/(suffixMap[suffix]+0.0)
    check = []
    for domain in domains:
        features = staticFeature(domain,suffixDomainMap=suffixMap,meanLenMap=meanLenMap)
        check.append(features)
    oneHotFeature(check, 'enddomain')
    train = pd.DataFrame(check)
    train.to_csv('domainFeature.csv',index=0)

def cluster():
    df = pd.read_csv('domainFeature.csv')
    domains = list(df['domain'])
    suffixGraph = {}
    del(df['domain'])
    X = df
    # X['subDomainNum'] = preprocessing.scale(X['subDomainNum'], axis=0, with_mean=True, with_std=True, copy=True)
    X = preprocessing.scale(X, axis=0, with_mean=True, with_std=True, copy=True)
    result = {}
    y_pred = DBSCAN(eps = 0.5).fit_predict(X)
    for i in range(len(X)):
        key = y_pred[i]
        if key not in result:
            result[key] = set()
        result[key].add(psl.get_public_suffix(domains[i]))
    for key in result:
        if key==-1:
            continue
        print(result[key])
        for suffix in result[key]:
            suffixGraph[suffix] = 'suffix_%d'%key
    json.dump(suffixGraph,open('suffix_graph.txt','w+'))

def write2graph():
    suffixMap = json.load(open('suffix_graph.txt'))
    suffix2sha = {}
    filename = '/home/malware/trace2_test/trojan_web.csv'
    inf = open(filename, encoding='utf-8')
    reader = csv.reader(inf)
    next(reader)
    for data in reader:
        sha = data[0]
        url = data[1]
        if '/' in url:
            vals = url.split('/')
            fulldomain = vals[2].split(':')[0]
        else:
            fulldomain = url
        if not checkip(fulldomain):
            suffix = psl.get_public_suffix(fulldomain)
            if suffix not in suffix2sha:
                suffix2sha[suffix] = set()
                suffix2sha[suffix].add(sha)
    outf_graph = open('graph.txt','a+')
    for suffix in suffix2sha:
        if suffix not in suffixMap:
            continue
        groupID = suffixMap[suffix]
        for sha in suffix2sha[suffix]:
            outf_graph.write('%s,%s\n'%(sha,groupID))




if __name__=='__main__':
    clusterDomains()
    cluster()
    write2graph()