import pandas as pd
import networkx as nx

def loadDF(filename):
    df = pd.read_csv(filename, header=None)
    df.columns = ['sha1', 'groupID']
    t = df.groupby('groupID', as_index=False)['sha1'].agg('count')
    t.columns = ['groupID','cnt']
    return pd.merge(df,t,on='groupID')

def checkTrace(filename = '/home/OpenCode/malware/geek/code/trace2.csv'):
    edges = []

    def addToEdges(data):
        edges.append((data['sha1'], data['groupID']))

    df = loadDF(filename)
    df.apply(addToEdges, axis=1)
    G = nx.Graph()
    G.add_edges_from(edges)
    shaSet = set(df.sha1)
    groupIndex = 0
    noCluster = 0
    for largest_components in sorted(nx.connected_components(G), key=len, reverse=False):
        g = G.subgraph(largest_components)
        relatedNodes = []
        nodeSet = set(g.nodes) & shaSet
        groupID = set(g.nodes) - shaSet
        groupLen = len(nodeSet)
        if len(nodeSet) > 0:
            groupIndex += 1
            if groupLen < 2:
                noCluster += groupLen
        print(noCluster, groupIndex)
        print(groupLen, groupID)

def compare():
    df1 = loadDF('/home/OpenCode/malware/geek/code/trace2.csv')
    df2 = loadDF('/home/OpenCode/malware/geekai/trace2.csv')
    set1 = set(df1[df1.cnt>1].sha1)
    set2 = set(df2[df2.cnt==1].sha1)
    bothset = set1 & set2
    print(df1[df1.sha1.map(lambda x:x in bothset)])


checkTrace('/home/malware/phase2_trace2.csv')
# compare()