insightconnect-plugin-runtime 5.5.2__py3-none-any.whl → 5.5.4__py3-none-any.whl

insightconnect_plugin_runtime/api/endpoints.py

@@ -634,14 +634,19 @@ class Endpoints:
             schema = conn.schema
             return jsonify(ConnectionDetailsSchema().dump(schema))
 
-        @v1.route("/connection/test", methods=["POST"])
-        def connection_test():
+        @v1.route("/connection/<string:connection_test_type>", methods=["POST"])
+        def connection_test(connection_test_type: str):
             """
             Run connection test endpoint
             ---
             post:
               summary: Run connection test
               description: Run InsightConnect plugin connection test
+              parameters:
+                - in: path
+                  name: connection_test_type
+                  description: Type of connection test to be run
+                  type: string
               responses:
                 200:
                   description: Connection test output to be returned
@@ -658,12 +663,26 @@ class Endpoints:
             status_code = 200
             output = None
 
+            supported_tests = ["test", "test_task"]
+            if connection_test_type not in supported_tests:
+                return make_response(
+                    jsonify(
+                        {
+                            "error": f"The requested endpoint is not available, only {supported_tests} are supported.",
+                            "code": 404,
+                            "method": "POST",
+                            "name": "Not Found"
+                        }
+                    ),
+                    404,
+                )
+
             input_message = request.get_json(force=True)
             Endpoints.validate_action_trigger_task_empty_input(input_message)
 
             try:
                 output = self.plugin.handle_step(
-                    input_message, is_debug=self.debug, is_test=True
+                    input_message, is_debug=self.debug, is_test=True, connection_test_type=connection_test_type
                 )
                 if output.get("body", {}).get("output") is None:
                     status_code = 204

insightconnect_plugin_runtime/clients/aws_client.py

@@ -36,6 +36,7 @@ class PaginationHelper:
         limit_key: str = None,
         more_results: str = None,
         non_aggregate_keys: List[str] = None,
+        max_pages : int = None,  # if we want to limit how many pages of data we pull from AWS
     ):
         self.input_token = input_token
         self.output_token = output_token
@@ -46,6 +47,7 @@ class PaginationHelper:
         self.keys_to_remove = []
         self.keys_to_remove.extend(input_token)
         self.keys_to_remove.extend(output_token)
+        self.max_pages = max_pages
         if more_results:
             self.keys_to_remove.append(self.more_results)
 
@@ -83,6 +85,20 @@ class PaginationHelper:
 
         return is_paginated
 
+    def check_total_results(self, params, response: dict[str, Any]) -> bool:
+        """
+        Check if we have ran the AWS call too many times, eg params[max_keys] * max_pages = len response
+        :param params: parameters used within the ListObjectsv2 command
+        :param response: response contents from ListObjectsv2
+        :return: boolean to indicate if we have reached the max amount of objects to be returned in one execution.
+        """
+        if self.max_pages:
+            max_allowed = params[self.limit_key] * self.max_pages
+            if max_allowed >= len(response[self.result_key[0]]):
+                return True
+
+        return False
+
     def merge_responses(
         self,
         input_: Dict[str, Any],
@@ -429,7 +445,7 @@ class AWSAction(Action):
         helper = self.connection.helper
 
         # Format the input parameters for the botocall call
-        self.logger.info(params)
+        self.logger.debug(params)
         try:
             params = helper.format_input(params)
         except Exception:
@@ -437,7 +453,7 @@ class AWSAction(Action):
             raise PluginException(cause="Unable to format input parameters")
 
         # Execute the botocore function
-        self.logger.info(params)
+        self.logger.debug(params)
         response = self._handle_botocore_function(client_function, params)
 
         # Format the output parameters for the komand action output schema
@@ -495,13 +511,21 @@ class AWSAction(Action):
         # Handle possible pagination if this action supports pagination.
         if self.pagination_helper:
             while self.pagination_helper.check_pagination(params, response):
+                if self.pagination_helper.check_total_results(params, response):
+                    self.logger.info("Reached max amount of pages per execution breaking...")
+                    break
+
                 self.logger.info("Response was paginated. Performing another call.")
                 response, max_hit = self.pagination_helper.merge_responses(
                     params, self.handle_rest_call(client_function, params), response
                 )
-                if max_hit:
+
+                if max_hit:  # all data possible in this particular AWS call
                     break
-            self.pagination_helper.remove_keys(response)
+
+            # if we have pagination enabled for AWS we want to take note of the tokens
+            if not self.pagination_helper.max_pages:
+                self.pagination_helper.remove_keys(response)
 
         return response
 

insightconnect_plugin_runtime/plugin.py

@@ -362,7 +362,7 @@ class Plugin(object):
         if "input" not in body:
             body["input"] = {}
 
-    def handle_step(self, input_message, is_test=False, is_debug=False):
+    def handle_step(self, input_message, is_test=False, is_debug=False, connection_test_type="test"):
         """
         Executes a single step, given the input message dictionary.
 
