Update app.py

app.py

@@ -5,6 +5,7 @@ import re
 import networkx as nx
 import matplotlib.pyplot as plt
 import io
+import graphviz
 
 # Mock existing resources (customizable)
 existing_resources = {
@@ -45,6 +46,7 @@ arm_schema = {
     "required": ["$schema", "contentVersion", "resources"]
 }
 
+# Function to validate the ARM template against the schema
 def validate_template(template):
     try:
         validate(instance=template, schema=arm_schema)
@@ -52,6 +54,7 @@ def validate_template(template):
     except ValidationError as e:
         return False, f"Template validation error: {e.message}"
 
+# Function to resolve template expressions (like parameters and variables)
 def resolve_template_expressions(template):
     def resolve_expression(match):
         expr = match.group(1)
@@ -66,6 +69,7 @@ def resolve_template_expressions(template):
     resolved_str = re.sub(r'\[(\w+\.\w+)\]', resolve_expression, template_str)
     return json.loads(resolved_str)
 
+# Function to estimate the cost of resources
 def estimate_cost(resources):
     # This is a simplified cost estimation. In a real scenario, you'd need to
     # integrate with Azure's pricing API or use a more comprehensive pricing model.
@@ -87,6 +91,7 @@ def estimate_cost(resources):
     
     return total_cost
 
+# Function to create a dependency graph of resources
 def create_dependency_graph(resources):
     G = nx.DiGraph()
     for resource in resources:
@@ -97,6 +102,7 @@ def create_dependency_graph(resources):
             G.add_edge(dep, resource_name)
     return G
 
+# Function to plot the dependency graph
 def plot_dependency_graph(G):
     plt.figure(figsize=(12, 8))
     pos = nx.spring_layout(G)
@@ -107,6 +113,7 @@ def plot_dependency_graph(G):
     plt.title("Resource Dependency Graph")
     return plt
 
+# Function to lint the ARM template
 def lint_template(template):
     lint_results = []
     
@@ -128,26 +135,30 @@ def lint_template(template):
     
     return lint_results
 
+# Function to generate a system diagram
 def generate_system_diagram(resources):
-    diagram = ["graph TD"]
-    resource_nodes = {}
-    
+    # Create a new Graphviz graph
+    dot = graphviz.Digraph()
+    dot.attr(rankdir='TB')  # Top to bottom layout
+
     # Create nodes for each resource
     for i, resource in enumerate(resources):
         resource_type = resource['type'].split('/')[-1]
         node_id = f"R{i}"
-        resource_nodes[resource['name']] = node_id
-        diagram.append(f"{node_id}[{resource_type}<br/>{resource['name']}]")
-    
+        dot.node(node_id, f"{resource_type}\n{resource['name']}")
+
     # Create connections based on dependencies
     for resource in resources:
         if 'dependsOn' in resource:
             for dependency in resource['dependsOn']:
-                if dependency in resource_nodes:
-                    diagram.append(f"{resource_nodes[dependency]} --> {resource_nodes[resource['name']]}")
-    
-    return "\n".join(diagram)
+                # Find the node id for the dependency
+                dep_id = next((f"R{j}" for j, r in enumerate(resources) if r['name'] == dependency), None)
+                if dep_id:
+                    dot.edge(dep_id, f"R{resources.index(resource)}")
+
+    return dot
 
+# Main function to simulate the deployment
 def simulate_deployment(arm_template):
     try:
         template = json.loads(arm_template)
@@ -155,7 +166,7 @@ def simulate_deployment(arm_template):
         # Validate template
         is_valid, validation_message = validate_template(template)
         if not is_valid:
-            return [], [], [], [], 0, None, [], "", validation_message
+            return [], [], [], [], 0, None, [], None, validation_message
 
         # Resolve expressions
         resolved_template = resolve_template_expressions(template)
@@ -208,7 +219,7 @@ def simulate_deployment(arm_template):
 
         return resources_to_create, resources_to_update, resources_to_delete, resource_details, estimated_cost, dependency_graph, lint_results, system_diagram, "Simulation completed successfully."
     except json.JSONDecodeError:
-        return [], [], [], [], 0, None, [], "", "Invalid JSON format. Please check your ARM template."
+        return [], [], [], [], 0, None, [], None, "Invalid JSON format. Please check your ARM template."
 
 # Streamlit UI
 st.title("Comprehensive ARM Template Simulator")
@@ -284,7 +295,7 @@ if st.button("Simulate Template"):
 
         st.write("### System Diagram:")
         if system_diagram:
-            st.mermaid(system_diagram)
+            st.graphviz_chart(system_diagram)
         else:
             st.write("Unable to generate system diagram.")