Duo admin Improve handling of 429s in task code

plugins/duo_admin/.CHECKSUM

@@ -1,5 +1,5 @@
 {
-	"spec": "4a94a77c37f17820e8768bb2850f6ee5",
+	"spec": "21815acfef998298925be284831306ce",
 	"manifest": "672f0da4df4edb87ab669d69e435c5c7",
 	"setup": "8a8919e13bd1afe4849427d3dae6dbf4",
 	"schemas": [
@@ -17,7 +17,7 @@
 		},
 		{
 			"identifier": "get_logs/schema.py",
-			"hash": "3502cb177351d18ff8a31266a49db228"
+			"hash": "412d7f456b6aba9f4a7b87092212e7d3"
 		},
 		{
 			"identifier": "get_phones_by_user_id/schema.py",
@@ -49,7 +49,7 @@
 		},
 		{
 			"identifier": "monitor_logs/schema.py",
-			"hash": "26f03015b329bc573a7e6f3a688fb861"
+			"hash": "ff4f7adf6cbae20cd793af79e763a06d"
 		}
 	]
 }

plugins/duo_admin/help.md

@@ -946,7 +946,7 @@ Example output:
 
 ## Troubleshooting
 
-* Many actions in this plugin take a User ID as input. A User ID is not the username - instead it's a unique identifier e.g. DU9I6T0F7R2S1J4XZHHA. A User ID can be obtained by passing a username to the Get User Status action.
+Many actions in this plugin take a User ID as input. A User ID is not the username - instead it's a unique identifier e.g. DU9I6T0F7R2S1J4XZHHA. A User ID can be obtained by passing a username to the Get User Status action.
 
 # Version History
 

plugins/duo_admin/komand_duo_admin/actions/get_logs/schema.py

@@ -4,7 +4,10 @@
 
 
 class Component:
