Basic Prometheus metrics for API requests (#908)

Co-authored-by: Michał Praszmo <michalpr@cert.pl>

docker-compose-dev.yml

@@ -31,6 +31,7 @@ services:
       MWDB_RECAPTCHA_SITE_KEY: "6LeIxAcTAAAAAJcZVRqyHh71UMIEGNQ_MXjiZKhI"
       MWDB_RECAPTCHA_SECRET: "6LeIxAcTAAAAAGG-vFI1TnRWxMZNFuojJ4WifJWe"
       MWDB_ENABLE_REGISTRATION: 1
+      MWDB_ENABLE_PROMETHEUS_METRICS: 1
       # Uncomment if you want to test S3 functions
       # MWDB_STORAGE_PROVIDER: s3
       # MWDB_HASH_PATHING: 0

mwdb/app.py

@@ -9,6 +9,7 @@
 from mwdb.core.app import api, app
 from mwdb.core.config import app_config
 from mwdb.core.log import getLogger, setup_logger
+from mwdb.core.metrics import metric_api_requests, metrics_enabled
 from mwdb.core.plugins import PluginAppContext, load_plugins
 from mwdb.core.static import static_blueprint
 from mwdb.core.util import token_hex
@@ -55,6 +56,7 @@
     MetakeyPermissionResource,
     MetakeyResource,
 )
