MongoDBAtlasEventCollector.py

from datetime import datetime

import demistomock as demisto
from CommonServerPython import *
from CommonServerUserPython import *
import urllib3
from requests.auth import HTTPDigestAuth

# Disable insecure warnings
urllib3.disable_warnings()

""" CONSTANTS """

DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ"  # ISO8601 format with UTC, default in XSOAR
VENDOR = "MongoDB"
PRODUCT = "Atlas"
DEFAULT_FETCH_LIMIT = 2500
ITEMS_PER_PAGE = 500

""" CLIENT CLASS """


class Client(BaseClient):
    """
    Client class to interact with the service API
    """

    def __init__(self, base_url, verify: bool, group_id: str, private_key: str = "", public_key: str = "") -> None:
        self.group_id = group_id
        auth = HTTPDigestAuth(public_key, private_key)
        headers = {"Accept": "application/vnd.atlas.2023-02-01+json"}
        super().__init__(base_url=base_url, verify=verify, headers=headers, auth=auth)

    def get_alerts_request(self, page_num: int = None, items_per_page: int = ITEMS_PER_PAGE) -> dict:
        """
        Fetch a paginated list of alerts from the service API.

        Args:
            page_num (int): The page number to retrieve.
            items_per_page (int): The number of alerts to retrieve per page.

        Returns:
            dict: A dictionary containing the paginated list of alerts and metadata such as total count.
        """
        params = assign_params(pageNum=page_num, itemsPerPage=items_per_page)
        results = self._http_request(method="GET", url_suffix=f"/api/atlas/v2/groups/{self.group_id}/alerts", params=params)
        return results

    def get_events_request(self, page_num: int = None, items_per_page: int = ITEMS_PER_PAGE, min_date: str = None) -> dict:
        """
        Fetch a paginated list of events from the service API.

        Args:
            page_num (int): The page number to retrieve.
            items_per_page (int): The number of events to retrieve per page.
            min_date (str): The minimum date to retrieve.

        Returns:
            dict: A dictionary containing the paginated list of events and metadata such as total count.
        """
        params = assign_params(pageNum=page_num, itemsPerPage=items_per_page, minDate=min_date)
        results = self._http_request(method="GET", url_suffix=f"/api/atlas/v2/groups/{self.group_id}/events", params=params)
        return results

    def get_response_from_page_link(self, page_link: str) -> dict:
        """
        Sends an HTTP GET request to fetch data using `page_link`.

        Args:
            page_link (str): The full URL to the specific data page to fetch.

        Returns:
            The API response containing the data retrieved from the provided page link.
        """
        results = self._http_request(method="GET", full_url=page_link)
        return results

    def get_events_first_run(self, fetch_limit: int) -> list:
        """
        Fetches events from multiple pages, ensuring that the total number of fetched events does not exceed the specified
        `fetch_limit`.

        Args:
            fetch_limit (int): The maximum number of events to fetch.

        Returns:
            list: A list of events, truncated to the `fetch_limit` if necessary.
        """
        results: List[dict] = []
        items_per_page = min(fetch_limit, ITEMS_PER_PAGE)
        page_num = 1

        while len(results) < fetch_limit:
            page_results = self.get_events_request(page_num=page_num, items_per_page=items_per_page).get("results")
            if not page_results:
                break

            results.extend(page_results)

            if len(results) >= fetch_limit:
                return results[:fetch_limit]

            items_per_page = min(fetch_limit - len(results), ITEMS_PER_PAGE)
            page_num += 1

        return results


""" HELPER FUNCTIONS """


################ ALERTS AND EVENTS FUNCTIONS ################


def add_entry_status_field(event: dict) -> None:
    """
    Adds a _ENTRY_STATUS field to an event by checking the event status.

    Args:
        event (dict): The event.
    """
    updated_str: str = str(event.get("updated"))
    created_str: str = str(event.get("created"))

    updated = datetime.strptime(updated_str, DATE_FORMAT)
    created = datetime.strptime(created_str, DATE_FORMAT)

    if updated == created:
        event["_ENTRY_STATUS"] = "new"
    elif updated > created:
        event["_ENTRY_STATUS"] = "updated"


