network_graph_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os

def generate_network_graph(json_input: str, output_format: str) -> str:
    """
    Generates a network graph from JSON input.

    Args:
        json_input (str): A JSON string describing the network graph structure.
                          It must follow the Expected JSON Format Example below.

    Expected JSON Format Example:
    {
      "nodes": [
        {
          "id": "server1",
          "label": "Web Server",
          "type": "server"
        },
        {
          "id": "db1",
          "label": "Database",
          "type": "database"
        },
        {
          "id": "user1",
          "label": "User",
          "type": "user"
        },
        {
          "id": "service1",
          "label": "API Service",
          "type": "service"
        }
      ],
      "connections": [
        {
          "from": "user1",
          "to": "server1",
          "label": "HTTP Request",
          "weight": 2
        },
        {
          "from": "server1",
          "to": "service1",
          "label": "API Call",
          "weight": 3
        },
        {
          "from": "service1",
          "to": "db1",
          "label": "Query",
          "weight": 1
        }
      ]
    }

    Returns:
        str: The filepath to the generated PNG image file.
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        if 'nodes' not in data or 'connections' not in data:
            raise ValueError("Missing required fields: nodes or connections")
            
        dot = graphviz.Graph(
            name='NetworkGraph',
            format='png',
            engine='neato',
            graph_attr={
                'overlap': 'false',
                'splines': 'true',
                'bgcolor': 'white',
                'pad': '1.0',
                'layout': 'neato',
                'sep': '+20',
                'esep': '+10',
                'nodesep': '1.5',
                'concentrate': 'false',
                'maxiter': '1000'
            },
            node_attr={
                'fixedsize': 'false'
            },
            edge_attr={
                'labeldistance': '3.0',
                'labelangle': '15',
                'labelfloat': 'true'
            }
        )
        
        type_colors = {
            'server': '#BEBEBE',
            'service': '#B8D4F1',
            'database': '#A8E6CF',
            'user': '#FFF9C4',
            'default': '#BEBEBE'
        }
        
        nodes = data.get('nodes', [])
        connections = data.get('connections', [])
        
        for node in nodes:
            node_id = node.get('id')
            label = node.get('label')
            node_type = node.get('type', 'default')
            
            if not all([node_id, label]):
                raise ValueError(f"Invalid node: {node}")
            
            node_color = type_colors.get(node_type, type_colors['default'])
            font_color = 'black'
            
            if node_type == 'server':
                shape = 'box'
                style = 'filled,rounded'
            elif node_type == 'database':
                shape = 'cylinder'
                style = 'filled,rounded'
            elif node_type == 'user':
                shape = 'ellipse'
                style = 'filled,rounded'
            elif node_type == 'service':
                shape = 'hexagon'
                style = 'filled,rounded'
            else:
                shape = 'circle'
                style = 'filled,rounded'
            
            dot.node(
                node_id,
                label,
                shape=shape,
                style=style,
                fillcolor=node_color,
                fontcolor=font_color,
                fontsize='12'
            )
        
        for connection in connections:
            from_node = connection.get('from')
            to_node = connection.get('to')
            label = connection.get('label', '')
            weight = connection.get('weight', 1)
            
            if not all([from_node, to_node]):
                raise ValueError(f"Invalid connection: {connection}")
            
            penwidth = str(max(1, min(5, weight)))
            
            dot.edge(
                from_node,
                to_node,
                label=label,
                color='#4a4a4a',
                fontcolor='#4a4a4a',
                fontsize='8',
                penwidth=penwidth,
                labeldistance='2.5',
                labelangle='10',
                labelfloat='true'
            )
            
        with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
            dot.render(tmp.name, format=output_format, cleanup=True)
            return f"{tmp.name}.{output_format}"
            
    except json.JSONDecodeError:
        return "Error: Invalid JSON format"
    except Exception as e:
        return f"Error: {str(e)}"