insightconnect-plugin-runtime 6.2.6__py3-none-any.whl → 6.3.0__py3-none-any.whl

insightconnect_plugin_runtime/api/endpoints.py

@@ -787,7 +787,7 @@ class Endpoints:
         return version
 
     def add_plugin_custom_config(
-        self, input_data: Dict[str, Any], org_id: str
+            self, input_data: Dict[str, Any], org_id: str
     ) -> Dict[str, Any]:
         """
         Using the retrieved configs pulled from komand-props, pass the configuration that matches the requesting
@@ -815,7 +815,7 @@ class Endpoints:
             # This means we still need to manually delete the state for plugins on a per org basis.
             # This also means first time customers for their 'initial' lookup would get the lookback value passed in.
             if input_data.get("body", {}).get("state") and additional_config.get(
-                "lookback"
+                    "lookback"
             ):
                 self.logger.info(
                     "Found an existing plugin state, not passing lookback value..."
@@ -842,8 +842,8 @@ class Endpoints:
             if isinstance(wrapped_exception, ClientException):
                 status_code = 400
             elif (
-                isinstance(wrapped_exception, PluginException)
-                and wrapped_exception.preset is PluginException.Preset.BAD_REQUEST
+                    isinstance(wrapped_exception, PluginException)
+                    and wrapped_exception.preset is PluginException.Preset.BAD_REQUEST
             ):
                 status_code = 400
             elif isinstance(wrapped_exception, (ConnectionTestException, ClientException)):
@@ -868,8 +868,8 @@ class Endpoints:
         if isinstance(wrapped_exception, (ConnectionTestException, ClientException)):
             return 400
         elif (
-            isinstance(wrapped_exception, PluginException)
-            and wrapped_exception.preset is PluginException.Preset.BAD_REQUEST
+                isinstance(wrapped_exception, PluginException)
+                and wrapped_exception.preset is PluginException.Preset.BAD_REQUEST
         ):
             return 400
         elif isinstance(wrapped_exception, ServerException):

insightconnect_plugin_runtime/server.py

@@ -1,17 +1,16 @@
-import os
-import sys
 import json
 import logging
+import os
+import sys
+from time import sleep
 
+import gunicorn.app.base
+import structlog
 from apispec import APISpec
 from apispec.ext.marshmallow import MarshmallowPlugin
 from apispec_webframeworks.flask import FlaskPlugin
-
 from flask import Flask, request_started, request
-import gunicorn.app.base
 from gunicorn.arbiter import Arbiter
-
-import structlog
 from pythonjsonlogger.jsonlogger import JsonFormatter
 from requests import get as request_get
 from requests.exceptions import (
@@ -21,9 +20,9 @@ from requests.exceptions import (
     JSONDecodeError,
     ConnectionError,
 )
-from time import sleep
 from werkzeug.utils import secure_filename
 
+from insightconnect_plugin_runtime.api.endpoints import Endpoints, handle_errors
 from insightconnect_plugin_runtime.api.schemas import (
     PluginInfoSchema,
     ActionTriggerOutputBodySchema,
@@ -39,9 +38,9 @@ from insightconnect_plugin_runtime.api.schemas import (
     ConnectionDetailsSchema,
     ConnectionTestSchema,
 )
-from insightconnect_plugin_runtime.api.endpoints import Endpoints, handle_errors
-from insightconnect_plugin_runtime.util import is_running_in_cloud
 from insightconnect_plugin_runtime.helper import clean_dict
+from insightconnect_plugin_runtime.telemetry import create_post_fork
+from insightconnect_plugin_runtime.util import is_running_in_cloud, OTEL_ENDPOINT
 
 API_TITLE = "InsightConnect Plugin Runtime API"
 API_VERSION = "1.0"
@@ -72,14 +71,14 @@ class PluginServer(gunicorn.app.base.BaseApplication):
     """
 
     def __init__(
-        self,
-        plugin,
-        port=10001,
-        workers=1,
-        threads=4,
-        debug=False,
-        worker_class="sync",
-        worker_connections=200,
+            self,
+            plugin,
+            port=10001,
+            workers=1,
+            threads=4,
+            debug=False,
+            worker_class="sync",
+            worker_connections=200,
     ):
 
         gunicorn_file = os.environ.get("GUNICORN_CONFIG_FILE")
@@ -120,6 +119,9 @@ class PluginServer(gunicorn.app.base.BaseApplication):
         self.get_plugin_properties_from_cps()
         self.app, self.blueprints = self.create_flask_app()
 
