From c13b8bd507ae63e5defbd4b8ca59a4534c2c13d0 Mon Sep 17 00:00:00 2001
From: Kumi <git@kumi.email>
Date: Wed, 5 Jun 2024 12:23:17 +0200
Subject: [PATCH] feat: add dynamic config for CryptPad exporter

Introduce support for configuring the CryptPad Prometheus exporter via environment variables and command line arguments. This enhancement moves away from the previous method that required direct modifications to the source code, offering a more flexible and user-friendly approach for setting the pins directory and server port. By enabling dynamic configuration, users can more easily adapt the exporter to various environments without altering the codebase. The update includes parsing arguments at runtime and prioritizing command line inputs over environment variables for greater control.

This change also encompasses improvements to code readability and structure, including better error handling and documentation within the exporter script.
---
 README.md   |  6 +++---
 exporter.py | 54 ++++++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index 41c12ad..85bbdf8 100644
--- a/README.md
+++ b/README.md
@@ -19,10 +19,10 @@ python exporter.py
 
 ## Configuration
 
-Currently, the exporter can only be configured by modifying the source code. The following variables can be changed:
+The exporter can be configured using environment variables or command line arguments.
 
-- `PINS_DIR`: The directory where the CryptPad pins are stored. Default: `/srv/cryptpad/data/pins`
-- `PORT`: The port the exporter listens on. Default: `8000`
+- `PINS_DIR`/`--pins-dir`: The directory where the CryptPad pins are stored. Default: `/srv/cryptpad/data/pins`
+- `PORT`/`--port`: The port the exporter listens on. Default: `8000`
 
 ## License
 
diff --git a/exporter.py b/exporter.py
index 13db130..9ba5684 100644
--- a/exporter.py
+++ b/exporter.py
@@ -1,22 +1,26 @@
 import os
+from argparse import ArgumentParser
 from prometheus_client import start_http_server, Gauge
 import time
 
 # Configuration (replace with your actual path)
-PINS_DIR = '/srv/cryptpad/data/pins/'
+PINS_DIR = "/srv/cryptpad/data/pins/"
 PORT = 8000
 
 # Prometheus metric
-registered_users_gauge = Gauge('cryptpad_registered_users', 'Number of registered users')
+registered_users_gauge = Gauge(
+    "cryptpad_registered_users", "Number of registered users"
+)
 
-def count_registered_users(pins_dir: os.PathLike) -> int:
-    """Returns the number of registered users by counting the number of folders in the user data directory
 
-    This emulates the way CryptPad counts registered users in the admin dashboard, see: 
+def count_registered_users(pins_dir: os.PathLike = PINS_DIR) -> int:
+    f"""Returns the number of registered users by counting the number of folders in the user data directory
+
+    This emulates the way CryptPad counts registered users in the admin dashboard, see:
     https://github.com/cryptpad/cryptpad/blob/7dbec1bd25b46e514ba82adad209961627410025/lib/commands/admin-rpc.js#L105
 
     Args:
-        pins_dir (os.PathLike): Path to the user pins directory (e.g. /srv/cryptpad/data/pins/)
+        pins_dir (os.PathLike): Path to the user pins directory (default: {PINS_DIR})
 
     Returns:
         int: Number of registered users
@@ -33,22 +37,46 @@ def count_registered_users(pins_dir: os.PathLike) -> int:
         print(f"Error counting registered users: {e}")
         return 0
 
-def update_metrics():
-    """Update the Prometheus metrics with the current number of registered users"""
-    user_count = count_registered_users(PINS_DIR)
+
+def update_metrics(pins_dir: os.PathLike = PINS_DIR):
+    f"""Update the Prometheus metrics with the current number of registered users
+
+    Args:
+        pins_dir (os.PathLike, optional): Path to the user pins directory (default: {PINS_DIR})
+    """
+    user_count = count_registered_users(pins_dir)
     registered_users_gauge.set(user_count)
     print(f"Updated registered users count: {user_count}")
 
+
 def main():
+    parser = ArgumentParser(description="CryptPad Prometheus exporter")
+    parser.add_argument(
+        "--pins-dir",
+        type=str,
+        default=PINS_DIR,
+        help=f"Path to the user pins directory (default: {PINS_DIR})",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=PORT,
+        help=f"Port to expose the Prometheus metrics (default: {PORT})",
+    )
+    args = parser.parse_args()
+
+    pins_dir = args.pins_dir or os.environ.get("PINS_DIR") or PINS_DIR
+    port = int(args.port or os.environ.get("PORT") or PORT)
+
     # Start up the server to expose the metrics.
-    start_http_server(PORT)
-    print("Prometheus exporter started on port " + PORT)
+    start_http_server(port)
+    print("Prometheus exporter started on port " + port)
 
     # Update metrics every 30 seconds
     while True:
-        update_metrics()
+        update_metrics(pins_dir)
         time.sleep(30)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()