import json
from collections import defaultdict

with open('cfgs.json', 'r') as f:
    cfgs = json.load(f)


def classify_link_loop(cfgs):
    classes = defaultdict(list)

    for func_name, cfg in cfgs.items():
        links = 0
        self_loops = 0

        for block, outs in cfg.items():
            for out in outs:
                links += 1
                if out == block:
                    self_loops += 1
        classes[(links, self_loops >= 1)].append(func_name)

    return classes

def classify_shortest_path_length(cfgs):
    import networkx as nx

    classes = defaultdict(list)

    for func_name, cfg in cfgs.items():
        graph = nx.DiGraph()
        for block, outs in cfg.items():
            for out in outs:
                graph.add_edge(block, out)
        classes[nx.average_shortest_path_length(graph)].append(func_name)

    return classes

# To use the other metric, use classify_shortest_path_length (requires NetworkX)
classes = classify_link_loop(cfgs)

output = ""

for k, vs in classes.items():
    for v in vs:
        start, end = v.split('_')[1:]
        output += f"{'T' if len(vs) > 100 else 'N'} 0x{start} 0x{end}\n"

with open('taint_data.txt', 'w') as f:
    f.write(output)