+        if os.environ.get(OTEL_ENDPOINT):
+            self.config_options[OTEL_ENDPOINT] = os.environ[OTEL_ENDPOINT]
+
     @staticmethod
     def configure_structlog_instance(is_debug: bool) -> None:
         structlog.configure(
@@ -166,6 +168,9 @@ class PluginServer(gunicorn.app.base.BaseApplication):
         for key, value in config.items():
             self.cfg.set(key.lower(), value)
 
+        post_fork = create_post_fork(lambda: self.app, lambda: self.plugin, lambda: self.config_options)
+        self.cfg.set("post_fork", post_fork)
+
     def create_flask_app(self):
         app = Flask(__name__)
 
@@ -193,7 +198,7 @@ class PluginServer(gunicorn.app.base.BaseApplication):
 
     def get_plugin_properties_from_cps(self):
         # Call out to komand-props to get configurations related to only the plugin pod running.
-        if is_running_in_cloud() and self.plugin.tasks:
+        if is_running_in_cloud():
             for attempt in range(1, CPS_RETRY + 1):
                 self.logger.info(
                     f"Getting plugin configuration information... (attempt {attempt}/{CPS_RETRY})"
@@ -206,7 +211,12 @@ class PluginServer(gunicorn.app.base.BaseApplication):
                     )  # match how we name our images
                     plugin_config = resp_json.get("plugins", {}).get(plugin, {})
 
-                    self.config_options = plugin_config
+                    if self.plugin.tasks:
+                        self.config_options = plugin_config
+
+                    if resp_json.get(OTEL_ENDPOINT, {}):
+                        self.config_options[OTEL_ENDPOINT] = resp_json.get(OTEL_ENDPOINT, {})
+
                     self.logger.info("Plugin configuration successfully retrieved...")
                     return
                 except MissingSchema as missing_schema:

insightconnect_plugin_runtime/telemetry.py

@@ -0,0 +1,73 @@
+import functools
+from typing import Any, Callable
+from flask.app import Flask
+
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.flask import FlaskInstrumentor
+from opentelemetry.instrumentation.requests import RequestsInstrumentor
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.trace import Status, StatusCode
+
+from insightconnect_plugin_runtime.plugin import Plugin
+from insightconnect_plugin_runtime.util import is_running_in_cloud, OTEL_ENDPOINT
+
+
+def init_tracing(app: Flask, plugin: Plugin, endpoint: str) -> None:
+    """
+    Initialize OpenTelemetry Tracing
+
+    The function sets up the tracer provider, span processor and exporter with auto-instrumentation
+
+    :param app: The Flask Application
+    :param plugin: The plugin to derive the service name from
+    :param endpoint: The Otel Endpoint to emit traces to
+    """
+
+    if not is_running_in_cloud():
+        return
+
+    resource = Resource(attributes={"service.name": f'{plugin.name.lower().replace(" ", "_")}-{plugin.version}'})
+
+    trace_provider = TracerProvider(resource=resource)
+    exporter = OTLPSpanExporter(endpoint=endpoint)
+    trace_provider.add_span_processor(BatchSpanProcessor(exporter))
+    trace.set_tracer_provider(trace_provider)
+
+    FlaskInstrumentor().instrument_app(app)
+
+    def requests_callback(span: trace.Span, _: Any, response: Any) -> None:
+        if hasattr(response, "status_code"):
+            span.set_status(Status(StatusCode.OK if response.status_code < 400 else StatusCode.ERROR))
+
+    RequestsInstrumentor().instrument(trace_provider=trace_provider, response_hook=requests_callback)
+
+
+def auto_instrument(func: Callable) -> Callable:
+    """
+    Decorator that auto-instruments a function with a trace
+
+    :param func: function to instrument
+    :return:
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        tracer = trace.get_tracer(__name__)
+        with tracer.start_as_current_span(func.__name__):
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+def create_post_fork(app_getter: Callable, plugin_getter: Callable, config_getter: Callable) -> Callable:
+    def post_fork(server, worker):
+        app = app_getter()
+        plugin = plugin_getter()
+        endpoint = config_getter().get(OTEL_ENDPOINT, None)
+        if endpoint:
+            init_tracing(app, plugin, endpoint)
+
+    return post_fork

insightconnect_plugin_runtime/util.py

@@ -10,6 +10,7 @@ import python_jsonschema_objects as pjs
 KEYS_TO_CHECK_SECRETS = ("secretKey", "password", "key", "token")
 DEFAULT_REPLACE_THRESHOLD = 0.8
 DEFAULT_NUMBER_OF_ITERATIONS = 50
+OTEL_ENDPOINT = "otel_tracing"
 
 
 class OutputMasker:

{insightconnect_plugin_runtime-6.2.6.dist-info → insightconnect_plugin_runtime-6.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: insightconnect-plugin-runtime
-Version: 6.2.6
+Version: 6.3.0
 Summary: InsightConnect Plugin Runtime
 Home-page: https://github.com/rapid7/komand-plugin-sdk-python
 Author: Rapid7 Integrations Alliance
@@ -26,6 +26,11 @@ Requires-Dist: apispec-webframeworks==1.0.0
 Requires-Dist: blinker==1.9.0
 Requires-Dist: structlog==24.4.0
 Requires-Dist: python-json-logger==2.0.7
+Requires-Dist: Jinja2==3.1.6
+Requires-Dist: opentelemetry-sdk==1.31.1
+Requires-Dist: opentelemetry-instrumentation-flask==0.52b1
+Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.31.1
+Requires-Dist: opentelemetry-instrumentation-requests==0.52b1
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier
@@ -219,6 +224,7 @@ contributed. Black is installed as a test dependency and the hook can be initial
 after cloning this repository.
 
 ## Changelog
+* 6.3.0 - Add Tracing Instrumentation
 * 6.2.6 - Remove setuptools after installation
 * 6.2.5 - Fixed bug related to failure to set default region to assume role method in `aws_client` for newer versions of boto3 | Updated alpine image packages on build
 * 6.2.4 - Update `make_request` helper to support extra parameter of `max_response_size` to cap the response

{insightconnect_plugin_runtime-6.2.6.dist-info → insightconnect_plugin_runtime-6.3.0.dist-info}/RECORD

@@ -8,14 +8,15 @@ insightconnect_plugin_runtime/helper.py,sha256=B0XqAXmn8CT1KQ6i5IoWLQrQ_HVOvuKrI
 insightconnect_plugin_runtime/metrics.py,sha256=hf_Aoufip_s4k4o8Gtzz90ymZthkaT2e5sXh5B4LcF0,3186
 insightconnect_plugin_runtime/plugin.py,sha256=Yf4LNczykDVc31F9G8uuJ9gxEsgmxmAr0n4pcZzichM,26393
 insightconnect_plugin_runtime/schema.py,sha256=6MVw5hqGATU1VLgwfOWfPsP3hy1OnsugCTsgX8sknes,521
-insightconnect_plugin_runtime/server.py,sha256=09fxsbKf2ZZvSqRP2Bv9e9-fspDyEFR8_YgIFeMnXqQ,12578
+insightconnect_plugin_runtime/server.py,sha256=DHooHBQa1M2z7aETTWK6u9B2jChi8vONzqK4n_c94f4,13138
 insightconnect_plugin_runtime/step.py,sha256=KdERg-789-s99IEKN61DR08naz-YPxyinPT0C_T81C4,855
 insightconnect_plugin_runtime/task.py,sha256=d-H1EAzVnmSdDEJtXyIK5JySprxpF9cetVoFGtWlHrg,123
+insightconnect_plugin_runtime/telemetry.py,sha256=6JgZOyhB0qcDYKPlQZ47M2ZPfDpfAE_zuAhBtw4GMvo,2586
 insightconnect_plugin_runtime/trigger.py,sha256=Zq3cy68N3QxAGbNZKCID6CZF05Zi7YD2sdy_qbedUY8,874
-insightconnect_plugin_runtime/util.py,sha256=qPkZ3LA55nYuNYdansEbnCnBccQkpzIpp9NA1B64Kvw,8444
+insightconnect_plugin_runtime/util.py,sha256=8cle29INhnshEcL2LWpaC0ZGqevjq8pW8TE0MFEiYYw,8475
 insightconnect_plugin_runtime/variables.py,sha256=7FjJGnU7KUR7m9o-_tRq7Q3KiaB1Pp0Apj1NGgOwrJk,3056
 insightconnect_plugin_runtime/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-insightconnect_plugin_runtime/api/endpoints.py,sha256=8QQrxzW8jmQIkalud8fqYwB05uUw8sTiDNgO5ZekOCA,33353
+insightconnect_plugin_runtime/api/endpoints.py,sha256=Kn9VfnpvcVO9oY5W07laJkXFNp0PfZQrY2AecYFbvNw,33377
 insightconnect_plugin_runtime/api/schemas.py,sha256=jRmDrwLJTBl-iQOnyZkSwyJlCWg4eNjAnKfD9Eko4z0,2754
 insightconnect_plugin_runtime/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 insightconnect_plugin_runtime/clients/aws_client.py,sha256=ViJF3klbD1YM5GVxme3BFYfpuADUlVGP8sdbX2Z6_Cw,23061
@@ -73,12 +74,12 @@ tests/unit/test_metrics.py,sha256=PjjTrB9w7uQ2Q5UN-893-SsH3EGJuBseOMHSD1I004s,79
 tests/unit/test_oauth.py,sha256=nbFG0JH1x04ExXqSe-b5BGdt_hJs7DP17eUa6bQzcYI,2093
 tests/unit/test_plugin.py,sha256=ZTNAZWwZhDIAbxkVuWhnz9FzmojbijgMmsLWM2mXQI0,4160
 tests/unit/test_schema.py,sha256=swWZPRo_Q4M6VHte-srmxcV2wH-XS7pgmNRxpaL0Qrg,642
-tests/unit/test_server_cloud_plugins.py,sha256=PuMDHTz3af6lR9QK1BtPScr7_cRbWhetowADieVlXdo,5096
+tests/unit/test_server_cloud_plugins.py,sha256=3hEdNrJmnDXqhxhd4uSdTJ0ksn2S11RsNdZCmnethCw,5340
 tests/unit/test_server_spec.py,sha256=je97BaktgK0Fiz3AwFPkcmHzYtOJJNqJV_Fw5hrvqX4,644
 tests/unit/test_trigger.py,sha256=E53mAUoVyponWu_4IQZ0IC1gQ9lakBnTn_9vKN2IZfg,1692
 tests/unit/test_variables.py,sha256=OUEOqGYZA3Nd5oKk5GVY3hcrWKHpZpxysBJcO_v5gzs,291
 tests/unit/utils.py,sha256=hcY0A2H_DMgCDXUTvDtCXMdMvRjLQgTaGcTpATb8YG0,2236
-insightconnect_plugin_runtime-6.2.6.dist-info/METADATA,sha256=EAB_1cINrkNbkoHu5n6xbhtqrT6HwPbBonogpA_eqsE,15934
-insightconnect_plugin_runtime-6.2.6.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-insightconnect_plugin_runtime-6.2.6.dist-info/top_level.txt,sha256=AJtyJOpiFzHxsbHUICTcUKXyrGQ3tZxhrEHsPjJBvEA,36
-insightconnect_plugin_runtime-6.2.6.dist-info/RECORD,,
+insightconnect_plugin_runtime-6.3.0.dist-info/METADATA,sha256=GCmcMcuLriTNb5PMLZc83UqlpxByddurw6G-WAwui-4,16225
+insightconnect_plugin_runtime-6.3.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+insightconnect_plugin_runtime-6.3.0.dist-info/top_level.txt,sha256=AJtyJOpiFzHxsbHUICTcUKXyrGQ3tZxhrEHsPjJBvEA,36
+insightconnect_plugin_runtime-6.3.0.dist-info/RECORD,,

{insightconnect_plugin_runtime-6.2.6.dist-info → insightconnect_plugin_runtime-6.3.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (76.0.0)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tests/unit/test_server_cloud_plugins.py

@@ -4,6 +4,7 @@ from unittest import TestCase, skip
 from unittest.mock import patch, MagicMock
 
 from insightconnect_plugin_runtime.server import PluginServer
+from insightconnect_plugin_runtime.util import OTEL_ENDPOINT
 from tests.plugin.hello_world import KomandHelloWorld
 from .utils import MockResponse, Logger
 
@@ -20,15 +21,19 @@ class TestServerCloudPlugins(TestCase):
     @parameterized.expand([["Set cloud to false", False], ["Set cloud to true", True]])
     @patch("insightconnect_plugin_runtime.server.request_get")
     def test_cloud_plugin_no_tasks_ignore_cps(self, _test_name, cloud, mocked_req, mock_cloud, _run):
+        fake_endpoint = "http://fake.endpoint.com"
+        mocked_req.return_value = MockResponse({OTEL_ENDPOINT: fake_endpoint}) if cloud else MockResponse({})
         mock_cloud.return_value = cloud  # Mock plugin running in cloud vs not
         self.plugin.tasks = None  # ensure still no tasks as other tests edit this and could fail before reverting
-        plugin_server = PluginServer(self.plugin)  # this plugin has no tasks by default
 
+        plugin_server = PluginServer(self.plugin)  # this plugin has no tasks by default
         plugin_server.start()
-        self.assertEqual(plugin_server.config_options, {})
 
-        # Plugin server never calls out to CPS as either we are not running in cloud mode or have no tasks.
-        self.assertFalse(mocked_req.called)
+        self.assertEqual(plugin_server.config_options, {OTEL_ENDPOINT: fake_endpoint} if cloud else {})
+
+        # Plugin server calls out to CPS when cloud to get tracing endpoint
+        self.assertEqual(mocked_req.called, cloud)
+
 
     @patch("insightconnect_plugin_runtime.server.request_get")
     def test_cloud_plugin_calls_cps(self, mocked_req, _mock_cloud, _run):