def get_page_url(links: list, page_type: str) -> str:
    """
    Retrieves the URL for the next or previous page from a list of links.

    Args:
        links (list): A list of dictionaries, each representing a link with "rel" and "href" keys.
        page_type (str): The type of the page, can be 'next', 'previous' or 'self'.
    Returns:
        str or empty string: The URL for the next or previous page if found, otherwise None.
    """
    for link in links:
        if link.get("rel") == page_type:
            return link.get("href")
    return ""


def add_time_field(event: dict) -> None:
    """
    Adds a `_time` field to an event based on its updated or created time.

    Args:
        event (dict): The event dictionary to add the `_time` field to. If the event has an 'updated' time,
         `_time` will be set to this value; otherwise, it will default to the 'created' time.
    """
    if event.get("updated"):
        event["_time"] = event.get("updated")
    else:
        event["_time"] = event.get("created")


def enrich_event(event: dict, event_type: str) -> None:
    """
    Enriches each event with additional information based on its type.

    Args:
        event (dict): The event dictionary to enrich.
        event_type (str): The type of event ('alerts' or 'events').
    """
    event["source_log_type"] = event_type
    add_time_field(event)
    if event_type == "alerts":
        add_entry_status_field(event)


################ ALERTS FUNCTIONS ################


def remove_alerts_by_ids(alerts: list, ids: list) -> list:
    """
    Removes alerts from a list based on specified IDs.

    Args:
        alerts (list): A list of alert dictionaries, each containing an "id" key.
        ids (list): A list of IDs to be removed from the alerts list.

    Returns:
        list: A filtered list of alerts excluding any alerts with IDs in the provided ids list.
    """
    results = []
    for alert in alerts:
        if alert.get("id") not in ids:
            results.append(alert)
    return results


def create_last_run_dict_for_alerts(links: list, last_page_alerts_ids: list) -> dict:
    """
    Updates the last_run dictionary with the current page's self URL and event IDs.

    Args:
        links (list): Current page links for pagination.
        last_page_alerts_ids (list): IDs of alerts on the last processed page.

    Returns:
        dict: Updated last_run dictionary.
    """
    return {"page_link": get_page_url(links, page_type="self"), "last_page_alerts_ids": last_page_alerts_ids}


def fetch_alerts_command(client: Client, fetch_limit: int, last_run: dict) -> tuple[list, dict]:
    """
    Fetches alerts until fetch_limit is reached, or no more alerts are available.

    Args:
        client (Client): MongoDBAtlasEventCollector client.
        fetch_limit: The maximum number of alerts to fetch.
        last_run (dict): Dictionary containing data from the previous run.

    Returns:
        - The list containing all fetched events.
        - The object to save for the next run.
    """

    demisto.debug("Start to fetch alerts")
    page_link = last_run.get("page_link", "")

    if page_link:
        demisto.debug(f"Getting a response from page {page_link}")
        response = client.get_response_from_page_link(page_link)
    else:
        demisto.debug("Initialize the first page")
        response = client.get_alerts_request(page_num=1, items_per_page=ITEMS_PER_PAGE)

    links = response.get("links", [])
    results = response.get("results", [])

    last_page_alerts_ids = last_run.get("last_page_alerts_ids", [])

    alerts = remove_alerts_by_ids(results, last_page_alerts_ids)
    demisto.debug(f"Those are the events ids from the last run {last_page_alerts_ids}")
    output = []

    while True:
        for alert in alerts:
            enrich_event(alert, event_type="alerts")
            output.append(alert)
            last_page_alerts_ids.append(alert.get("id"))

            if len(output) >= fetch_limit:
                last_run_new_dict = create_last_run_dict_for_alerts(links, last_page_alerts_ids)
                demisto.debug(f"The limit is reached. Amount of fetched alerts is {len(output)}")
                return output, last_run_new_dict

        next_url = get_page_url(links, page_type="next")
        if next_url:
            response = client.get_response_from_page_link(next_url)
            alerts = response.get("results", [])
            links = response.get("links", [])
            last_page_alerts_ids.clear()
        else:
            break

    demisto.debug(f"No more pages left to fetch. Total alerts fetched: {len(output)}")
    last_run_new_dict = create_last_run_dict_for_alerts(links, last_page_alerts_ids)
    return output, last_run_new_dict


