From 90df924a7ffdaa3cb35eb729cfef1633907e8888 Mon Sep 17 00:00:00 2001
From: Philippe Merle <philippe.merle@inria.fr>
Date: Sat, 28 Dec 2024 17:38:28 +0100
Subject: [PATCH] Add support for custom icons

---
 bin/kube-diagrams | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/bin/kube-diagrams b/bin/kube-diagrams
index e9ff934..ac38188 100755
--- a/bin/kube-diagrams
+++ b/bin/kube-diagrams
@@ -3,6 +3,7 @@ import argparse
 import copy
 from diagrams import Diagram, Edge, Cluster
 from diagrams.aws.enablement import ManagedServices
+from diagrams.custom import Custom
 import os
 import yaml
 from pprint import pprint
@@ -17,9 +18,12 @@ def query_path(data, path, default=None):
             return default
     return tmp.get(paths[-1], default)
 
+# Directory where this script is.
+DIRNAME = os.path.dirname(__file__)
+
 # Load configuration.
 config = {}
-with open(os.path.dirname(__file__) + "/kube-diagrams.yaml") as f:
+with open(DIRNAME + "/kube-diagrams.yaml") as f:
     config = yaml.safe_load(f) # load YAML config file
 
 # Undefined edge
@@ -47,6 +51,13 @@ def create_diagram_node(resource, fontcolor="black"):
     diagram_node_class = ManagedServices # default diagram node class
     node_config = get_node_config(resource)
     if node_config != None: # config found
+        custom_icon = node_config.get("custom_icon")
+        if custom_icon != None: # custom_icon defined
+            return Custom(
+                resource["metadata"]["name"],
+                custom_icon.replace("$KD", DIRNAME),
+                fontcolor=fontcolor
+            )
         # Get diagram node class name
         diagram_node_classname = node_config.get("diagram_node_classname")
         if diagram_node_classname != None: # classname defined