diff --git a/proxy/common/constants.py b/proxy/common/constants.py index 673f9a903c..39e52188b7 100644 --- a/proxy/common/constants.py +++ b/proxy/common/constants.py @@ -105,6 +105,7 @@ def _env_threadless_compliant() -> bool: DEFAULT_ENABLE_SSH_TUNNEL = False DEFAULT_ENABLE_DEVTOOLS = False DEFAULT_ENABLE_EVENTS = False +DEFAULT_ENABLE_METRICS = False DEFAULT_EVENTS_QUEUE = None DEFAULT_ENABLE_STATIC_SERVER = False DEFAULT_ENABLE_WEB_SERVER = False diff --git a/proxy/core/event/manager.py b/proxy/core/event/manager.py index 6a100f904a..21c662bf4a 100644 --- a/proxy/core/event/manager.py +++ b/proxy/core/event/manager.py @@ -20,7 +20,7 @@ from .queue import EventQueue from .dispatcher import EventDispatcher from ...common.flag import flags -from ...common.constants import DEFAULT_ENABLE_EVENTS +from ...common.constants import DEFAULT_ENABLE_EVENTS, DEFAULT_ENABLE_METRICS logger = logging.getLogger(__name__) @@ -33,6 +33,13 @@ help='Default: False. Enables core to dispatch lifecycle events. ' 'Plugins can be used to subscribe for core events.', ) +flags.add_argument( + '--enable-metrics', + action='store_true', + default=DEFAULT_ENABLE_METRICS, + help='Default: False. Enables core to dispatch metrics. ' + 'Plugins can be used to subscribe for metrics.', +) class EventManager: diff --git a/proxy/core/event/names.py b/proxy/core/event/names.py index 369724aac4..de4a01f31b 100644 --- a/proxy/core/event/names.py +++ b/proxy/core/event/names.py @@ -28,6 +28,7 @@ ('RESPONSE_HEADERS_COMPLETE', int), ('RESPONSE_CHUNK_RECEIVED', int), ('RESPONSE_COMPLETE', int), + ('METRIC', int), ], ) -eventNames = EventNames(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11) +eventNames = EventNames(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) diff --git a/proxy/http/metric_emisor.py b/proxy/http/metric_emisor.py new file mode 100644 index 0000000000..aec67896ea --- /dev/null +++ b/proxy/http/metric_emisor.py @@ -0,0 +1,59 @@ +# -*- coding: utf-8 -*- +""" + proxy.py + ~~~~~~~~ + ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on + Network monitoring, controls & Application development, testing, debugging. + + :copyright: (c) 2013-present by Abhinav Singh and contributors. + :license: BSD, see LICENSE for more details. +""" +import logging +from typing import Set, Union, Optional + +from proxy import metrics +from proxy.core.event import eventNames + + +logger = logging.getLogger(__name__) + + +class MetricEmisorMixin: + """MetricEmisorMixin provides methods to publish metrics.""" + + def _can_emit_metrics(self): + if self.flags.enable_events and self.flags.enable_metrics: + assert self.event_queue + return True + logging.info('Metrics disabled') + return False + + + def emit_metric(self, metric: metrics.Metric) -> None: + if self._can_emit_metrics(): + self.event_queue.publish( + request_id=self.uid, + event_name=eventNames.METRIC, + event_payload=metric, + publisher_id=self.__class__.__qualname__, + ) + + def emit_metric_counter( + self, + name: str, + increment: int | float=1, + description: Optional[str]=None, + tags: Optional[set[str]]=None, + ) -> None: + if self._can_emit_metrics(): + self.emit_metric(metrics.Counter(name, increment, description, tags)) + + def emit_metric_gauge( + self, + name: str, + value: Union[int, float], + description: Optional[str]=None, + tags: Optional[Set[str]]=None, + ) -> None: + if self._can_emit_metrics(): + self.emit_metric(metrics.Gauge(name, value, description, tags)) diff --git a/proxy/http/plugin.py b/proxy/http/plugin.py index 754fb28024..22bd5b4dc6 100644 --- a/proxy/http/plugin.py +++ b/proxy/http/plugin.py @@ -11,13 +11,14 @@ import socket import argparse from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List, Union, Optional +from typing import TYPE_CHECKING, Set, List, Union, Optional from .parser import HttpParser from .connection import HttpClientConnection from ..core.event import EventQueue from .descriptors import DescriptorsHandlerMixin from ..common.utils import tls_interception_enabled +from .metric_emisor import MetricEmisorMixin if TYPE_CHECKING: # pragma: no cover @@ -26,6 +27,7 @@ class HttpProtocolHandlerPlugin( DescriptorsHandlerMixin, + MetricEmisorMixin, ABC, ): """Base HttpProtocolHandler Plugin class. diff --git a/proxy/http/server/plugin.py b/proxy/http/server/plugin.py index d3536dd10b..f65b030a3b 100644 --- a/proxy/http/server/plugin.py +++ b/proxy/http/server/plugin.py @@ -22,6 +22,7 @@ from ..descriptors import DescriptorsHandlerMixin from ...common.types import RePattern from ...common.utils import bytes_ +from ..metric_emisor import MetricEmisorMixin from ...http.server.protocols import httpProtocolTypes @@ -29,7 +30,11 @@ from ...core.connection import TcpServerConnection, UpstreamConnectionPool -class HttpWebServerBasePlugin(DescriptorsHandlerMixin, ABC): +class HttpWebServerBasePlugin( + DescriptorsHandlerMixin, + MetricEmisorMixin, + ABC, +): """Web Server Plugin for routing of requests.""" def __init__( diff --git a/proxy/metrics/__init__.py b/proxy/metrics/__init__.py new file mode 100644 index 0000000000..39d22c5d6a --- /dev/null +++ b/proxy/metrics/__init__.py @@ -0,0 +1,49 @@ +# -*- coding: utf-8 -*- +""" + proxy.py + ~~~~~~~~ + ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on + Network monitoring, controls & Application development, testing, debugging. + + :copyright: (c) 2013-present by Abhinav Singh and contributors. + :license: BSD, see LICENSE for more details. +""" +from typing import Set, Union, Optional +from datetime import datetime + + +class Metric: + def __init__( + self, + name: str, + description: Optional[str]=None, + tags:Set[str] = None, + ): + self.timestamp = datetime.utcnow().timestamp() + self.name = name + self.description = description + self.tags = tags + + +class Counter(Metric): + def __init__( + self, + name:str, + increment: Union[int, float]=1, + description: Optional[str]=None, + tags:Set[str] = None, + ): + super().__init__(name, description) + self.increment = increment + + +class Gauge(Metric): + def __init__( + self, + name:str, + value: Union[int, float]=1, + description: Optional[str]=None, + tags:Set[str] = None, + ): + super().__init__(name, description) + self.value = value diff --git a/proxy/plugin/prometheus.py b/proxy/plugin/prometheus.py new file mode 100644 index 0000000000..6aba944203 --- /dev/null +++ b/proxy/plugin/prometheus.py @@ -0,0 +1,73 @@ +# -*- coding: utf-8 -*- +""" + proxy.py + ~~~~~~~~ + ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on + Network monitoring, controls & Application development, testing, debugging. + + :copyright: (c) 2013-present by Abhinav Singh and contributors. + :license: BSD, see LICENSE for more details. + + .. spelling:: + + ws + onmessage +""" +import logging +from typing import Any, Dict, List, Tuple + +from prometheus_client import Counter +from prometheus_client.registry import REGISTRY +from prometheus_client.exposition import generate_latest + +from proxy import metrics +from proxy.core.event import EventSubscriber +from proxy.http.parser import HttpParser +from proxy.http.server import HttpWebServerBasePlugin, httpProtocolTypes +from proxy.http.responses import okResponse + + +logger = logging.getLogger(__name__) + + +class Prometheus(HttpWebServerBasePlugin): + """ + Expose metrics on prometheus format. + Requires to install prometheus client (`pip install prometheus-client`) + """ + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.subscriber = EventSubscriber( + self.event_queue, + callback=self.process_metric_event, + ) + self.subscriber.setup() + self.metrics: Dict[str, Any] = {} + + def routes(self) -> List[Tuple[int, str]]: + return [ + (httpProtocolTypes.HTTP, r'/metrics$'), + (httpProtocolTypes.HTTPS, r'/metrics$'), + ] + + def handle_request(self, request: HttpParser) -> None: + self.emit_metric_counter('prometheus_requests') + if request.path == b'/metrics': + self.client.queue(okResponse(generate_latest())) + + def process_metric_event(self, event: Dict[str, Any]) -> None: + payload = event['event_payload'] + if not isinstance(payload, metrics.Metric): + return + try: + logger.info(event) + if isinstance(payload, metrics.Counter): + name = f"counter_{payload.name}" + + if name not in REGISTRY._names_to_collectors: + Counter(name, payload.description) + counter = REGISTRY._names_to_collectors.get(name) + counter.inc(payload.increment) + + except: + logger.exception('Problems')