Add support for prometheus metrics

Author: righettod

scripts/generate-kubernetes-cluster-representation-from-kubernetes-metrics-server-data.py

@@ -2,13 +2,13 @@
 import argparse
 import re
 import collections
+from sqlalchemy import true
 from termcolor import colored
 
 """
-Script to generate a representations from information leaked from an instance 
-of Kubernetes Metrics Server via its endpoints "/metrics".
+Script to generate a representations from information leaked from an instance of Kubernetes Metrics Server via its endpoints "/metrics".
 
-Also support output from an instance of Prometheus.
+Support also the output from an instance of an exporter for Prometheus that is using the Prometheus data format.
 
 Command to get the data:
 
@@ -29,25 +29,36 @@
 Dependencies:
     None
 
-References:
+Resources:
     https://github.com/projectdiscovery/nuclei-templates/blob/main/http/technologies/kubernetes/kubelet/kubelet-metrics.yaml
     https://github.com/projectdiscovery/nuclei-templates/blob/main/http/exposures/configs/prometheus-metrics.yaml
 """
-K8S_RESOURCE_TYPES = ["pod", "deployment", "service", "configmap", "node", "secret"]
-K8S_OTHERS_TYPES = ["controller", "host", "webhook"]
+K8S_RESOURCE_TYPES = ["pod", "deployment", "service", "configmap", "node", "secret", "replicaset", "daemonset", "networkpolicy", "statefulset"]
+K8S_OTHERS_TYPES = ["controller", "host", "webhook", "lease", "pod_cidr"]
+K8S_RESOURCE_TYPES.sort()
+K8S_OTHERS_TYPES.sort()
 NAMESPACE_REGEX = r'namespace="(.*?)"'
 RESOURCE_KEY = "resources"
 OTHER_KEY = "others"
 MISC_KEY = "misc"
 
 
+def is_name_not_none(name):
+    return (name.lower().strip(" \t") != "<none>")
+
+
 def load_records(kms_data_file):
     # KEY is the namespace
     # VALUE is a DICT for which KEY is the K8S resource type and VALUE is a list of the related K8S resource/component names
     records = {}
     with open(kms_data_file, mode="r", encoding="utf-8") as f:
         lines = f.read().splitlines()
+    lines_count = len(lines)
+    line_id = 0
     for line in lines:
+        line_id += 1
+        loading_progress = round((line_id / lines_count) * 100)
+        print(f"\rExtracting data: {str(loading_progress).zfill(2)}%", end="", flush=True)
         # Skip comment
         if line.startswith("#"):
             continue
@@ -64,31 +75,35 @@ def load_records(kms_data_file):
             results = re.findall(resource_type_regex, line, re.IGNORECASE)
             if len(results) > 0:
                 resource_name = results[0]
-                if resource_type not in records[k8s_ns][RESOURCE_KEY]:
-                    records[k8s_ns][RESOURCE_KEY][resource_type] = []
-                if resource_name not in records[k8s_ns][RESOURCE_KEY][resource_type]:
-                    records[k8s_ns][RESOURCE_KEY][resource_type].append(resource_name)
+                if is_name_not_none(resource_name):
+                    if resource_type not in records[k8s_ns][RESOURCE_KEY]:
+                        records[k8s_ns][RESOURCE_KEY][resource_type] = []
+                    if resource_name not in records[k8s_ns][RESOURCE_KEY][resource_type]:
+                        records[k8s_ns][RESOURCE_KEY][resource_type].append(resource_name)
         # Extract the other types of components
         for component_type in K8S_OTHERS_TYPES:
             component_type_regex = f"{component_type}=\"(.*?)\""
             results = re.findall(component_type_regex, line, re.IGNORECASE)
             if len(results) > 0:
                 component_name = results[0]
-                if component_type not in records[k8s_ns][OTHER_KEY]:
-                    records[k8s_ns][OTHER_KEY][component_type] = []
-                if component_name not in records[k8s_ns][OTHER_KEY][component_type]:
-                    records[k8s_ns][OTHER_KEY][component_type].append(component_name)
+                if is_name_not_none(component_name):
+                    if component_type not in records[k8s_ns][OTHER_KEY]:
+                        records[k8s_ns][OTHER_KEY][component_type] = []
+                    if component_name not in records[k8s_ns][OTHER_KEY][component_type]:
+                        records[k8s_ns][OTHER_KEY][component_type].append(component_name)
         # Extract special item like this one:
         # gotk_reconcile_duration_seconds_bucket{kind="Kustomization",name="workspace",namespace="kommander",le="2"}
-        if 'kind=' in line and 'name=' in line:
+        if "kind=" in line and "name=" in line:
             kind_regex = "kind=\"(.*?)\""
             kind_name = re.findall(kind_regex, line, re.IGNORECASE)[0]
-            name_regex = "name=\"(.*?)\""
-            name_name = re.findall(name_regex, line, re.IGNORECASE)[0]
-            if kind_name not in records[k8s_ns][MISC_KEY]:
-                records[k8s_ns][MISC_KEY][kind_name] = []
-            if name_name not in records[k8s_ns][MISC_KEY][kind_name]:
-                records[k8s_ns][MISC_KEY][kind_name].append(name_name)
+            if kind_name.lower() not in (K8S_RESOURCE_TYPES + K8S_OTHERS_TYPES) and is_name_not_none(kind_name):
+                name_regex = "name=\"(.*?)\""
+                name_name = re.findall(name_regex, line, re.IGNORECASE)[0]
+                if kind_name not in records[k8s_ns][MISC_KEY]:
+                    records[k8s_ns][MISC_KEY][kind_name] = []
+                if name_name not in records[k8s_ns][MISC_KEY][kind_name]:
+                    records[k8s_ns][MISC_KEY][kind_name].append(name_name)
+    print("\r", end="", flush=True)
     return collections.OrderedDict(sorted(records.items()))
 
 
@@ -100,12 +115,12 @@ def generate_tree_representation(records):
         for category in [RESOURCE_KEY, OTHER_KEY, MISC_KEY]:
             category_items = namespace_items[category]
             if len(category_items) > 0:
-                print(colored(f"=> {category.capitalize()}", "light_cyan"))
+                print(colored(f"=> {category.capitalize()}", "light_magenta"))
                 for category_item_type in category_items:
                     item_collection = category_items[category_item_type]
                     item_count = len(item_collection)
                     item_collection.sort()
-                    print(colored(f"==> {category_item_type.capitalize()} ({item_count}): ", "light_magenta"), end="")
+                    print(colored(f"==> {category_item_type.capitalize()} ({item_count}): ", "light_cyan"), end="")
                     print(", ".join(item_collection))