-    DESCRIPTION = "This action is used to get auth logs, limited to past 180 days.[Currentmillis.com](https://currentmillis.com/) is useful for finding a usable UNIX timestamp.Available inputs for parameters can be found in [Duo Admin API docs](https://duo.com/docs/adminapi#logs:~:text=The%20factor%20or%20method%20used%20for%20an%20authentication%20attempt.%20One%20of%3A)"
+    DESCRIPTION = "This action is used to get auth logs, limited to past 180 days.
+[Currentmillis.com](https://currentmillis.com/) is useful for finding a usable UNIX timestamp.
+
+Available inputs for parameters can be found in [Duo Admin API docs](https://duo.com/docs/adminapi#logs:~:text=The%20factor%20or%20method%20used%20for%20an%20authentication%20attempt.%20One%20of%3A)"
 
 
 class Input:

plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/schema.py

@@ -64,7 +64,9 @@ class MonitorLogsOutput(insightconnect_plugin_runtime.Output):
   "type": "array",
   "title": "Logs",
   "description": "List of administrator, authentication and trust monitor event logs within the specified time range",
-  "items": {},
+  "items": {
+    "$ref": {}
+  },
   "required": [
     "logs"
   ],

plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/task.py

@@ -1,21 +1,25 @@
-from typing import Dict
+from datetime import datetime, timedelta, timezone
+from hashlib import sha1
+from time import time
+from typing import Dict, Tuple, Any
 
 import insightconnect_plugin_runtime
 from insightconnect_plugin_runtime.exceptions import PluginException
 
-from .schema import MonitorLogsInput, MonitorLogsOutput, MonitorLogsState, Component, Input
-
 # Custom imports below
+from komand_duo_admin.util.constants import Assistance
 from komand_duo_admin.util.exceptions import ApiException
-from datetime import datetime, timedelta, timezone
-from hashlib import sha1
+from komand_duo_admin.util.util import Utils
+
+from .schema import MonitorLogsInput, MonitorLogsOutput, MonitorLogsState, Component, Input
 
 ADMIN_LOGS_LOG_TYPE = "Admin logs"
 AUTH_LOGS_LOG_TYPE = "Auth logs"
 TRUST_MONITOR_EVENTS_LOG_TYPE = "Trust monitor events"
 INITIAL_CUTOFF_HOURS = 24
 MAX_CUTOFF_HOURS = 168
 API_CUTOFF_HOURS = 4320
+RATE_LIMIT_DELAY = 600
 
 
 class MonitorLogs(insightconnect_plugin_runtime.Task):
@@ -30,6 +34,7 @@ class MonitorLogs(insightconnect_plugin_runtime.Task):
     PREVIOUS_ADMIN_LOG_HASHES = "previous_admin_log_hashes"
     PREVIOUS_AUTH_LOG_HASHES = "previous_auth_log_hashes"
     PREVIOUS_TRUST_MONITOR_EVENT_HASHES = "previous_trust_monitor_event_hashes"
+    RATE_LIMIT_DATETIME = "rate_limit_datetime"
 
     def __init__(self):
         super(self.__class__, self).__init__(
@@ -103,7 +108,44 @@ def get_parameters_for_query(
         self.logger.info(f"Retrieve data from {mintime} to {maxtime}. Get next page is set to {get_next_page}")
         return mintime, maxtime, get_next_page
 
+    def check_rate_limit(self, state: Dict):
+        rate_limited = state.get(self.RATE_LIMIT_DATETIME)
+        now = time()
+        if rate_limited:
+            rate_limit_string = Utils.convert_epoch_to_readable(rate_limited)
+            log_msg = f"Rate limit value stored in state: {rate_limit_string}. "
+            if rate_limited > now:
+                log_msg += "Still within rate limiting period, skipping task execution..."
+                self.logger.info(log_msg)
+                error = PluginException(
+                    cause=PluginException.causes.get(PluginException.Preset.RATE_LIMIT),
+                    assistance=Assistance.RATE_LIMIT,
+                )
+                return error
+
+            log_msg += "However no longer in rate limiting period, so task can be executed..."
+            del state[self.RATE_LIMIT_DATETIME]
+            self.logger.info(log_msg)
+
+    def check_rate_limit_error(self, error: ApiException, status_code: int, state: dict, rate_limit_delay: int) -> Tuple[int, Any]:
+        if status_code == 429:
+            new_run_time = time() + rate_limit_delay  # default to wait 10 minutes before the next run
+            try:
+                new_run_time_string = Utils.convert_epoch_to_readable(new_run_time)
+                self.logger.error(f"A rate limit error has occurred, task will resume after {new_run_time_string}")
+                state[self.RATE_LIMIT_DATETIME] = new_run_time
+            except Exception as err:
+                self.logger.error(
+                    f"Unable to calculate new run time, no rate limiting applied to the state. Error: {repr(err)}",
+                    exc_info=True,
+                )
+            return 200, None
+        return status_code, error
+
     def run(self, params={}, state={}, custom_config={}):  # noqa: C901
+        rate_limit_delay = custom_config.get("rate_limit_delay", RATE_LIMIT_DELAY)
+        if rate_limited := self.check_rate_limit(state):
+            return [], state, False, 429, rate_limited
         self.connection.admin_api.toggle_rate_limiting = False
         has_more_pages = False
         backward_comp_first_run = False
@@ -257,7 +299,8 @@ def run(self, params={}, state={}, custom_config={}):  # noqa: C901
                 state[self.PREVIOUS_TRUST_MONITOR_EVENT_HASHES] = []
                 state[self.PREVIOUS_ADMIN_LOG_HASHES] = []
                 state[self.PREVIOUS_AUTH_LOG_HASHES] = []
-                return [], state, False, error.status_code, error
+                status_code, error = self.check_rate_limit_error(error, error.status_code, state, rate_limit_delay)
+                return [], state, False, status_code, error
         except Exception as error:
             self.logger.info(f"An Exception has been raised. Error: {error}")
             state[self.PREVIOUS_TRUST_MONITOR_EVENT_HASHES] = []

plugins/duo_admin/komand_duo_admin/util/constants.py

@@ -19,6 +19,7 @@ class Assistance:
     VERIFY_INPUT = (
         "Verify your input is correct and not malformed and try again. If the issue persists, please contact support."
     )
+    RATE_LIMIT = "Task will resume collection of logs after the rate limiting period has expired."
 
 
 class PossibleInputs:

plugins/duo_admin/komand_duo_admin/util/util.py

@@ -0,0 +1,7 @@
+from datetime import datetime
+
+
+class Utils:
+    @staticmethod
+    def convert_epoch_to_readable(epoch_time: float) -> str:
+        return datetime.utcfromtimestamp(epoch_time).strftime("%Y-%m-%d %H:%M:%S")

plugins/duo_admin/plugin.spec.yaml

@@ -47,8 +47,7 @@ links:
 - "[Duo Security](https://duo.com/)"
 references:
 - "[Duo Admin API](https://duo.com/docs/adminapi)"
-troubleshooting:
-  - "Many actions in this plugin take a User ID as input. A User ID is not the username - instead it's a unique identifier e.g. DU9I6T0F7R2S1J4XZHHA. A User ID can be obtained by passing a username to the Get User Status action."
+troubleshooting: "Many actions in this plugin take a User ID as input. A User ID is not the username - instead it's a unique identifier e.g. DU9I6T0F7R2S1J4XZHHA. A User ID can be obtained by passing a username to the Get User Status action."
 version_history:
 - "5.0.2 - Updated SDK to the latest version (v6.2.2) | Address vulnerabilities"
 - "5.0.1 - Update to enable Plugin as FedRAMP ready | Update SDK (`6.1.2`)"

plugins/duo_admin/unit_test/expected/monitor_logs_rate_limit.json.exp

@@ -0,0 +1,10 @@
+{
+  "logs": [],
+  "state": {
+    "previous_admin_log_hashes": [],
+    "previous_auth_log_hashes": [],
+    "rate_limit_datetime": 3050853836.506636,
+    "previous_trust_monitor_event_hashes": []
+},
+  "status_code": 429
+}

plugins/duo_admin/unit_test/inputs/monitor_logs_with_rate_limit.json.inp

@@ -0,0 +1,6 @@
+{
+    "previous_admin_log_hashes": [],
+    "previous_auth_log_hashes": [],
+    "rate_limit_datetime": 3050853836.506636,
+    "previous_trust_monitor_event_hashes": []
+}

plugins/duo_admin/unit_test/test_monitor_logs.py

@@ -21,7 +21,7 @@
 class TestMonitorLogs(TestCase):
     @classmethod
     def setUpClass(cls) -> None:
-        cls.action = Util.default_connector(MonitorLogs())
+        cls.task = Util.default_connector(MonitorLogs())
         cls.custom_config = {"cutoff": {"date": "2023-04-30T08:34:46.000Z"}, "lookback": "2023-04-30T08:34:46.000Z"}
 
     @parameterized.expand(
@@ -37,6 +37,12 @@ def setUpClass(cls) -> None:
                     "lookback": "2023-04-30T08:34:46.000Z",
                 },
             ],
+            [
+                "with_rate_limit",
+                Util.read_file_to_dict("inputs/monitor_logs_with_rate_limit.json.inp"),
+                Util.read_file_to_dict("expected/monitor_logs_rate_limit.json.exp"),
+                {},
+            ],
             [
                 "with_state",
                 Util.read_file_to_dict("inputs/monitor_logs_with_state.json.inp"),
@@ -94,7 +100,7 @@ def test_monitor_logs(
         expected,
         config,
     ):
-        actual, actual_state, has_more_pages, status_code, _ = self.action.run(
+        actual, actual_state, has_more_pages, status_code, _ = self.task.run(
             state=current_state, custom_config=config
         )
         self.assertEqual(actual, expected.get("logs"))