@@ -427,6 +427,7 @@ class Plugin(object):
                     log_stream,
                     is_test,
                     is_debug,
+                    connection_test_type=connection_test_type
                 )
             elif message_type == "trigger_start":
                 out_type = "trigger_event"
@@ -437,6 +438,7 @@ class Plugin(object):
                     log_stream,
                     is_test,
                     is_debug,
+                    connection_test_type=connection_test_type
                 )
             elif message_type == "task_start":
                 out_type = "task_event"
@@ -449,6 +451,7 @@ class Plugin(object):
                         log_stream,
                         is_test,
                         is_debug,
+                        connection_test_type=connection_test_type
                     )
                 else:
                     (
@@ -464,6 +467,7 @@ class Plugin(object):
                         log_stream,
                         is_test,
                         is_debug,
+                        connection_test_type=connection_test_type
                     )
             elif message_type == "connection_test":
                 out_type = "connection_test"
@@ -475,6 +479,7 @@ class Plugin(object):
                     is_test,
                     is_debug,
                     is_connection_test=True,
+                    connection_test_type=connection_test_type
                 )
         except (ClientException, ServerException, PluginException, Exception) as e:
             success = False
@@ -513,6 +518,7 @@ class Plugin(object):
         is_test=False,
         is_debug=False,
         is_connection_test=False,
+        connection_test_type="test"
     ):
         """
         Starts an action.
@@ -523,6 +529,7 @@ class Plugin(object):
         :param is_test: True if the action's test method should execute
         :param is_debug: True if debug is enabled
         :param is_connection_test: True if connection test is running
+        :param connection_test_type: The type of connection test to be run
         :return: An action_event message
         """
         connection = self.connection_cache.get(message_body["connection"], logger)
@@ -534,7 +541,9 @@ class Plugin(object):
                     plugin_version=connection.meta.version,
                 )
             )
