insightconnect-plugin-runtime 6.3.10__tar.gz → 6.4.1__tar.gz

{insightconnect_plugin_runtime-6.3.10/insightconnect_plugin_runtime.egg-info → insightconnect_plugin_runtime-6.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: insightconnect-plugin-runtime
-Version: 6.3.10
+Version: 6.4.1
 Summary: InsightConnect Plugin Runtime
 Home-page: https://github.com/rapid7/komand-plugin-sdk-python
 Author: Rapid7 Integrations Alliance
@@ -199,11 +199,11 @@ Running a specific test file:
 
 |                   | Plugin  | Slim Plugin |
 |:------------------|:-------:|:-----------:|
-| Python Version    | 3.11.13 |   3.11.13   |
-| OS                | Alpine  |  Bullseye   |
+| Python Version    | 3.13.9  |   3.13.9    |
+| OS                | Alpine  |   Trixie    |
 | Package installer |   apk   |     apt     |
 | Shell             | /bin/sh |  /bin/bash  |
-| Image Size        | ~370MB  |   ~230MB    |
+| Image Size        | ~500MB  |   ~300MB    |
 
 Note that for the plugin image, we run `apk update` and `apk add ..` which leads to a longer build time.
 
@@ -226,6 +226,8 @@ contributed. Black is installed as a test dependency and the hook can be initial
 after cloning this repository.
 
 ## Changelog
+* 6.4.1 - Fix logging on cloud mode to remove `taskName` | Close boto client after assuming role
+* 6.4.0 - When running in cloud mode, always create a new connection object | Bump Python version to 3.13.9 | Docker build tidyup to reduce image size
 * 6.3.10 - Fixed tracing name to better allign otel standards
 * 6.3.9 - Fixed `monitor_task_delay` decorator handling of millisecond epoch timestamps | Allow `test_task` connection method to accept task name parameter | Updated base images to Python 3.11.13
 * 6.3.8 - Update exception string representation methods to remove newline characters

{insightconnect_plugin_runtime-6.3.10 → insightconnect_plugin_runtime-6.4.1}/README.md

@@ -155,11 +155,11 @@ Running a specific test file:
 
 |                   | Plugin  | Slim Plugin |
 |:------------------|:-------:|:-----------:|
-| Python Version    | 3.11.13 |   3.11.13   |
-| OS                | Alpine  |  Bullseye   |
+| Python Version    | 3.13.9  |   3.13.9    |
+| OS                | Alpine  |   Trixie    |
 | Package installer |   apk   |     apt     |
 | Shell             | /bin/sh |  /bin/bash  |
-| Image Size        | ~370MB  |   ~230MB    |
+| Image Size        | ~500MB  |   ~300MB    |
 
 Note that for the plugin image, we run `apk update` and `apk add ..` which leads to a longer build time.
 
@@ -182,6 +182,8 @@ contributed. Black is installed as a test dependency and the hook can be initial
 after cloning this repository.
 
 ## Changelog
+* 6.4.1 - Fix logging on cloud mode to remove `taskName` | Close boto client after assuming role
+* 6.4.0 - When running in cloud mode, always create a new connection object | Bump Python version to 3.13.9 | Docker build tidyup to reduce image size
 * 6.3.10 - Fixed tracing name to better allign otel standards
 * 6.3.9 - Fixed `monitor_task_delay` decorator handling of millisecond epoch timestamps | Allow `test_task` connection method to accept task name parameter | Updated base images to Python 3.11.13
 * 6.3.8 - Update exception string representation methods to remove newline characters

{insightconnect_plugin_runtime-6.3.10 → insightconnect_plugin_runtime-6.4.1}/insightconnect_plugin_runtime/clients/aws_client.py

@@ -19,6 +19,7 @@ from insightconnect_plugin_runtime.exceptions import (
     ConnectionTestException,
 )
 from insightconnect_plugin_runtime.helper import clean
+from insightconnect_plugin_runtime.util import is_running_in_cloud
 
 REGION = "region"
 EXTERNAL_ID = "external_id"
@@ -371,6 +372,7 @@ class AWSAction(Action):
         aws_service: str,
         aws_command: str,
         pagination_helper: PaginationHelper = None,
+        close_client: bool = None,
     ):
         """
 
@@ -382,6 +384,8 @@ class AWSAction(Action):
         :param output: The output schema object
         :param aws_service: The AWS service. Should be snake case.
         :param aws_command: The type of request to invoke. Should be snake case.
+        :param pagination_helper: Paginating helper indicate attrs like max_pages and token location.
+        :param close_client: Determine if the created client should be closed at the end of the action.
         """
 
         super().__init__(
@@ -391,6 +395,11 @@ class AWSAction(Action):
         self.aws_command = aws_command
         self.pagination_helper = pagination_helper
 
+        # when running in cloud mode we won't hold this connection, so each call to action.run() will spawn a new
+        # client, in this case we should close the client unless otherwise specified. A use case to not close
+        # is on a task plugin that re-uses the client for list_objects_v2 and get_bucket_content.
+        self.close_client = is_running_in_cloud() if close_client is None else close_client
+
     def _handle_botocore_function(
         self, client_function: Callable, params: Dict
     ) -> Dict:
@@ -441,6 +450,9 @@ class AWSAction(Action):
             raise PluginException(
                 cause="Error occurred when invoking the aws-cli.", data=error
             )
+        finally:
+            if self.close_client:
+                self.connection.client.close()
         return response
 
     def _handle_format_output(self, response: Dict, helper: ActionHelper) -> Dict:
@@ -542,6 +554,9 @@ class AWSAction(Action):
             if not self.pagination_helper.max_pages:
                 self.pagination_helper.remove_keys(response)
 
+        if self.close_client:
+            self.connection.client.close()
+
         return response
 
     def test(self, params={}):
@@ -610,14 +625,14 @@ class AWSAction(Action):
                     }
                 )
             )
