malware-classification/convert.py at main · thanhmcisai/malware-classification · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
from networkx.classes.multidigraph import MultiDiGraph
from networkx.classes.multigraph import MultiGraph
from networkx.drawing.nx_agraph import to_agraph
from networkx.readwrite import json_graph
import networkx as nx
from typing import Union
import json
import os


def convert_dot_to_networkx(
        dot_file_path: str,
        saved_file_path: str = '') -> Union[MultiDiGraph, MultiGraph]:
    if not os.path.exists(dot_file_path):
        print(f'{dot_file_path} does not exist')
        assert(os.path.exists(dot_file_path))
    if os.path.getsize(dot_file_path) == 0:
        print(f'{dot_file_path} is empty')
        assert(os.path.getsize(dot_file_path) > 0)

    print(f'converting {dot_file_path}...')

    # https://networkx.org/documentation/stable/reference/generated/networkx.drawing.nx_pydot.read_dot.html?highlight=read_dot#networkx.drawing.nx_pydot.read_dot
    G = nx.drawing.nx_pydot.read_dot(dot_file_path)

    print(G.nodes())
    print(G.edges())

    if saved_file_path != '':
        A = to_agraph(G)
        A.draw(saved_file_path, format='png', prog='dot')

    return G

if __name__ == "__main__":
    print(convert_dot_to_networkx('m1.dot', 'result.png'))