From 052593cb60e4e8338a5af99a7b36a5f06571108b Mon Sep 17 00:00:00 2001 From: lcwiklinski-r7 Date: Tue, 14 Jan 2025 13:46:56 +0100 Subject: [PATCH] Duo admin Improve handling of 429s in task code --- plugins/duo_admin/.CHECKSUM | 6 +- plugins/duo_admin/help.md | 2 +- .../actions/get_logs/schema.py | 5 +- .../tasks/monitor_logs/schema.py | 4 +- .../tasks/monitor_logs/task.py | 55 +++++++++++++++++-- .../komand_duo_admin/util/constants.py | 1 + .../duo_admin/komand_duo_admin/util/util.py | 7 +++ plugins/duo_admin/plugin.spec.yaml | 3 +- .../expected/monitor_logs_rate_limit.json.exp | 10 ++++ .../monitor_logs_with_rate_limit.json.inp | 6 ++ .../duo_admin/unit_test/test_monitor_logs.py | 10 +++- 11 files changed, 93 insertions(+), 16 deletions(-) create mode 100644 plugins/duo_admin/komand_duo_admin/util/util.py create mode 100644 plugins/duo_admin/unit_test/expected/monitor_logs_rate_limit.json.exp create mode 100644 plugins/duo_admin/unit_test/inputs/monitor_logs_with_rate_limit.json.inp diff --git a/plugins/duo_admin/.CHECKSUM b/plugins/duo_admin/.CHECKSUM index 67ebe691da..a60a868759 100644 --- a/plugins/duo_admin/.CHECKSUM +++ b/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" } ] } \ No newline at end of file diff --git a/plugins/duo_admin/help.md b/plugins/duo_admin/help.md index b665dc020a..2f8ae6e8b5 100644 --- a/plugins/duo_admin/help.md +++ b/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 diff --git a/plugins/duo_admin/komand_duo_admin/actions/get_logs/schema.py b/plugins/duo_admin/komand_duo_admin/actions/get_logs/schema.py index 1e93938303..0ff701033c 100755 --- a/plugins/duo_admin/komand_duo_admin/actions/get_logs/schema.py +++ b/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: diff --git a/plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/schema.py b/plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/schema.py index fc0dd65573..10d4b377e5 100755 --- a/plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/schema.py +++ b/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" ], diff --git a/plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/task.py b/plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/task.py index f6db9f23e9..297a1a56b4 100755 --- a/plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/task.py +++ b/plugins/duo_admin/komand_duo_admin/tasks/monitor_logs/task.py @@ -1,14 +1,17 @@ -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" @@ -16,6 +19,7 @@ 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] = [] diff --git a/plugins/duo_admin/komand_duo_admin/util/constants.py b/plugins/duo_admin/komand_duo_admin/util/constants.py index 76a42056c5..ca1c86c749 100644 --- a/plugins/duo_admin/komand_duo_admin/util/constants.py +++ b/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: diff --git a/plugins/duo_admin/komand_duo_admin/util/util.py b/plugins/duo_admin/komand_duo_admin/util/util.py new file mode 100644 index 0000000000..2507dd328c --- /dev/null +++ b/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") diff --git a/plugins/duo_admin/plugin.spec.yaml b/plugins/duo_admin/plugin.spec.yaml index b63e28d6eb..9380a8ea3b 100644 --- a/plugins/duo_admin/plugin.spec.yaml +++ b/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`)" diff --git a/plugins/duo_admin/unit_test/expected/monitor_logs_rate_limit.json.exp b/plugins/duo_admin/unit_test/expected/monitor_logs_rate_limit.json.exp new file mode 100644 index 0000000000..c3ea7bf5c2 --- /dev/null +++ b/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 +} diff --git a/plugins/duo_admin/unit_test/inputs/monitor_logs_with_rate_limit.json.inp b/plugins/duo_admin/unit_test/inputs/monitor_logs_with_rate_limit.json.inp new file mode 100644 index 0000000000..5731e01b98 --- /dev/null +++ b/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": [] +} diff --git a/plugins/duo_admin/unit_test/test_monitor_logs.py b/plugins/duo_admin/unit_test/test_monitor_logs.py index f74215a19e..e204bcaae8 100644 --- a/plugins/duo_admin/unit_test/test_monitor_logs.py +++ b/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"))