import pandas as pd
import re
import codecs
from publicsuffix import PublicSuffixList
import Levenshtein
from urllib import parse


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 filterName(filename):
    if '%' in filename:
        return parse.unquote(filename)
    else:
        return filename

def checkSet(df,key,set):
    return df[df[key].map(lambda x:x in set)]

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 checkip(ip):
    if p.match(ip):
        return True
    else:
        return False

psl_file = codecs.open('/home/OpenCode/malware/geekai/suffix.dat', encoding='utf8')
psl = PublicSuffixList(psl_file)

phase1 = '/home/OpenCode/malware/trace2_test/trojan_web.csv'
phase2 = '/home/OpenCode/malware/phase2/trace2_test/trojan_web.csv'
df1 = pd.read_csv(phase1)
df2 = pd.read_csv(phase2)
df = pd.concat([df1,df2])
df['host'] = df.url.map(lambda x: x.split('://')[-1].split('/')[0].split(':')[0])
df['clean_url'] = df.url.map(filterName)
dfIP = df[df.host.map(checkip)]
ipCnt = dfIP.groupby('host').count()

dfDomain = df[~df.host.map(checkip)]
dfDomain['psl'] = dfDomain.host.map(lambda x:psl.get_public_suffix(x))
dfDomain['top'] = dfDomain.psl.map(lambda x: '.'.join(x.split('.')[1:]))
dfDomain['url_len'] = dfDomain.url.map(lambda x:len(x))


t1 = dfDomain[['sha1','host']].drop_duplicates().groupby('host',as_index=False).agg({'sha1':'count'})
t1.columns = ['host','host_sha_cnt']
t2 = dfDomain[['psl','host']].drop_duplicates().groupby('psl',as_index=False).agg({'host':'count'})
t2.columns = ['psl','subname_cnt']
t3 = dfDomain[['sha1','psl']].drop_duplicates().groupby('psl',as_index=False).agg({'sha1':'count'})
t3.columns = ['psl','psl_sha_cnt']

dfDomain = pd.merge(dfDomain,t1,on='host')
dfDomain = pd.merge(dfDomain,t2,on='psl')
dfDomain = pd.merge(dfDomain,t3,on='psl')

whiteSet = whiteDomainSet()

firstSet = set(dfDomain[dfDomain.psl_sha_cnt==1].psl) - whiteSet
firstDF = checkSet(dfDomain,'psl',firstSet)
firstDF.to_csv('first_url.csv',index=False)

secondDF = dfDomain[dfDomain.psl.map(lambda x:x not in whiteSet)]


fqdnDF1 = pd.read_csv('/home/OpenCode/malware/trace2_test/fqdn_ip.csv')
fqdnDF2 = pd.read_csv('/home/OpenCode/malware/phase2/trace2_test/fqdn_ip.csv')
fqdn = pd.concat([fqdnDF1,fqdnDF2])
fqdn = fqdn[~fqdn.site.map(checkip)]
fqdn['psl'] = fqdn.site.map(lambda x:psl.get_public_suffix(x))

t = fqdn[['site','ip']].drop_duplicates().groupby('site',as_index=False).count()
t.columns = ['host','ip_cnt']

nxSet = set(dfDomain.psl)-set(fqdn.psl)

apkDomain = dfDomain[dfDomain.url.map(lambda x:x.endswith('apk'))]