################ EVENTS FUNCTIONS ################


def get_latest_date(date1: str, date2: str) -> str:
    """
    Compares two dates and returns the latest one.

    Args:
        date1 (str): The first date in string format.
        date2 (str): The second date in string format.

    Returns:
        str: The latest date in string format.
    """
    if not date1:
        return date2
    if not date2:
        return date1

    dt1 = datetime.strptime(date1, DATE_FORMAT)
    dt2 = datetime.strptime(date2, DATE_FORMAT)

    return date1 if dt1 >= dt2 else date2


def get_last_page_of_events(client: Client, results: dict) -> dict:
    """
    Iterates through paginated event data, following "next" links provided in the API response until it reaches the last page.

    Args:
        client (Client): The API client instance used to retrieve the event data.
        results (dict): The initial response dictionary containing event data and pagination links.

    Returns:
        dict: The final page of events retrieved from the API.
    """
    links = results.get("links", [])
    next_url = get_page_url(links, page_type="next")
    last_response = results

    while next_url:
        last_response = client.get_response_from_page_link(next_url)
        links = last_response.get("links")
        next_url = get_page_url(links, page_type="next")

    return last_response


def save_events_ids_with_specific_created_date(events: list, created_date: str) -> list:
    """
    Filters event IDs by a given creation date.

    Args:
        events (list): List of event dictionaries with "id" and "created" date fields.
        created_date (str): Target creation date to filter by.

    Returns:
        list: IDs of events matching the specified creation date.
    """
    results = []
    for event in events:
        if event.get("created") == created_date:
            results.append(event.get("id"))
    return results


def create_last_run_dict_for_events(output: list, new_min_time: str) -> dict:
    """
    Creates a dictionary to store the last run information for events.

    Args:
        output (list): List of event dictionaries.
        new_min_time (str): The minimum time to filter events by.

    Returns:
        dict: A dictionary with `min_time` and `events_with_created_min_time` keys.
    """
    events_with_created_min_time = save_events_ids_with_specific_created_date(output, new_min_time)
    return {"min_time": new_min_time, "events_with_created_min_time": events_with_created_min_time}


def first_time_fetching_events(client: Client, fetch_limit: int) -> tuple[list, str]:
    """
    Fetches and enriches the first batch of events, returning them with the minimum creation time.

    Args:
        client (Client): The client instance used to retrieve events.
        fetch_limit (int): The maximum number of events to retrieve.

    Returns:
        tuple: A tuple containing:
            - results (list): A list of event dictionaries.
            - new_min_time (str): The creation time of the latest fetched event.
    """
    results = client.get_events_first_run(fetch_limit)
    for event in results:
        enrich_event(event, event_type="events")
    last_fetched_event = results[0] if results else None
    new_min_time = last_fetched_event.get("created") if last_fetched_event else None
    return results, new_min_time  # type: ignore[return-value]