+from mwdb.resources.metrics import MetricsResource
 from mwdb.resources.oauth import (
     OpenIDAccountIdentitiesResource,
     OpenIDAuthenticateResource,
@@ -156,6 +158,15 @@ def log_request(response):
         response_time = None
     response_size = response.calculate_content_length()
 
+    if metrics_enabled():
+        user = g.auth_user.login if g.auth_user else request.remote_addr
+        metric_api_requests.inc(
+            method=request.method,
+            endpoint=request.endpoint,
+            user=str(user),
+            status_code=str(response.status_code),
+        )
+
     getLogger().debug(
         "request",
         extra={
@@ -387,6 +398,9 @@ def require_auth():
     RemoteTextBlobPushResource, "/remote/<remote_name>/push/blob/<hash64:identifier>"
 )
 
+if metrics_enabled():
+    api.add_resource(MetricsResource, "/varz")
+
 setup_logger()
 
 # Load plugins

mwdb/core/capabilities.py

@@ -49,6 +49,8 @@ class Capabilities(object):
     karton_unassign = "karton_unassign"
     # Can mark object as shareable with 3rd parties
     modify_3rd_party_sharing = "modify_3rd_party_sharing"
+    # Can access Prometheus metrics
+    access_prometheus_metrics = "access_prometheus_metrics"
 
     @classmethod
     def all(cls):

mwdb/core/config.py

@@ -132,6 +132,7 @@ class MWDBConfig(Config):
 
     enable_json_logger = key(cast=intbool, required=False, default=False)
     enable_sql_profiler = key(cast=intbool, required=False, default=False)
+    enable_prometheus_metrics = key(cast=intbool, required=False, default=False)
     log_only_slow_sql = key(cast=intbool, required=False, default=False)
     use_x_forwarded_for = key(cast=intbool, required=False, default=False)
     enable_debug_log = key(cast=intbool, required=False, default=False)

mwdb/core/metrics/__init__.py

@@ -0,0 +1,11 @@
+from .redis_counter import RedisCounter, collect, metrics_enabled
+
+metric_api_requests = RedisCounter(
+    "api_requests", "API request metrics", ("method", "endpoint", "user", "status_code")
+)
+
+__all__ = [
+    "collect",
+    "metrics_enabled",
+    "metric_api_requests",
+]

mwdb/core/metrics/redis_counter.py

@@ -0,0 +1,87 @@
+from typing import Dict, Tuple
+
+from prometheus_client import CollectorRegistry
+from prometheus_client import Gauge as PrometheusGauge
+from prometheus_client import generate_latest
+from redis import Redis
+
+from mwdb.core.config import app_config
+from mwdb.core.log import getLogger
+
+METRIC_REGISTRY = CollectorRegistry()
+METRIC_EXPIRATION_TIME = 1 * 60 * 60
+COUNTERS = []
+
+logger = getLogger()
+
+redis = None
+if app_config.mwdb.enable_prometheus_metrics:
+    if not app_config.mwdb.redis_uri:
+        logger.warning(
+            "metrics: Prometheus metrics are disabled because redis_uri is not set"
+        )
+    else:
+        redis = Redis.from_url(app_config.mwdb.redis_uri, decode_responses=True)
+
+
+class RedisCounter:
+    KEY_PREFIX = "METRICS"
+
+    def __init__(self, name: str, documentation: str, labelnames: Tuple[str] = ()):
+        self._gauge = PrometheusGauge(
+            name, documentation, labelnames, registry=METRIC_REGISTRY
+        )
+        self._name = name
+        self._labelnames = labelnames
+        COUNTERS.append(self)
+
+    def inc(self, amount: int = 1, **labelkwargs) -> None:
+        if not redis:
+            return
+        redis_key = self._key_from_labelkwargs(**labelkwargs)
+        p = redis.pipeline()
+        p.incr(redis_key, amount)
+        p.expire(redis_key, METRIC_EXPIRATION_TIME)
+        p.execute()
+
+    def _key_from_labelkwargs(self, **labelkwargs):
+        elements = [labelkwargs[name].replace(":", "_") for name in self._labelnames]
+        return ":".join([self.KEY_PREFIX, self._name] + elements)
+
+    def _labelkwargs_from_key(self, key):
+        parts = key.split(":")
+        name, *elements = parts[1:]
+        if name != self._name:
+            return None
+        if len(elements) != len(self._labelnames):
+            logger.warning(
+                f"metrics: Got {len(elements)} label parts "
+                f"but expected {len(self._labelnames)}"
+            )
+            return None
+        return {self._labelnames[idx]: value for idx, value in enumerate(elements)}
+
+    def update_counters(self, counters: Dict[str, int]):
+        """
+        This method updates counters based on
+        """
+        self._gauge.clear()
+
+        for key, value in counters.items():
+            labelkwargs = self._labelkwargs_from_key(key)
+            if labelkwargs is None:
+                return
+            self._gauge.labels(**labelkwargs).set(value)
+
+
+def collect():
+    keys = redis.keys(f"{RedisCounter.KEY_PREFIX}:*")
+    values = redis.mget(keys)
+    counters = dict(zip(keys, values))
+    for counter in COUNTERS:
+        counter.update_counters(counters)
+    return generate_latest(METRIC_REGISTRY)
+
+
+def metrics_enabled():
+    return redis is not None

mwdb/resources/metrics.py

@@ -0,0 +1,36 @@
+from flask import Response
+from flask_restful import Resource
+
+from mwdb.core.capabilities import Capabilities
+from mwdb.core.metrics import collect
+
+from . import requires_authorization, requires_capabilities
+
+
+class MetricsResource(Resource):
+    @requires_authorization
+    @requires_capabilities(Capabilities.access_prometheus_metrics)
+    def get(self):
+        """
+        ---
+        summary: Get Prometheus metrics
+        description: |
+            Returns metrics for Prometheus.
+
+            Requires 'access_prometheus_metrics' privilege
+        security:
+            - bearerAuth: []
+        tags:
+            - metrics
+        responses:
+            200:
+                description: Metrics in Prometheus format
+            403:
+                description: |
+                    User don't have 'access_prometheus_metrics' privilege.
+        """
+        metrics = collect()
+        return Response(
+            metrics,
+            content_type="application/octet-stream",
+        )

mwdb/web/src/commons/auth/capabilities.tsx

@@ -35,6 +35,8 @@ export let capabilitiesList: Record<Capability, string> = {
     [Capability.removingKarton]: "Can remove analysis from object",
     [Capability.modify3rdPartySharing]:
         "Can mark objects as shareable with 3rd parties",
+    [Capability.accessPrometheusMetrics]:
+        "Can access Prometheus metrics on /api/varz endpoint",
 };
 
 afterPluginsLoaded(() => {

mwdb/web/src/types/types.ts

@@ -26,6 +26,7 @@ export enum Capability {
     kartonReanalyze = "karton_reanalyze",
     removingKarton = "karton_unassign",
     modify3rdPartySharing = "modify_3rd_party_sharing",
+    accessPrometheusMetrics = "access_prometheus_metrics",
 }
 
 export type User = {

requirements.txt

@@ -29,3 +29,4 @@ python-dateutil==2.8.2
 pyzipper==0.3.5
 pycryptodomex==3.19.1
 ssdeep==3.4
+prometheus-client==0.19.0