From 4879d0f3fbdd245290b7d1394862b0c15561da6b Mon Sep 17 00:00:00 2001
From: Gregor Kleen <gkleen@yggdrasil.li>
Date: Sat, 11 Mar 2023 12:42:44 +0100
Subject: cake-prometheus-exporter

---
 .../cake-prometheus-exporter.py                    | 124 +++++++++++++++++++++
 overlays/cake-prometheus-exporter/default.nix      |  29 +++++
 2 files changed, 153 insertions(+)
 create mode 100644 overlays/cake-prometheus-exporter/cake-prometheus-exporter.py
 create mode 100644 overlays/cake-prometheus-exporter/default.nix

diff --git a/overlays/cake-prometheus-exporter/cake-prometheus-exporter.py b/overlays/cake-prometheus-exporter/cake-prometheus-exporter.py
new file mode 100644
index 00000000..07197149
--- /dev/null
+++ b/overlays/cake-prometheus-exporter/cake-prometheus-exporter.py
@@ -0,0 +1,124 @@
+#!@python@/bin/python
+
+from os import environ
+import sys
+
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+import subprocess
+import json
+
+from urllib.parse import urlparse
+
+
+def _format_prom_attrs(**attrs):
+    if not attrs:
+        return ''
+
+    return '{' + ','.join(map(lambda k: f'{k}="{attrs[k]}"', attrs)) + '}'
+
+def _format_prom_metrics(metricName, metricType, metrics, metricHelp=''):
+    metricStr = dedent(f'''
+      # HELP {metricName} {metricHelp}
+      # TYPE {metricName} {metricType}
+    ''').lstrip()
+    for (attrs, val) in metrics:
+        attrs_str = _format_prom_attrs(**attrs)
+        metricStr += dedent(f'''
+            {metricName}{attrs_str} {val}
+        ''').lstrip()
+    return metricStr
+
+
+class CAKEMetrics:
+    _instance = None
+
+    @classmethod
+    def instance(cls):
+        if cls._instance is None:
+            cls._instance = cls.__new__(cls)
+            cls._instance.attrs = None
+        return cls._instance
+
+    def __init__(self):
+        raise RuntimeError('Call instance() instead')
+
+    def update(self):
+        attrs = dict()
+
+        tc_output = None
+        with subprocess.Popen(['tc', '-s', '-j', 'qdisc', 'show'], stdout=subprocess.PIPE) as proc:
+            tc_output = json.load(proc.stdout)
+
+        for qdisc in tc_output:
+            if 'kind' not in qdisc or qdisc['kind'] != 'cake':
+                continue
+
+            tin_names = []
+            if len(qdisc['tins']) == 4:
+                tin_names = ['bulk', 'best-effort', 'video', 'voice']
+            tins = {}
+            for tin_name, tin_data in zip(tin_names, qdics['tins']):
+                tins[tin_name] = {
+                    'bytes': tin_data['sent_bytes'],
+                    'packets': tin_data['sent_packets']
+                }
+
+            attrs[qdisc['dev']] = {
+                'bytes': qdisc['bytes'],
+                'packets': qdisc['packets'],
+                'tins': tins
+            }
+
+        self.attrs = attrs
+
+    def json_text(self):
+        return json.dumps(self.attrs)
+
+    def prometheus(self):
+        metrics = ''
+
+        metrics += _format_prom_metrics('cake_bytes', 'counter', [({'dev': dev}, self.attrs[dev]['bytes']) for dev in self.attrs])
+        metrics += _format_prom_metrics('cake_packets', 'counter', [({'dev': dev}, self.attrs[dev]['packets']) for dev in self.attrs])
+
+        metrics += _format_prom_metrics('cake_tin_bytes', 'counter', [({'dev': dev, 'tin': tin}, self.attrs[dev]['tins'][tin]['bytes']) for tin in self.attrs[dev]['tins'] for dev in self.attrs])
+        metrics += _format_prom_metrics('cake_tin_packets', 'counter', [({'dev': dev, 'tin': tin}, self.attrs[dev]['tins'][tin]['packets']) for tin in self.attrs[dev]['tins'] for dev in self.attrs])
+
+        return metrics.encode('utf-8')
+
+
+class CAKEMetricsServer(BaseHTTPRequestHandler):
+    def log_message(self, format, *args):
+        pass
+
+    def do_GET(self):
+        cake_metrics = CAKEMetrics.instance()
+        cake_metrics.update()
+
+        url = urlparse(self.path)
+
+        match url.path:
+            case '/metrics.json':
+                self.send_response(200)
+                self.send_header("Content-type", "application/json")
+                self.end_headers()
+
+                self.wfile.write(cake_metrics.json_text().encode('utf-8'))
+            case '/metrics':
+                self.send_response(200)
+                self.send_header("Content-type", "text/plain")
+                self.end_headers()
+
+                self.wfile.write(cake_metrics.prometheus())
+            case _:
+                self.send_response(404)
+                self.end_headers()
+
+
+def main():
+    webServer = HTTPServer((str(environ.get('CAKE_HOSTNAME')), int(environ.get('CAKE_PORT'))), CAKEMetricsServer)
+
+    webServer.serve_forever()
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/overlays/cake-prometheus-exporter/default.nix b/overlays/cake-prometheus-exporter/default.nix
new file mode 100644
index 00000000..3d0acc2d
--- /dev/null
+++ b/overlays/cake-prometheus-exporter/default.nix
@@ -0,0 +1,29 @@
+{ final, prev, ... }:
+let
+  inpPython = final.python310.override {};
+in {
+  cake-prometheus-exporter = prev.stdenv.mkDerivation rec {
+    pname = "cake-prometheus-exporter";
+    version = "0.0.0";
+
+    src = ./cake-prometheus-exporter.py;
+
+    phases = [ "buildPhase" "checkPhase" "installPhase" ];
+
+    python = inpPython.withPackages (ps: with ps; []);
+
+    buildPhase = ''
+      substituteAll $src cake-prometheus-exporter
+    '';
+
+    doCheck = true;
+    checkPhase = ''
+      ${python}/bin/python -m py_compile cake-prometheus-exporter
+    '';
+
+    installPhase = ''
+      install -m 0755 -D -t $out/bin \
+        cake-prometheus-exporter
+    '';
+  };
+}
-- 
cgit v1.2.3