def fetch_events_command(client: Client, fetch_limit: int, last_run: dict) -> tuple[list, dict]:
    """
    Fetches events until fetch_limit is reached, or no more events are available.

    Args:
        client (Client): MongoDBAtlasEventCollector client.
        fetch_limit (int): The maximum number of events to fetch.
        last_run (dict): Dictionary containing data from the previous run.

    Returns:
        - The list containing all fetched events.
        - The object to save for the next run.
    """
    min_date = last_run.get("min_time")
    events_with_created_min_time = last_run.get("events_with_created_min_time") or []

    if not min_date:  # first time fetching events
        output, new_min_time = first_time_fetching_events(client, fetch_limit)
        new_last_run_obj = create_last_run_dict_for_events(output, new_min_time)
        return output, new_last_run_obj

    demisto.debug(f"Start to fetch events with {min_date}")

    results: dict = client.get_events_request(min_date=min_date)
    response = get_last_page_of_events(client, results)
    links = response.get("links", [])
    events = response.get("results", [])

    output = []
    new_min_date = min_date

    while True:
        for event in reversed(events):
            event_id = event.get("id")
            if event_id in events_with_created_min_time:
                continue

            enrich_event(event, "events")
            output.append(event)
            new_min_date = get_latest_date(new_min_date, event.get("created"))

            if len(output) >= fetch_limit:
                demisto.debug(f"Fetch limit reached. Total events fetched: {len(output)}")
                new_last_run_obj = create_last_run_dict_for_events(output, new_min_date)
                return output, new_last_run_obj

        previous_page = get_page_url(links, page_type="previous")
        if previous_page:
            response = client.get_response_from_page_link(previous_page)
            events = response.get("results", [])
            links = response.get("links", [])
        else:
            break

    demisto.debug(f"No more events left to fetch. Total events fetched: {len(output)}")
    new_last_run_obj = create_last_run_dict_for_events(output, new_min_date)
    return output, new_last_run_obj


""" COMMAND FUNCTIONS """


def test_module(client: Client) -> str:
    """
    Returning 'ok' indicates that the integration works like it suppose to. Connection to the service is successful.

    Args:
        client: MongoDB Atlas client

    Returns:
        'ok' if test passed, anything else will fail the test
    """
    fetch_events_command(client=client, fetch_limit=1, last_run={})
    return "ok"


def fetch_events(client: Client, fetch_limit: int) -> tuple[list, dict]:
    last_run = demisto.getLastRun()

    alerts_output, last_run_alerts = fetch_alerts_command(client, fetch_limit, last_run)
    events_output, last_run_events = fetch_events_command(client, fetch_limit, last_run)

    last_run_new_obj = {**last_run_alerts, **last_run_events}
    return (alerts_output + events_output), last_run_new_obj


def get_events(client: Client, args) -> tuple[list, CommandResults]:
    fetch_limit = int(args.get("limit"))

    output, _ = fetch_events(client, fetch_limit)

    filtered_events = []
    for event in output:
        filtered_event = {
            "ID": event.get("id"),
            "Event Type": event.get("source_log_type"),
            "Time": event.get("_time"),
            "Created": event.get("created"),
        }
        filtered_events.append(filtered_event)

    human_readable = tableToMarkdown(name="MongoDB Atlas Events", t=filtered_events, removeNull=True)
    command_results = CommandResults(
        readable_output=human_readable,
        outputs=output,
        outputs_prefix="MongoDBAtlasEventCollector",
    )
    return output, command_results


""" MAIN FUNCTION """


def main() -> None:  # pragma: no cover
    """main function, parses params and runs command functions"""
    params = demisto.params()
    args = demisto.args()
    command = demisto.command()

    demisto.debug(f"MongoDB Command being called is {demisto.command()}")
    try:
        credentials = params.get("credentials", {})
        public_key = credentials.get("identifier")
        private_key = credentials.get("password")
        group_id = params.get("group_id")

        base_url = params.get("url")
        verify = not params.get("insecure", False)
        fetch_limit = arg_to_number(params.get("max_events_per_fetch")) or DEFAULT_FETCH_LIMIT

        client = Client(base_url=base_url, verify=verify, public_key=public_key, private_key=private_key, group_id=group_id)

        if command == "test-module":
            result = test_module(client)
            return_results(result)
        elif command == "mongo-db-atlas-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)
        elif command == "fetch-events":
            events, last_run_new_obj = fetch_events(client, fetch_limit)
            if events:
                demisto.debug(f"Sending {len(events)} events to Cortex XSIAM")
                send_events_to_xsiam(events=events, vendor=VENDOR, product=PRODUCT)
                demisto.setLastRun(last_run_new_obj)
                demisto.debug(f"Successfully saved last_run= {demisto.getLastRun()}")

    except Exception as e:
        return_error(f"Failed to execute {demisto.command()} command.\nError:\n{str(e)}")


""" ENTRY POINT """

if __name__ in ("__main__", "__builtin__", "builtins"):
    main()