+            sts_client.close()
         except ClientError as error:
             raise PluginException(
                 cause=f"Boto3 raised following error during assume role: {error.response['Error']['Code']}",
                 assistance="Please verify your role ARN and external ID are correct",
             )
         credentials = assumed_role_object["Credentials"]
-        session = botocore.session.Session()
-        boto_session = session.create_client(
+        boto_client = boto3.client(
             service_name,
             aws_access_key_id=credentials["AccessKeyId"],
             aws_secret_access_key=credentials["SecretAccessKey"],
@@ -625,4 +640,4 @@ class AWSAction(Action):
             region_name=assume_role_params[REGION],
         )
 
-        return boto_session
+        return boto_client

{insightconnect_plugin_runtime-6.3.10 → insightconnect_plugin_runtime-6.4.1}/insightconnect_plugin_runtime/connection.py

@@ -5,7 +5,7 @@ import json
 
 from jsonschema import validate
 
-from .util import sample as utilsample
+from .util import sample as utilsample, is_running_in_cloud
 
 
 def key(parameters):
@@ -26,28 +26,27 @@ def key(parameters):
 class ConnectionCache(object):
     def __init__(self, prototype):
         self.connections = {}
-        self.prototype = prototype
+        self.prototype = prototype  # connection JSON which does not contain a logger or validation on the values
 
     def get(self, parameters, logger):
-        conn_key = key(parameters)
-
-        # we could 'lock' this data  structure
-        # but honestly i  don't see the point
-        # worst case in a race condition is that
-        # we create 2 copies of the connection, and
-        # the later one overrides the first one.
-        #
-        # for real connection pooling we will need to fix this.
-        if conn_key in self.connections:
-            return self.connections[conn_key]
+        # when we're running in cloud mode we don't want to create and store connections to reduce the number
+        # of connection objects lying around in memory as this isn't safe and can cause OOM errors.
+        if is_running_in_cloud():
+            conn = self.create_and_validate_connection(parameters, logger)
+        else:
+            # first check if we have an existing connection obj we can return as it's been validated already
+            conn_key = key(parameters)
+            if not (conn := self.connections.get(conn_key)):  # otherwise create a new conn obj and save in cache
+                conn = self.create_and_validate_connection(parameters, logger)
+                self.connections[conn_key] = conn
+        return conn
 
+    def create_and_validate_connection(self, parameters, logger):
         conn = copy.copy(self.prototype)
         conn.logger = logger
         conn.set_(parameters)
-        # i don't know why this is needed twice..
-        # i think for backwards compat reasons
         conn.connect(parameters)
-        self.connections[conn_key] = conn
+
         return conn
 
 

{insightconnect_plugin_runtime-6.3.10 → insightconnect_plugin_runtime-6.4.1}/insightconnect_plugin_runtime/server.py

@@ -11,7 +11,7 @@ from apispec.ext.marshmallow import MarshmallowPlugin
 from apispec_webframeworks.flask import FlaskPlugin
 from flask import Flask, request_started, request
 from gunicorn.arbiter import Arbiter
-from pythonjsonlogger.jsonlogger import JsonFormatter
+from pythonjsonlogger.jsonlogger import JsonFormatter, RESERVED_ATTRS
 from requests import get as request_get
 from requests.exceptions import (
     HTTPError,
@@ -146,7 +146,8 @@ class PluginServer(gunicorn.app.base.BaseApplication):
         console_handler = logging.StreamHandler()
         if is_running_in_cloud():
             console_handler.setFormatter(
-                JsonFormatter()
+                # since python bump to 3.13 the logging lib now includes a bunch of extra keywords we want to discard
+                JsonFormatter(reserved_attrs=RESERVED_ATTRS + ("taskName",))
             )  # Only log in JSON if running in cloud
 
         logger.addHandler(console_handler)

{insightconnect_plugin_runtime-6.3.10 → insightconnect_plugin_runtime-6.4.1/insightconnect_plugin_runtime.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: insightconnect-plugin-runtime
-Version: 6.3.10
+Version: 6.4.1
 Summary: InsightConnect Plugin Runtime
 Home-page: https://github.com/rapid7/komand-plugin-sdk-python
 Author: Rapid7 Integrations Alliance
@@ -199,11 +199,11 @@ Running a specific test file:
 
 |                   | Plugin  | Slim Plugin |
 |:------------------|:-------:|:-----------:|
-| Python Version    | 3.11.13 |   3.11.13   |
-| OS                | Alpine  |  Bullseye   |
+| Python Version    | 3.13.9  |   3.13.9    |
+| OS                | Alpine  |   Trixie    |
 | Package installer |   apk   |     apt     |
 | Shell             | /bin/sh |  /bin/bash  |
-| Image Size        | ~370MB  |   ~230MB    |
+| Image Size        | ~500MB  |   ~300MB    |
 
 Note that for the plugin image, we run `apk update` and `apk add ..` which leads to a longer build time.
 
@@ -226,6 +226,8 @@ contributed. Black is installed as a test dependency and the hook can be initial
 after cloning this repository.
 
 ## Changelog
+* 6.4.1 - Fix logging on cloud mode to remove `taskName` | Close boto client after assuming role
+* 6.4.0 - When running in cloud mode, always create a new connection object | Bump Python version to 3.13.9 | Docker build tidyup to reduce image size
 * 6.3.10 - Fixed tracing name to better allign otel standards
 * 6.3.9 - Fixed `monitor_task_delay` decorator handling of millisecond epoch timestamps | Allow `test_task` connection method to accept task name parameter | Updated base images to Python 3.11.13
 * 6.3.8 - Update exception string representation methods to remove newline characters

{insightconnect_plugin_runtime-6.3.10 → insightconnect_plugin_runtime-6.4.1}/insightconnect_plugin_runtime.egg-info/SOURCES.txt

@@ -75,6 +75,7 @@ tests/unit/__init__.py
 tests/unit/test_action.py
 tests/unit/test_api.py
 tests/unit/test_aws_action.py
+tests/unit/test_connection.py
 tests/unit/test_custom_encoder.py
 tests/unit/test_endpoints.py
 tests/unit/test_exceptions.py

{insightconnect_plugin_runtime-6.3.10 → insightconnect_plugin_runtime-6.4.1}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setup(
     name="insightconnect-plugin-runtime",
-    version="6.3.10",
+    version="6.4.1",
     description="InsightConnect Plugin Runtime",
     long_description=long_description,
     long_description_content_type="text/markdown",

{insightconnect_plugin_runtime-6.3.10 → insightconnect_plugin_runtime-6.4.1}/tests/unit/test_aws_action.py

@@ -3,6 +3,7 @@ import io
 import json
 import unittest
 import unittest.mock
+from os import environ
 from pathlib import Path
 
 import botocore.exceptions as be
@@ -28,6 +29,10 @@ class Boto3Stub:
             }
         }
 
+    def close(self):
+        # stub for closing the sts client after assuming the role
+        pass
+
 
 class TestAwsAction(unittest.TestCase):
     def setUp(self) -> None:
@@ -37,6 +42,12 @@ class TestAwsAction(unittest.TestCase):
         }
         self.region = "us-east"
 
+        self.assume_role_params = {
+            "role_arn": "test_role",
+            "external_id": "test_id",
+            "region": "test-region",
+        }
+
         self.aws_action = AWSAction(
             "NewAction", "Description", None, None, "ec2", "service"
         )
@@ -45,19 +56,11 @@ class TestAwsAction(unittest.TestCase):
         self.aws_action.input = Input({})
         self.aws_action.output = Output({})
 
-    @unittest.mock.patch("botocore.session.Session", return_value=unittest.mock.Mock())
     @unittest.mock.patch("boto3.client", return_value=Boto3Stub())
-    def test_assume_role(self, mock_session, mock_sts_client):
-        assume_role_params = {
-            "role_arn": "test_role",
-            "external_id": "test_id",
-            "region": "test-region",
-        }
-        aws_session = unittest.mock.Mock()
-
-        AWSAction.try_to_assume_role("ec2", assume_role_params, self.auth_params)
+    def test_assume_role(self, mock_sts_client):
+        AWSAction.try_to_assume_role("ec2", self.assume_role_params, self.auth_params)
 
-        mock_sts_client.assert_called_once()
+        self.assertEqual(mock_sts_client.call_count, 2)  # twice for assume role and create client afterwards
 
     @unittest.mock.patch("botocore.session.Session", return_value=unittest.mock.Mock())
     @unittest.mock.patch("boto3.client", return_value=Boto3Stub())
@@ -135,6 +138,58 @@ class TestAwsAction(unittest.TestCase):
         with self.assertRaises(PluginException):
             self.aws_action.handle_rest_call(mock_call, {})
 
+    @unittest.mock.patch.dict(environ, {"PLUGIN_RUNTIME_ENVIRONMENT": "orchestrator"})
+    @unittest.mock.patch.object(AWSAction, "handle_rest_call", return_value={"mock_key": "mock_value"})
+    def test_run_client_non_cloud_mode_in_action_default_behaviour(self, _mock_handle_rest_call):
+        # Test AWSAction called for customers running on an orchestrator that their client can remain open
+        aws_action = AWSAction("NewAction", "Description", None, None, "s3", "service")
+        aws_action.connection = unittest.mock.create_autospec(Connection)
+        aws_action.connection.assume_role_params = self.assume_role_params
+        aws_action.connection.auth_params = self.auth_params
+        aws_action.connection.client = unittest.mock.create_autospec(Boto3Stub)
+        aws_action.connection.client.service = unittest.mock.MagicMock()
+        aws_action.run()
+        aws_action.connection.client.close.assert_not_called()
+
+    @unittest.mock.patch.dict(environ, {"PLUGIN_RUNTIME_ENVIRONMENT": "cloud"})
+    @unittest.mock.patch.object(AWSAction, "handle_rest_call", return_value={"mock_key": "mock_value"})
+    def test_run_client_cloud_mode_in_normal_action(self, _mock_handle_rest_call):
+        # Test the same AWSAction as above but now running in cloud - we should close this client
+        aws_action = AWSAction("NewAction", "Description", None, None, "s3", "service")
+        aws_action.connection = unittest.mock.create_autospec(Connection)
+        aws_action.connection.assume_role_params = self.assume_role_params
+        aws_action.connection.auth_params = self.auth_params
+        aws_action.connection.client = unittest.mock.create_autospec(Boto3Stub)
+        aws_action.connection.client.service = unittest.mock.MagicMock()
+        aws_action.run()
+        aws_action.connection.client.close.assert_called_once()
+
+    @unittest.mock.patch.dict(environ, {"PLUGIN_RUNTIME_ENVIRONMENT": "orchestrator"})
+    @unittest.mock.patch.object(AWSAction, "handle_rest_call", return_value={"mock_key": "mock_value"})
+    def test_run_client_non_cloud_mode_in_action_default_override(self, _mock_handle_rest_call):
+        # Test AWSAction when we have specified to close the client
+        aws_action = AWSAction("NewAction", "Description", None, None, "s3", "service", close_client=True)
+        aws_action.connection = unittest.mock.create_autospec(Connection)
+        aws_action.connection.assume_role_params = self.assume_role_params
+        aws_action.connection.auth_params = self.auth_params
+        aws_action.connection.client = unittest.mock.create_autospec(Boto3Stub)
+        aws_action.connection.client.service = unittest.mock.MagicMock()
+        aws_action.run()
+        aws_action.connection.client.close.assert_called_once()
+
+    @unittest.mock.patch.dict(environ, {"PLUGIN_RUNTIME_ENVIRONMENT": "cloud"})
+    @unittest.mock.patch.object(AWSAction, "handle_rest_call", return_value={"mock_key": "mock_value"})
+    def test_run_client_cloud_mode_in_typical_c2c_task(self, _mock_handle_rest_call):
+        # Test AWSAction used within a C2C task, we want to keep the client for subsequent calls
+        aws_action = AWSAction("NewAction", "Description", None, None, "s3", "service", close_client=False)
+        aws_action.connection = unittest.mock.create_autospec(Connection)
+        aws_action.connection.assume_role_params = self.assume_role_params
+        aws_action.connection.auth_params = self.auth_params
+        aws_action.connection.client = unittest.mock.create_autospec(Boto3Stub)
+        aws_action.connection.client.service = unittest.mock.MagicMock()
+        aws_action.run()
+        aws_action.connection.client.close.assert_not_called()
+
     def mocked_requests_get(*args, **kwargs):
         class MockResponse:
             def __init__(self, json_data, status_code):

insightconnect_plugin_runtime-6.4.1/tests/unit/test_connection.py

@@ -0,0 +1,113 @@
+from os import environ
+from unittest import TestCase
+from unittest.mock import patch, MagicMock
+
+from insightconnect_plugin_runtime.connection import Connection, ConnectionCache, key
+
+"""
+Any time an action / task / trigger API is called and we route through `plugin.py` -> `start_step`
+the request body, containing the connection JSON is parsed and we make use of the `ConnectionCache` to retrieve
+a connection that has already been validated or we create the new `Connection` object and call the `connect` method
+which should be implemented within the paryicular plugin. 
+
+These unit test file is testing that the validation, and caching logic works as expected both when running
+on an orchestrator or when running on the cloud. Cloud enabled plugins we do not want to re-use a previously
+validated connection as more customers use the plugin, more and more connection objects can be left on the pod
+which increases the memory usage and instead these should always be freshly created. 
+"""
+
+
+class TestCloudConnections(TestCase):
+    def setUp(self):
+        # Create a sample connection schema that will be validated
+        self.connection_schema = {
+            "type": "object",
+            "properties": {
+                "username": {"type": "string"},
+            },
+            "required": ["username"],
+        }
+        # these initialisations happen during plugin start up
+        self.connection = Connection(self.connection_schema)
+        self.connection_cache = ConnectionCache(self.connection)
+
+        self.connection = {
+            "username": "test_user@rapid7.com"
+        }
+
+        self.logger = MagicMock()
+
+    def stub_con(self, params):
+        # very basic example of what our connections do in plugins
+        self.logger.info("Connect: Connecting...")
+        self.username = params.get("username")
+
+    @patch.dict(environ, {"PLUGIN_RUNTIME_ENVIRONMENT": "cloud"})
+    @patch.object(Connection, "connect", new=stub_con)
+    def test_running_cloud_has_no_cache(self):
+        conn = self.connection_cache.get(self.connection, self.logger)
+        self.assertEqual(conn.parameters, self.connection)
+        self.assertDictEqual(self.connection_cache.connections, {})
+
+    @patch.dict(environ, {"PLUGIN_RUNTIME_ENVIRONMENT": "cloud"})
+    @patch.object(Connection, "connect", new=stub_con)
+    def test_running_cloud_has_no_cache_on_subsequent_run(self):
+        # call twice and we should still have no connection cache values
+        with patch("insightconnect_plugin_runtime.connection.key", MagicMock(side_effect=key)) as stub_key:
+            conn = self.connection_cache.get(self.connection, self.logger)
+            _ = self.connection_cache.get(self.connection, self.logger)
+        self.assertEqual(conn.parameters, self.connection)
+        self.assertDictEqual(self.connection_cache.connections, {})
+        self.assertEqual(stub_key.call_count, 0)  # shouldn't call this at all
+
+    @patch.object(Connection, "connect", new=stub_con)
+    def test_running_onprem_has_cached_connection(self):
+        conn = self.connection_cache.get(self.connection, self.logger)
+        self.assertEqual(conn.parameters, self.connection)
+        self.assertNotEqual(self.connection_cache.connections, {})
+
+        # the connection cache key should exist
+        hashed_key = key(conn.parameters)
+        self.assertIn(hashed_key, self.connection_cache.connections)
+        self.assertTrue(type(self.connection_cache.connections[hashed_key] == Connection))
+
+    @patch.object(ConnectionCache, "create_and_validate_connection")
+    def test_running_onprem_uses_cache_on_subsequent_calls(self, stub_create):
+        with patch("insightconnect_plugin_runtime.connection.key", MagicMock(side_effect=key)) as stub_key:
+            # Mock the return value of create_and_validate_connection
+            pretend_obj = 'for this unit test pretend this is a conn object'
+            stub_create.return_value = pretend_obj
+
+            # first call the connection cache for the first time adding a new entry
+            conn = self.connection_cache.get(self.connection, self.logger)
+            self.assertEqual(len(self.connection_cache.connections), 1)
+            self.assertEqual(list(self.connection_cache.connections.values())[0], pretend_obj)
+            self.assertEqual(conn, pretend_obj)
+            self.assertEqual(stub_key.call_count, 1)
+
+            # on a second call we get the same value but the create_and_validate shouldn't be called again
+            conn2 = self.connection_cache.get(self.connection, self.logger)
+            self.assertEqual(conn, conn2)
+            stub_create.assert_called_once_with(self.connection, self.logger)
+
+            # on a third call we add a new entry to the cached connections
+            new_conn = self.connection.copy()
+            new_conn["username"] = "test_user_2@rapid7.com"
+            _ = self.connection_cache.get(new_conn, self.logger)
+            self.assertEqual(len(self.connection_cache.connections), 2)  # new entry added
+
+            self.assertEqual(stub_key.call_count, 3)  # should have been called each time
+
+    def test_connection_not_implemented_raised(self):
+        with self.assertRaises(NotImplementedError):  # we haven't stubbed this so the error is expected
+            _conn = self.connection_cache.get(self.connection, self.logger)
+
+        # this means the connection object is not saved in cache
+        self.assertEqual(self.connection_cache.connections, {})
+
+    def test_connection_validation_raised(self):
+        with self.assertRaises(Exception):  # not the expected input to match the connection schema
+            _conn = self.connection_cache.get({"not_expected_key": "value"}, self.logger)
+
+        # this means the connection object is not saved in cache
+        self.assertEqual(self.connection_cache.connections, {})