grapher.py

import json
import argparse
from graphviz import Digraph
import colorsys
import os
import sys


def load_data(file_path):
    with open(file_path, 'r') as file:
        return json.load(file)

def load_colors():
    with open("colors.json", 'r') as file:
        return json.load(file)

def format_label(label, max_length=10):
    words = label.split()
    lines = []
    current_line = []
    current_length = 0

    for word in words:
        if current_length + len(word) + 1 > max_length:
            lines.append(" ".join(current_line))
            current_line = [word]
            current_length = len(word)
        else:
            current_line.append(word)
            current_length += len(word) + 1

    if current_line:
        lines.append(" ".join(current_line))

    return "\n".join(lines)

def count_children(data, child_counts=None):
    if child_counts is None:
        child_counts = {}

    for key, value in data.items():
        child_counts[key] = 0
        if isinstance(value, dict):
            child_counts[key] += len(value)
            for child_key in value.keys():
                count_children(value, child_counts)
        elif isinstance(value, list):
            child_counts[key] += len(value)

    return child_counts

def darken_color(hex_color, amount=0.2):
    hex_color = hex_color.lstrip('#')
    r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
    h, l, s = colorsys.rgb_to_hls(r / 255, g / 255, b / 255)
    l = max(0, l - amount)
    r, g, b = colorsys.hls_to_rgb(h, l, s)
    return f"#{int(r * 255):02X}{int(g * 255):02X}{int(b * 255):02X}"

def add_nodes_edges(graph, data, parent=None, color="#FFFFFF", child_counts=None):
    for key, value in data.items():
        formatted_key = format_label(key)
        node_color = color
        border_color = darken_color(node_color)
        shape = "circle" if child_counts.get(key, 0) > 4 else "box"

        graph.node(
            formatted_key,
            style='rounded,filled',
            fillcolor=node_color,
            color=border_color,
            penwidth="2",
            shape=shape,
        )

        if parent:
            formatted_parent = format_label(parent)
            # Use "dotted" style for edges
            graph.edge(formatted_parent, formatted_key, color=border_color, style="dotted")

        if isinstance(value, dict):
            add_nodes_edges(graph, value, key, node_color, child_counts)
        elif isinstance(value, list):
            for item in value:
                formatted_item = format_label(item)
                graph.node(formatted_item, style='rounded,filled', fillcolor=node_color, color=border_color, penwidth="2", shape="box")
                graph.edge(formatted_key, formatted_item, color=border_color, style="dotted")


def create_colored_mindmap(data, output_type="png"):
    colors = load_colors()
    graph = Digraph(format=output_type, engine='twopi')
    graph.attr(overlap='false', splines='true', rankdir='TB', dpi='300')

    child_counts = count_children(data)

    for root, content in data.items():
        formatted_root = format_label(root)
        graph.node(formatted_root, shape='circle', style='filled', color='#A9CCE3')

        for n1_key, n1_content in content.items():
            n1_color = colors.get(n1_key, "#FFFFFF")
            border_color = darken_color(n1_color)
            formatted_n1_key = format_label(n1_key)
            graph.node(
                formatted_n1_key,
                shape="circle",
                fontname="Helvetica-Bold",
                style="bold,filled",
                fillcolor=n1_color,
                color=border_color,
                penwidth="2"
            )
            graph.edge(formatted_root, formatted_n1_key, color=border_color, style="bold")

            if isinstance(n1_content, dict):
                add_nodes_edges(graph, n1_content, n1_key, n1_color, child_counts)
            elif isinstance(n1_content, list):
                for item in n1_content:
                    formatted_item = format_label(item)
                    graph.node(formatted_item, style='rounded,filled', fillcolor=n1_color, color=border_color, penwidth="2", shape="box")
                    graph.edge(formatted_n1_key, formatted_item, color=border_color, style="bold")

    return graph

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate the Cyber Security MindMap based on the Json file")
    parser.add_argument("-data", type=str, required=True, help="Path to the JSON data file.")
    parser.add_argument("-type", type=str, choices=["png", "pdf"], default="png", help="Output file type (png or pdf).")
    parser.add_argument("-title", type=str, help="Title for the mindmap.")

    args = parser.parse_args()

    data = load_data(args.data)
    mindmap = create_colored_mindmap(data, output_type=args.type)

    if args.title:
        formatted_title = format_label(f"{args.title} - clb 2024", max_length=20)
        mindmap.node(
            "title",
            formatted_title,
            shape="box",
            style="bold,filled",
            color="#FFFFFF",
            fontname="Helvetica-Bold",
            fontsize="18",
            fixedsize="false"
        )

        mindmap.attr(overlap="false")
        mindmap.attr(splines="true")

output_file = f"cybersecurity_mindmap.{args.type}"
with open(os.devnull, 'w') as fnull:
    sys.stderr = fnull
    try:
        mindmap.render(output_file, view=True)
    finally:
        sys.stderr = sys.__stderr__  # Restore original stderr