import os
import re
import pandas as pd
from publicsuffix import PublicSuffixList
from sklearn import preprocessing
from urllib import parse
import codecs
import networkx as nx

psl_file = codecs.open('suffix.dat', encoding='utf8')
psl = PublicSuffixList(psl_file)
p = re.compile('^((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)$')  
def checkip(ip):
    if p.match(ip):
        return True
    else:
        return False
phase1 = 'D:/tempdata/phase1/phase1/trace2_test'
phase1 = '/home/OpenCode/malware/trace2_test'
phase2 = '/home/OpenCode/malware/phase2/trace2_test'



def noHeadDomain(domain):
    vals = domain.split('.')
    if vals[0] in ['www','wap']:
        return '.'.join(vals[1:])
    return domain

def processDomain(data):
    return data.map(noHeadDomain)

def readPSLDomain(data):
    data['pslDomain'] = data.map(lambda x:psl.get_public_suffix(x))



def whiteDomainSet():
    whiteSet = set()
    for line in open('prank.top.1m.20180322').readlines():
        domain = line.split('\t')[0]
        whiteSet.add(psl.get_public_suffix(domain))
    for line in open('top-1m.csv').readlines():
        domain = line.split(',')[1].strip()
        whiteSet.add(psl.get_public_suffix(domain))
    print('all white %d'%len(whiteSet))
    return whiteSet

def url2file(df):
    result = {}
    def addUrlMap(data):
        domain = data['clean_domain']
        if checkip(domain):
            return -1
        fileurl = data['url']
        filename = fileurl[fileurl.find(domain)+len(domain):]
        if '%' in filename:
            filename = parse.unquote(filename)
        result[domain] = filename
    df.apply(addUrlMap,axis=1)
    return result

def cleanDF(df,blackSet):
    return df[df['clean_domain'].map(lambda x : psl.get_public_suffix(x) in blackSet)][['sha1','clean_domain']]

def loadBlackNodes():
    G = nx.read_adjlist('black.txt')
    return set(G.nodes)

def url2domain(x):
    func = lambda x: x.split('://')[-1].split('/')[0].split(':')[0]
    try:
        result = func(x)
    except:
        result = None
        print(x,result)
    return result

def loadDomainsFromDir(basedir,outname):
    webfile = os.path.join(basedir,'trojan_web.csv')
    ipfile = os.path.join(basedir,'fqdn_ip.csv')
    reginfofile = os.path.join(basedir,'fqdn_reginfo.csv')
    domain_ip = pd.read_csv(ipfile)
    domain_ip['clean_domain'] = processDomain((domain_ip['site']))
    domain_reg = pd.read_csv(reginfofile)
    domain_null = pd.isnull(domain_reg['email_address'])
    domain_no_user = set(domain_reg[~domain_null].domain)
    domain_reg['clean_domain'] = processDomain(domain_reg['domain'])
    webdf = pd.read_csv(webfile)
    webdf['clean_domain'] = processDomain(webdf['url'].map(url2domain))
    ipSet = set(domain_ip['clean_domain'])
    regSet = set(domain_reg['clean_domain'])
    webSet = set(webdf['clean_domain'])
    allDomains = webSet | ipSet | regSet
    nxDomains = webSet - ipSet - regSet
    fileMap = url2file(webdf)
    whiteSet = whiteDomainSet()
    result = []
    pslMap = {}
    allPslSet = set(psl.get_public_suffix(domain) for domain in allDomains)
    blackSet = loadBlackNodes() & allPslSet
    # for domain in nxDomains | blackSet:
    for domain in nxDomains:
        if checkip(domain):
            continue
        pslDomain = psl.get_public_suffix(domain)
        if pslDomain in whiteSet:
            continue
        result.append({
            'domain':domain,
            'host':pslDomain,
            'top':pslDomain[pslDomain.find('.'):],
            'file':fileMap[domain] if domain in fileMap else 'None'
        })
        if pslDomain not in pslMap:
            pslMap[pslDomain] = 0
        pslMap[pslDomain]+=1
    for data in result:
        data['subCount'] = pslMap[data['host']]
    df = pd.DataFrame(result)
    df.to_csv(outname,index=False)
    # with open('blackDomain.txt','w+') as outf:
    #     for domain in domain_no_user:
    #         outf.write(domain+'\n')
    #
    # pslSet = set(df.host)
    # domainMap = pd.concat([
    #     cleanDF(domain_ip,pslSet),cleanDF(domain_reg,pslSet),cleanDF(webdf,pslSet)
    #     ])
    # domainMap.drop_duplicates().to_csv('sha2domain.csv',index=False)

if __name__=='__main__':
    loadDomainsFromDir(phase1,'cluster_domain_1.csv')
    loadDomainsFromDir(phase2,'cluster_domain_2.csv')
