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

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-5.5.3.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.3
+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,7 @@ 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

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

@@ -18,7 +18,7 @@ insightconnect_plugin_runtime/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeR
 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.3.dist-info/METADATA,sha256=md4T9Yk6AkiGS9oQfd-kSxfL7QIoSqZedyp9abVmWL8,13720
-insightconnect_plugin_runtime-5.5.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-insightconnect_plugin_runtime-5.5.3.dist-info/top_level.txt,sha256=AJtyJOpiFzHxsbHUICTcUKXyrGQ3tZxhrEHsPjJBvEA,36
-insightconnect_plugin_runtime-5.5.3.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"})