import traceback from datetime import datetime import demistomock as demisto # noqa: F401 import urllib3 from CommonServerPython import * # noqa: F401 from CommonServerUserPython import * # noqa urllib3.disable_warnings() """ CONSTANTS """ VENDOR = "Proofpoint" PRODUCT = "Isolation" DEFAULT_FETCH_LIMIT = 50000 ITEMS_PER_PAGE = 10000 DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ" """ CLIENT CLASS """ class Client(BaseClient): """ Client class to interact with the service API """ def __init__(self, base_url, verify: bool, api_key: str) -> None: self.api_key = api_key super().__init__(base_url=base_url, verify=verify) def get_events(self, start_date: str, end_date: str) -> dict: """ Gets events from the specified start date to the end date using the API. Args: start_date (str): The start date for the data retrieval in YYYY-MM-DD format. end_date (str): The end date for the data retrieval in YYYY-MM-DD format. Returns: dict: The API response containing the usage data events. """ demisto.debug(f"[HTTP Call] Making request: {start_date=}, {end_date=}, pageSize={ITEMS_PER_PAGE}") results = self._http_request( method="GET", url_suffix=f"/api/v2/reporting/usage-data?key={self.api_key}&pageSize={ITEMS_PER_PAGE}" f"&from={start_date}&to={end_date}", retries=3, ) demisto.debug( "[HTTP Call] Request completed. Response keys: ", f"{list(results.keys()) if isinstance(results, dict) else type(results)}", ) return results """ HELPER FUNCTIONS """ def get_and_parse_date(event: dict) -> str | None: """ Parses the date string from an event dictionary and formats it according to the specified date format. Args: event (dict): A dictionary containing event data. Returns: str: The formatted date string if parsing is successful. Raises: ValueError: If the 'date' value in the event dictionary is invalid or cannot be parsed. """ date_str = event.get("date") try: start = parse_date_string(date_str, DATE_FORMAT) return start.strftime(DATE_FORMAT) except ValueError: raise ValueError("Invalid date format") def sort_events_by_date(events: list) -> list: """ Sorts a list of events by their date in ascending order. Args: events (list): A list of dictionaries. Returns: list: The sorted list of events based on the 'date' field. """ return sorted(events, key=lambda x: datetime.strptime(x["date"], "%Y-%m-%dT%H:%M:%S.%f%z")) def hash_user_name_and_url(event: dict) -> str: """ Generates a hash-like string by concatenating the 'url' and 'userName' fields from an event dictionary. Args: event (dict): A dictionary containing event data. Returns: str: A string in the format '&'. """ url = event.get("url", "") user_name = event.get("userName", "") return f"{url}&{user_name}" def remove_duplicate_events(start_date, ids: set, events: list) -> None: """ Removes duplicate events from a list of events based on a set of unique identifiers and a specified start date. Args: start_date (str): The date to check against, in the same format as the event dates. ids (set): A set of hashed identifiers for detecting duplicates. events (list): A list of sorted event dictionaries to process. """ events_copy = events.copy() for event in events_copy: event_date = get_and_parse_date(event) if event_date != start_date: break hashed_id = hash_user_name_and_url(event) if hashed_id in ids: demisto.debug(f"[Dedup] Removing duplicated event with hash_id {hashed_id}") events.remove(event) def get_and_reorganize_events(client: Client, start: str, end: str, ids: set) -> list: """ Fetches events, sorts them by date, and removes duplicates. Args: client (Client): The client to fetch events from. start (str): The start date for fetching events. end (str): The end date for fetching events. ids (set): A set of already processed event IDs to filter out duplicates. Returns: list: A list of sorted and deduplicated events. """ demisto.debug(f"[API Fetch] Fetching events: {start=}, {end=}, existing_ids_count={len(ids)}") events: list = client.get_events(start, end).get("data", []) demisto.debug(f"[API Fetch] Got {len(events)} raw events from API.") events = sort_events_by_date(events) remove_duplicate_events(start, ids, events) demisto.debug(f"[Dedup] {len(events)} events remain after dedup.") return events """ COMMAND FUNCTIONS """ def test_module(client: Client) -> str: """ Tests the connection to the service by attempting to fetch events within a date range. Args: client (Client): The client object used to interact with the service. Returns: str: 'ok' if the connection is successful. If an authorization error occurs, an appropriate error message is returned. """ try: current_time = get_current_time() start_date = (current_time - timedelta(minutes=1)).strftime(DATE_FORMAT) end_date = current_time.strftime(DATE_FORMAT) demisto.debug(f"[Test Module] Testing connection: {start_date=}, {end_date=}") fetch_events(client, 1, {"start_date": start_date, "end_date": end_date}) demisto.debug("[Test Module] Success") message = "ok" except DemistoException as e: raise e return message def fetch_events(client: Client, fetch_limit: int, get_events_args: dict = None) -> tuple[list, dict]: output: list = [] if get_events_args: # handle get_event command event_date = get_events_args.get("start_date", "") end = get_events_args.get("end_date", "") ids: set = set() demisto.debug(f"[Fetch] get-events mode: {event_date=}, {end=}") else: # handle fetch_events case last_run = demisto.getLastRun() or {} demisto.debug(f"[Fetch] fetch-events mode: {last_run=}") event_date = last_run.get("start_date", "") if not event_date: event_date = get_current_time().strftime(DATE_FORMAT) end = get_current_time().strftime(DATE_FORMAT) ids = set(last_run.get("ids", [])) current_start_date = event_date demisto.debug(f"[Fetch] Starting fetch loop: {current_start_date=}, {end=}, {fetch_limit=}") has_more_events = True iteration = 0 while has_more_events and len(output) < fetch_limit: iteration += 1 demisto.debug(f"[Fetch Pagination Loop] Loop iteration {iteration}: {event_date=}") events = get_and_reorganize_events(client, event_date, end, ids) if not events: demisto.debug(f"[Fetch] No more events found. Total: {len(output)}") event_date = end ids = set() has_more_events = False continue demisto.debug(f"[Fetch] Processing {len(events)} events.") for event in events: event["_TIME"] = event.get("date") output.append(event) event_date = get_and_parse_date(event) if event_date != current_start_date: current_start_date = event_date ids = set() hashed_id = hash_user_name_and_url(event) ids.add(hashed_id) if len(output) >= fetch_limit: demisto.debug(f"[Fetch] Reached fetch limit. Total: {len(output)}") break new_last_run = {"start_date": event_date, "ids": list(ids)} demisto.debug(f"[Fetch] Fetch complete. Total: {len(output)}") return output, new_last_run def get_events(client: Client, args: dict) -> tuple[list, CommandResults]: """ Fetches events within the specified date range and returns them. Args: client (Client): The client to fetch events from. args (dict): A dictionary containing the start and end dates for the query. Returns: list: A list of events fetched within the specified date range. """ start_date = args.get("start_date") end_date = args.get("end_date") limit: int = arg_to_number(args.get("limit")) or DEFAULT_FETCH_LIMIT demisto.debug(f"[Get Events Command]: {start_date=}, {end_date=}, {limit=}") output, _ = fetch_events(client, limit, {"start_date": start_date, "end_date": end_date}) demisto.debug(f"[Get Events Command] fetched {len(output)} events.") filtered_events = [] for event in output: filtered_event = { "User ID": event.get("userId"), "User Name": event.get("userName"), "URL": event.get("url"), "Date": event.get("date"), } filtered_events.append(filtered_event) human_readable = tableToMarkdown(name="Proofpoint Isolation Events", t=filtered_events, removeNull=True) command_results = CommandResults( readable_output=human_readable, outputs=output, outputs_prefix="ProofpointIsolationEventCollector", ) return output, command_results """ MAIN FUNCTION """ def main() -> None: # pragma: no cover """main function, parses params and runs command functions""" params = demisto.params() command = demisto.command() args = demisto.args() demisto.debug(f"[Main] Command being called is {command}") try: base_url = params.get("base_url") verify = not params.get("insecure", False) api_key = params.get("credentials").get("password") fetch_limit = arg_to_number(params.get("max_events_per_fetch")) or DEFAULT_FETCH_LIMIT client = Client(base_url=base_url, verify=verify, api_key=api_key) if command == "test-module": result = test_module(client) return_results(result) elif command == "fetch-events": demisto.debug(f"[Fetch] Starting with last_run={demisto.getLastRun()}") events, new_last_run_dict = fetch_events(client, fetch_limit) if events: demisto.debug(f"[Fetch] Sending {len(events)} events to XSIAM.") send_events_to_xsiam(events=events, vendor=VENDOR, product=PRODUCT) demisto.debug(f"[Fetch] Successfully sent {len(events)} events to XSIAM.") else: demisto.debug("[Fetch] No events fetched in this cycle.") demisto.setLastRun(new_last_run_dict) demisto.debug(f"[Fetch] Last run updated: {new_last_run_dict}") elif command == "proofpoint-isolation-get-events": events, command_results = get_events(client, args) if events and argToBoolean(args.get("should_push_events")): send_events_to_xsiam(events=events, vendor=VENDOR, product=PRODUCT) return_results(command_results) except Exception as e: return_error(f"Failed to execute {command} command.\nError:\n{e!s}\nTraceback:\n{traceback.format_exc()}") """ ENTRY POINT """ if __name__ in ("__main__", "__builtin__", "builtins"): main()