Create network_graph_generator.py

network_graph_generator.py

@@ -0,0 +1,118 @@
+import graphviz
+import json
+from tempfile import NamedTemporaryFile
+import os
+
+def generate_network_graph(json_input: str, output_format: str) -> str:
+    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': '0.5',
+                'layout': 'neato'
+            },
+            node_attr={
+                'fixedsize': 'false'
+            }
+        )
+        
+        base_color = '#19191a'
+        lightening_factor = 0.12
+        
+        nodes = data.get('nodes', [])
+        connections = data.get('connections', [])
+        
+        for i, node in enumerate(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}")
+            
+            current_depth = node.get('depth', i % 5)
+            
+            if not isinstance(base_color, str) or not base_color.startswith('#') or len(base_color) != 7:
+                base_color_safe = '#19191a'
+            else:
+                base_color_safe = base_color
+                
+            base_r = int(base_color_safe[1:3], 16)
+            base_g = int(base_color_safe[3:5], 16)
+            base_b = int(base_color_safe[5:7], 16)
+            
+            current_r = base_r + int((255 - base_r) * current_depth * lightening_factor)
+            current_g = base_g + int((255 - base_g) * current_depth * lightening_factor)
+            current_b = base_b + int((255 - base_b) * current_depth * lightening_factor)
+            
+            current_r = min(255, current_r)
+            current_g = min(255, current_g)
+            current_b = min(255, current_b)
+            
+            node_color = f'#{current_r:02x}{current_g:02x}{current_b:02x}'
+            font_color = 'white' if current_depth * lightening_factor < 0.6 else 'black'
+            font_size = max(9, 14 - (current_depth * 1))
+            
+            if node_type == 'server':
+                shape = 'box'
+            elif node_type == 'database':
+                shape = 'cylinder'
+            elif node_type == 'user':
+                shape = 'ellipse'
+            elif node_type == 'service':
+                shape = 'hexagon'
+            else:
+                shape = 'circle'
+            
+            dot.node(
+                node_id,
+                label,
+                shape=shape,
+                style='filled',
+                fillcolor=node_color,
+                fontcolor=font_color,
+                fontsize=str(font_size)
+            )
+        
+        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='10',
+                penwidth=penwidth
+            )
+
+        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)}"