-            if hasattr(connection, "test"):
+            if hasattr(connection, "test_task") and connection_test_type == "test_task":
+                func = connection.test_task
+            elif hasattr(connection, "test"):
                 func = connection.test
             else:
                 raise NotImplementedError(
@@ -617,7 +626,10 @@ class Plugin(object):
             if is_test:
                 # Check if connection test func available. If so - use it (preferred). Else fallback to action/trigger test
                 if hasattr(step.connection, "test"):
-                    func = step.connection.test
+                    if hasattr(step.connection, "test") and connection_test_type == "test_task":
+                        func = step.connection.task_test
+                    else:
+                        func = step.connection.test
                 else:
                     func = step.test
             else:

{insightconnect_plugin_runtime-5.5.2.dist-info → insightconnect_plugin_runtime-5.5.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: insightconnect-plugin-runtime
-Version: 5.5.2
+Version: 5.5.4
 Summary: InsightConnect Plugin Runtime
 Home-page: https://github.com/rapid7/komand-plugin-sdk-python
 Author: Rapid7 Integrations Alliance
@@ -211,6 +211,8 @@ contributed. Black is installed as a test dependency and the hook can be initial
 after cloning this repository.
 
 ## Changelog
+* 5.5.4 - Support pagination parameters within AWS client.
+* 5.5.3 - Adding in a new endpoint that can be called to run a task connection test
 * 5.5.2 - Address bug with typing for type `List` in Python 3.8
 * 5.5.1 - Address bug with typing for type `Tuple` in Python 3.8
 * 5.5.0 - Updated helper class to add `make_request`, `response_handler`, `extract_json`, and `request_error_handling` for HTTP requests, and `hash_sha1` and `compare_and_dedupe_hashes` to provide support for hash comparisons | Add `METHOD_NOT_ALLOWED`, `CONFLICT`, `REDIRECT_ERROR`, and `CONNECTION_ERROR` to PluginException presets

{insightconnect_plugin_runtime-5.5.2.dist-info → insightconnect_plugin_runtime-5.5.4.dist-info}/RECORD

@@ -6,7 +6,7 @@ insightconnect_plugin_runtime/dispatcher.py,sha256=ru7njnyyWE1-oD-VbZJ-Z8tELwvDf
 insightconnect_plugin_runtime/exceptions.py,sha256=7aYNoGgmV6SugHAQqeQYm1zo2LZm0vgXEGqHmYD7NCo,8260
 insightconnect_plugin_runtime/helper.py,sha256=XlMGpW2LQYo5KjhJeJtaYRCEDDWjrA9ZW1jP9FDpKpE,31155
 insightconnect_plugin_runtime/metrics.py,sha256=hf_Aoufip_s4k4o8Gtzz90ymZthkaT2e5sXh5B4LcF0,3186
-insightconnect_plugin_runtime/plugin.py,sha256=9k01QqEh78OTs0pMRfnwZcRr-vFe7iyAtLz_QBa23e4,24127
+insightconnect_plugin_runtime/plugin.py,sha256=s2YVyygPcNeXeDi3X_JmaCiLHXZWIqz4mlLJy6p7-UA,24904
 insightconnect_plugin_runtime/schema.py,sha256=jTNc6KAMqFpaDVWrAYhkVC6e8I63P3X7uVlJkAr1hiY,583
 insightconnect_plugin_runtime/server.py,sha256=09fxsbKf2ZZvSqRP2Bv9e9-fspDyEFR8_YgIFeMnXqQ,12578
 insightconnect_plugin_runtime/step.py,sha256=KdERg-789-s99IEKN61DR08naz-YPxyinPT0C_T81C4,855
@@ -15,10 +15,10 @@ insightconnect_plugin_runtime/trigger.py,sha256=Zq3cy68N3QxAGbNZKCID6CZF05Zi7YD2
 insightconnect_plugin_runtime/util.py,sha256=qPkZ3LA55nYuNYdansEbnCnBccQkpzIpp9NA1B64Kvw,8444
 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=57pSoVSyZU6ERPZ6joMUf6eTDXVRCROLJkmtw4DJhUg,31720
+insightconnect_plugin_runtime/api/endpoints.py,sha256=rqieWL71Eu0jIKrspizof4xSyp672xCQxbrQUYdpbbA,32560
 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=p_s7YLYOeSrBnJl-1gStuNwse6nw4_sl7kz6AKThboY,20465
+insightconnect_plugin_runtime/clients/aws_client.py,sha256=1hAhPD89ELlqtND2Vn46G4yaDG920gBlPg1Tyi02aAU,21658
 insightconnect_plugin_runtime/clients/oauth.py,sha256=bWtAGRMwdK4dw9vMPcw9usklyIHBDtZh55kMZ7sWROc,2453
 insightconnect_plugin_runtime/data/input_message_schema.json,sha256=7_BcHi6UOBiVWGrrJHHn5IoddteXjL7GOKETdO9T2DE,1770
 insightconnect_plugin_runtime/data/output_message_schema.json,sha256=Qya6U-NR5MfOlw4V98VpQzGBVq75eGMUQhI-j3yxOHI,1137
@@ -64,7 +64,7 @@ tests/plugin/hello_world/tests/test_server.py,sha256=7QfMNVsNcRoX5aAL0xM9tXS-ciy
 tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/test_action.py,sha256=0SVen1qNrFLKKAXMurN9T4NpItpftnSpDZ71AdrGZeo,1683
 tests/unit/test_api.py,sha256=uZ1dWMmgQ-ZePYjmcZfjc-qOTJsjs20Wic0Uf4U12-8,4441
-tests/unit/test_aws_action.py,sha256=uerrNB_cH0J9lnxEaVuVaMxTFZoPXuT9lmtpcwoG34Y,7952
+tests/unit/test_aws_action.py,sha256=eei4uPDx1Y8WcYgAWhiyqBYv5OOHGiZrCvajXC7RHdM,9597
 tests/unit/test_custom_encoder.py,sha256=KLYyVOTq9MEkZXyhVHqjm5LVSW6uJS4Davgghsw9DGk,2207
 tests/unit/test_endpoints.py,sha256=LuXOfLBu47rDjGa5YEsOwTZBEdvQdl_C6-r46oxWZA8,6401
 tests/unit/test_exceptions.py,sha256=t8c67n3ZQPpyzzYPN6TZHpNeE8uKlOSfHJPk71QONuU,4668
@@ -78,7 +78,7 @@ tests/unit/test_server_spec.py,sha256=je97BaktgK0Fiz3AwFPkcmHzYtOJJNqJV_Fw5hrvqX
 tests/unit/test_trigger.py,sha256=E53mAUoVyponWu_4IQZ0IC1gQ9lakBnTn_9vKN2IZfg,1692
 tests/unit/test_variables.py,sha256=OUEOqGYZA3Nd5oKk5GVY3hcrWKHpZpxysBJcO_v5gzs,291
 tests/unit/utils.py,sha256=VooVmfpIgxmglNdtmT32AkEDFxHxyRHLK8RsCWjjYRY,2153
-insightconnect_plugin_runtime-5.5.2.dist-info/METADATA,sha256=OHfxHfvU8vRAiFkDevFvTnw39pzjpWuYNkKQWQD9mlE,13636
-insightconnect_plugin_runtime-5.5.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-insightconnect_plugin_runtime-5.5.2.dist-info/top_level.txt,sha256=AJtyJOpiFzHxsbHUICTcUKXyrGQ3tZxhrEHsPjJBvEA,36
-insightconnect_plugin_runtime-5.5.2.dist-info/RECORD,,
+insightconnect_plugin_runtime-5.5.4.dist-info/METADATA,sha256=4T8uMLK4wr9KtC6eSxa-vn77qsWBbbDfCPA0K6uT3KQ,13779
+insightconnect_plugin_runtime-5.5.4.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+insightconnect_plugin_runtime-5.5.4.dist-info/top_level.txt,sha256=AJtyJOpiFzHxsbHUICTcUKXyrGQ3tZxhrEHsPjJBvEA,36
+insightconnect_plugin_runtime-5.5.4.dist-info/RECORD,,

tests/unit/test_aws_action.py

@@ -10,7 +10,7 @@ import botocore.response as br
 from parameterized import parameterized
 
 from insightconnect_plugin_runtime import Input, Output, Connection
-from insightconnect_plugin_runtime.clients.aws_client import AWSAction, ActionHelper
+from insightconnect_plugin_runtime.clients.aws_client import AWSAction, ActionHelper, PaginationHelper
 from insightconnect_plugin_runtime.exceptions import PluginException
 
 
@@ -80,6 +80,37 @@ class TestAwsAction(unittest.TestCase):
         self.aws_action.connection.helper.format_input.assert_called_once()
         self.aws_action.connection.helper.format_output.assert_called_once()
 
+    @unittest.mock.patch("botocore.session.Session", return_value=unittest.mock.Mock())
+    @unittest.mock.patch("boto3.client", return_value=Boto3Stub())
+    def test_run_action_pagination(self, mock_session, mock_sts_client):
+        """
+        Test when max_pages attribute is on the pagination_helper that the response is not cleaned
+        allowing the two tokens (next_continuation_token & contination_token) to be returned.
+        """
+        client_function = getattr(mock_session, "service")
+        # Mock connection as we're calling action.run() in this unit test
+        self.aws_action.connection.assume_role_params = {}
+        self.aws_action.connection.auth_params = {}
+        self.aws_action.connection.client = {}
+
+        self.aws_action.pagination_helper = unittest.mock.create_autospec(PaginationHelper)
+        self.aws_action.connection.client = mock_session
+
+        # First run has no pagination set so we clean the response
+        self.aws_action.pagination_helper.max_pages = None
+        self.aws_action.run()
+
+        client_function.assert_called_once()
+        self.aws_action.connection.helper.format_input.assert_called_once()
+        self.aws_action.connection.helper.format_output.assert_called_once()
+        self.aws_action.pagination_helper.remove_keys.assert_called_once()
+
+        # Call again with max_pages set and this means remove_keys is not called
+        self.aws_action.pagination_helper.max_pages = 1
+        self.aws_action.pagination_helper.remove_keys.reset_mock()
+        self.aws_action.run()
+        self.aws_action.pagination_helper.remove_keys.assert_not_called()
+
     def _mock_call_raise_endpoint_connection_error(self):
         raise be.EndpointConnectionError(**{"endpoint_url": "test_url"})