malwareGraphExtract.py

import r2pipe
import os
import re
import json
import networkx as nx
import numpy as np
import sys
import matplotlib.pyplot as plt
import argparse
from networkx.algorithms import bipartite
import networkx.algorithms.isomorphism as iso
import networkx.algorithms.graph_hashing as graph_hashing
import pandas as pd

np.set_printoptions(threshold=np.inf)

parser = argparse.ArgumentParser(description='Extract Adjacency matrix / graph from PE')
parser.add_argument('file', nargs='*',  help="The file to analyze")
parser.add_argument('-e', action='store_true',  help="export function to *.gml")
parser.add_argument('-i', help="import gml file to compare with all other function")
p = parser.parse_args()

printFunction = {}
matching_d = {}
iso_l = []
bi_l = []
nodes_l = []
hash_l = []


def allFunction(f):
    """ 
    Print all function from PE 
    sym : Symbol from symbol table
    sub : subroutine, part of a function
    fcn : regular function
    """
    result = []
    x = f.cmdj("aflj")
    for i in range(0, int(f.cmd("aflc"))):
        data = json.dumps(x[i]['name'])
        if ("fcn." in data) or ("entry0" in data) :
            result.append(data.replace('"',''))
    return result


def ActionBasicBloc(nameFunction, *args):
    """  
        Print adjacency_matrix of Graph
        if args = save, save the graph to png
    """
    G = initGraph()
    cpt = 0
    bb = f.cmd("afb " + nameFunction)
 

    for i in bb.splitlines():
        index = i.split(' ')

        if cpt == 0:
            G.add_node(index[0], color='green', style='filled')
        elif cpt == G.number_of_nodes() -1 :
            G.add_node(index[0], color='red', style='filled')
        else:
            G.add_node(index[0])
        cpt = cpt + 1
        if ("j" in index) and ("f" in index) :
            G.add_edges_from([(index[0], index[5]),(index[0],index[7])], weight=1)
        elif ("j" in index) or ("f" in index) :
            G.add_edge(index[0], index[5], weight=1)
        else :
            pass

    color_d = nx.get_node_attributes(G, 'color')
    default_color = 'blue'
    color_l = [color_d.get(node, default_color) for node in G.nodes()]
    if 'save' in args :
        print(nx.to_numpy_matrix(G))
        nx.draw(G, node_color=color_l)
        plt.savefig(str(p.file) +"_graph_" + str(nameFunction) + ".png")
    elif 'print' in args :
        printFunction[str(nameFunction)] = str(G.number_of_nodes())
        #debug
        #if G.number_of_nodes() < 10:
            #print(str(nameFunction))
            #print(nx.to_numpy_matrix(G))
            #[print(nameFunction, n, nbrdict) for n, nbrdict in G.adjacency()]
            #print(bipartite.is_bipartite(G))
            #print(nx.graph_edit_distance(G1,G2)) #is_isomorphic
    elif 'export' in args :
        nx.write_gml(G, "_"+ str(nameFunction)+".gml")           
    elif 'matching' in args :
        if p.i  :
            if os.path.isfile(p.i):
                try :
                    G_up =nx.read_gml(p.i)

                    if G_up.number_of_nodes() < 10 and G_up.number_of_nodes() == G.number_of_nodes() :
                        matching_d[str(nameFunction)] = str(nx.graph_edit_distance(G, G_up)) 
                        #print("Distance between " + p.i +" and " + str(nameFunction) + " : "+  str(nx.graph_edit_distance(G, G_up)))
                        nodes_l.append(G.number_of_nodes())
                        iso_l.append(1) if (nx.is_isomorphic(G, G_up) == True) else iso_l.append(0)
                        bi_l.append(1) if (bipartite.is_bipartite(G) == True) else bi_l.append(0)
                        hash_l.append(graph_hashing.weisfeiler_lehman_graph_hash(G))
                        #print(graph_hashing.weisfeiler_lehman_graph_hash(G_up))
                    else:
                        pass 
                        
                except FileNotFoundError:
                    print("error opening file")
            else:
                pass
    else:
        print(nx.to_numpy_matrix(G))
    
    
def initGraph():
    return nx.MultiDiGraph()


if __name__=="__main__":
    pd.set_option('display.max_rows', None)
    print(p.file)

    if p.file  :
        if os.path.isfile(p.file[0]):
            try :
                f = r2pipe.open(p.file[0])
                f.cmd("e asm.bytes = 0")
                f.cmd("e scr.utf8 = true")
                f.cmd("aaa")
                f.cmd("aan") #print alias function like -A
                listeDesFonctions = allFunction(f)
                print("Number of methods : " + str(len(listeDesFonctions)))
                if p.e :
                    [ActionBasicBloc(fc, "export") for fc in listeDesFonctions]
                elif p.i :
                    [ActionBasicBloc(fc, "matching") for fc in listeDesFonctions]
                    label = ['Function Name', 'GED']
                    df = pd.DataFrame(matching_d.items(), columns=label)
                    df["node"] = pd.Series(nodes_l, dtype='float64')
                    df["Isomoprhism"] = pd.Series(iso_l, dtype='float64')
                    df["Bipartite"] = pd.Series(bi_l, dtype='float64')
                    df["weisfeiler_lehman"] = pd.Series(hash_l)
                    print(df)
                else:
                    [ActionBasicBloc(fc, "print") for fc in listeDesFonctions]
                    df = pd.DataFrame(printFunction.items(), columns=['Method','Node'])
                    print(df[pd.to_numeric(df['Node']) < int(10) ])

            except FileNotFoundError:
                print("error opening file")
        else :
            print("file doesn't exist")
    else:
        parser.print_help()