import demistomock as demisto from CommonServerPython import * """ CONSTANTS """ DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ" VENDOR = "onelogin" PRODUCT = "iam" DEFAULT_LIMIT = 1000 """ HELPER FUNCTIONS """ def arg_to_strtime(value: Any) -> Optional[str]: if datetime_obj := arg_to_datetime(value): return datetime_obj.strftime(DATE_FORMAT) return None def prepare_query_params(params: dict, last_run: dict = {}) -> dict: """ Parses the given inputs into OneLogin Events API expected format. """ query_params = { "limit": arg_to_number(params.get("limit", DEFAULT_LIMIT)), "since": last_run.get("since") or arg_to_strtime(params.get("since")), "until": arg_to_strtime(params.get("until")), "event_type_id": params.get("event_type_id"), "after_cursor": last_run.get("after_cursor") or params.get("after_cursor"), } return query_params def check_response(raw_response): """ Checks the status of the API raw respones and raises an exceptionif it's an error. """ status = raw_response.get("status") if status and status.get("code") != 200 and status.get("error"): demisto.error(f"Failed to get events from OneLogin API. Error message: {status.get('message')}") raise Exception( f"Error code: {status.get('code')}. Error type: {status.get('type')}.\nError message: {status.get('message')}" ) def get_last_event_ids(events: List[dict]) -> list: """ Gets the last event ids with the same created time from the last fetch to prevent duplications in the next fetches. """ ids = [] last_time = events[-1].get("created_at", "").rsplit(".")[0] # Compare the times by seconds and not milliseconds for event in reversed(events): if event.get("created_at", "").rsplit(".")[0] != last_time: break ids.append(event.get("id")) return ids """ CLIENT CLASS """ class Client(BaseClient): def get_access_token_request(self): """ Send request to get an access token from OneLogin. """ raw_response = self._http_request( method="POST", url_suffix="/auth/oauth2/v2/token", params={"grant_type": "client_credentials"} ) check_response(raw_response) demisto.info("Succesfully got access token.") return raw_response.get("access_token") def get_events_request(self, access_token: str, query_params: dict) -> tuple: """ Send request to get events from OneLogin. """ self._headers["Authorization"] = f"Bearer:{access_token}" raw_response = self._http_request(method="GET", url_suffix="/api/1/events", params=query_params) check_response(raw_response) events = raw_response.get("data", []) cursor = raw_response.get("pagination", {}).get("after_cursor") demisto.info(f"Succesfully got {len(events)} events.") return raw_response, events, cursor def get_event_types_request(self) -> list: """ Send request to get event types from OneLogin. """ raw_response = self._http_request(method="GET", url_suffix="/api/1/events/types") check_response(raw_response) demisto.info("Succesfully got event types.") return raw_response.get("data", []) def get_event_types_from_last_run(self, last_run: dict): """ Gets the event types from the LastRun obj. For the first fetch, will send a call to get the event types from OneLogin and will update the LastRun as well. """ if "event_types" not in last_run: event_types = self.get_event_types_request() last_run["event_types"] = {str(event_type["id"]): event_type["name"] for event_type in event_types} return last_run.get("event_types") def convert_type_id_to_name(self, event: dict, event_types: dict, last_run: dict) -> str: """ Gets the event type name by the event type id. Args: event: The event getting from the OneLogin events API call. event_types: Event types getting from the LastRun obj. last_run: The LastRun obj. Returns: (str): The event type name. """ if event_type_name := event_types.get(str(event["event_type_id"])): return event_type_name demisto.info( f"Could not find the event type id {event['event_type_id']!s}. Trying to request the event types from OneLogin." ) event_types_res = self.get_event_types_request() event_types = {str(event_type["id"]): event_type["name"] for event_type in event_types_res} if event_type_name := event_types.get(str(event["event_type_id"])): last_run["event_types"] = event_types return event_type_name demisto.info( f"Could not find the event type id '{event['event_type_id']!s}' " f"in the event types from OneLogin, returning an empty value" ) return "" def handle_pagination_first_batch(self, access_token: str, query_params: dict, last_run: dict) -> tuple: """ Makes the first events API call in the current fetch run. If `first_id` or `last_event_ids` exists in the lastRun obj, finds it in the response and returns only the subsequent events (that weren't collected yet). """ _, events, cursor = self.get_events_request(access_token, query_params) if last_run.get("first_id"): for idx, event in enumerate(events): if event.get("id") == last_run["first_id"]: events = events[idx:] break last_run.pop("first_id", None) # removing to make sure it won't be used in future runs if event_ids := last_run.get("last_event_ids"): events_from_index = 0 for idx, event in enumerate(events): if event.get("id") in event_ids: events_from_index = idx + 1 events = events[events_from_index:] return events, cursor def fetch_events(self, access_token: str, query_params: dict, last_run: dict) -> List[dict]: """ Aggregates events using cursor-based pagination, until one of the following occurs: 1. Reaches the user-defined limit (parameter). In this case, stores the last used `cursor` and `since` and the id of the next event to collect (`first_id`) and returns the events that have been accumulated so far. 2. Reaches the end of the pagination. In this case, the lastRun obj will be updated with the `cursor` that will be None and with a new value for `since` taken from the last event that were fetched. 3. Reaches a rate limit. In this case, stores the last cursor used in the lastRun obj and returns the events that have been accumulated so far. """ aggregated_events: List[dict] = [] event_types = self.get_event_types_from_last_run(last_run) user_defined_limit = query_params.pop("limit") query_params["limit"] = 1000 # Constant limit value for the request try: events, cursor = self.handle_pagination_first_batch(access_token, query_params, last_run) while events: for event in events: if len(aggregated_events) == user_defined_limit: demisto.info(f"Reached the user-defined limit ({user_defined_limit}) - stopping.") last_run["first_id"] = event.get("id") cursor = query_params["after_cursor"] break event["event_type_name"] = self.convert_type_id_to_name(event, event_types, last_run) aggregated_events.append(event) else: # Finished iterating through all events in this batch if cursor: demisto.info("Using the cursor from the last API call to execute the next call.") query_params["after_cursor"] = cursor _, events, cursor = self.get_events_request(access_token, query_params) continue demisto.info("Finished iterating through all events in this fetch run.") break except DemistoException as e: if not e.res or e.res.status_code != 429: raise e demisto.info("Reached API rate limit, storing last used cursor.") cursor = query_params["after_cursor"] since = query_params["since"] if cursor or not aggregated_events else aggregated_events[-1].get("created_at") event_ids = last_run["last_event_ids"] if not aggregated_events else get_last_event_ids(aggregated_events) last_run.update({"after_cursor": cursor, "since": since, "last_event_ids": event_ids}) return aggregated_events """ COMMAND FUNCTIONS """ def test_module_command(client: Client, params: dict) -> str: """ Tests connection to OneLogin. Args: client (Client): The client implementing the API to OneLogin. params (dict): The configuration parameters. Returns: (str) 'ok' if success. """ access_token = client.get_access_token_request() params = prepare_query_params(params) client.get_events_request(access_token, params) return "ok" def get_events_command(client: Client, args: dict) -> tuple[list, CommandResults]: """ Gets log events from OneLogin. Args: cilent (Client): the client implementing the API to OneLogin. args (dict): the command arguments. Returns: (list) the events retrieved from the API call. (CommandResults) the CommandResults object holding the collected events information. """ access_token = client.get_access_token_request() query_params = prepare_query_params(args) raw_response, events, cursor = client.get_events_request(access_token, query_params) results = CommandResults( raw_response=raw_response, readable_output=tableToMarkdown( "OneLogin Events", events, metadata=f"Cursor: {cursor}" if cursor else None, date_fields=["created_at"], removeNull=True, ), ) return events, results def fetch_events_command(client: Client, params: dict, last_run: dict) -> tuple[list, dict]: """ Collects log events from OneLogin using pagination. Args: client (Client): the client implementing the API to OneLogin. params (dict): the instance configuration parameters. last_run (dict): the lastRun object, holding information from the previous run. Returns: (list) the events retrieved from the API call. (dict) the updated lastRun object. """ access_token = client.get_access_token_request() query_params = prepare_query_params(params, last_run) events = client.fetch_events(access_token, query_params, last_run) return events, last_run """ MAIN FUNCTION """ def main() -> None: command = demisto.command() params = demisto.params() args = demisto.args() demisto.info(f"Command being called is {command}") try: client_id = params.get("credentials", {}).get("identifier") client_secret = params.get("credentials", {}).get("password") client = Client( base_url=params.get("url"), verify=not params.get("insecure", False), proxy=params.get("proxy", False), headers={ "Content-Type": "application/json", "Authorization": f"client_id:{client_id}, client_secret:{client_secret}", }, ) if command == "test-module": return_results(test_module_command(client, params)) elif command == "onelogin-get-events": events, results = get_events_command(client, args) return_results(results) if argToBoolean(args.get("should_push_events", "false")): send_events_to_xsiam(events, vendor=VENDOR, product=PRODUCT) elif command == "fetch-events": last_run = demisto.getLastRun() events, last_run = fetch_events_command(client, params, last_run) send_events_to_xsiam(events, vendor=VENDOR, product=PRODUCT) demisto.setLastRun(last_run) except Exception as e: return_error(f"Failed to execute {command} command.\nError:\n{e}") """ ENTRY POINT """ if __name__ in ("__main__", "__builtin__", "builtins"): main()