import pandas as pd
import os
from domain import whiteDomainSet,checkip,psl
import networkx as nx

phase1 = '/home/OpenCode/malware/trace2_test'
basedir = phase1
webfile = os.path.join(basedir, 'trojan_web.csv')
ipfile = os.path.join(basedir, 'fqdn_ip.csv')
reginfofile = os.path.join(basedir, 'fqdn_reginfo.csv')
phone = os.path.join(basedir, 'trojan_phone.csv')
email = os.path.join(basedir, 'trojan_email.csv')

def domain2host(domain):
    if checkip(domain):
        return domain
    else:
        return psl.get_public_suffix(domain)

whiteSet = whiteDomainSet()
ipDF = pd.read_csv(ipfile)
ipDF = ipDF.drop_duplicates()
ipDF['psl'] = ipDF['site'].map(domain2host)

domainDF = pd.read_csv(reginfofile)
domainDF = domainDF.drop_duplicates()
domainDF['psl'] = domainDF['domain'].map(domain2host)


edges = []
def addIPEdge(data):
    edge = (data['ip'],data['psl'])
    if edge[0] != edge[1]:
        edges.append(edge)


def addMailEdge(data):
    edge = (data['email_address'], data['psl'])
    if edge[0] != edge[1]:
        edges.append(edge)


ipDF.apply(addIPEdge,axis=1)
domainDF.apply(addMailEdge,axis=1)

G = nx.Graph()
G.add_edges_from(edges)
initNodes = set(G.nodes) & whiteSet
nodes = []
for n in initNodes:
    edges = nx.dfs_successors(G, n)
    for k,v in edges.items():
        nodes.append(k)
        nodes.extend(v)
nodes = set(G.nodes)-set(nodes)
pslSet = set(domainDF.psl) & set(ipDF.psl) & set(nodes)
newG = G.subgraph(nodes)
nx.write_adjlist(newG,'black.txt')

webDF = pd.read_csv(webfile)
webDF['domain'] = webDF['url'].map(lambda x: x.split('://')[-1].split('/')[0].split(':')[0])
webDF['psl'] =  webDF['domain'].map(domain2host)


def addEdges(df,key,edges):
    t = df[df.psl.map(lambda x:x in pslSet)]
    def addEdge(data):
        edge = (data[key],data['sha1'])
        edges.append(edge)
    t.apply(addEdge,axis=1)


shaSet = set(webDF.sha1) | set(ipDF.sha1) | set(domainDF.sha1) | set(pd.read_csv(reginfofile).sha1) | set(pd.read_csv(phone).sha1)
print('shaSet%d'%len(shaSet))
edges = []
# addEdges(webDF,'psl',edges)
# addEdges(ipDF,'psl',edges)
# addEdges(domainDF,'psl',edges)


resultG = nx.Graph()
for line in open('/home/OpenCode/malware/geek/code/graph.txt').readlines():
    vals = line.split(',')
    edges.append((vals[0],vals[1]))
groupIndex = 0
resultG.add_edges_from(edges)
outf = open('/home/malware/trace2.csv', 'w+')
noCluster = 0
print('nodes:%d'%len(resultG))
for largest_components in sorted(nx.connected_components(resultG), key=len, reverse=False):
    g = resultG.subgraph(largest_components)
    nodeSet = set(g.nodes) & shaSet
    groupLen = len(nodeSet)
    if len(nodeSet) > 0:
        groupIndex += 1
        if groupLen < 2:
            noCluster += groupLen
        for node in nodeSet:
            outf.write('%s,%d\n' % (node, groupIndex))
print('maxLen:%d' % groupLen)
print('totalGroup:%d' % groupIndex)