AWS Security Hub Event Collector

import datetime as dt
from typing import TYPE_CHECKING, cast

import demistomock as demisto  # noqa: F401
from AWSApiModule import *
from CommonServerPython import *  # noqa: F401

# The following import are used only for type hints and autocomplete.
# It is not used at runtime, and not exist in the docker image.
if TYPE_CHECKING:
    from mypy_boto3_securityhub import SecurityHubClient
    from mypy_boto3_securityhub.type_defs import AwsSecurityFindingTypeDef


VENDOR = "AWS"
PRODUCT = "Security Hub"
TIME_FIELD = "CreatedAt"
DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
DATETIME_FORMAT_NO_MS = "%Y-%m-%dT%H:%M:%SZ"  # Fallback format without milliseconds
DEFAULT_FIRST_FETCH = "3 days"
DEFAULT_MAX_RESULTS = 1000
API_MAX_PAGE_SIZE = 100  # The API only allows a maximum of 100 results per request. Using more raises an error.
RETRY_LIMIT_INCREMENT = 100  # How much to increase limit when AWS returns all events as duplicates
MAX_AWS_LIMIT = 10000  # Maximum limit to prevent excessive API calls


def parse_aws_timestamp(timestamp_str: str) -> dt.datetime:
    """
    Parse AWS timestamp with flexible format support.
    The AWS API returns timestamps in two formats:
    - With milliseconds: "2023-01-01T00:00:00.000Z"
    - Without milliseconds: "2023-01-01T00:00:00Z"

    Args:
        timestamp_str (str): Timestamp string from AWS API

    Returns:
        datetime: Parsed datetime object

    Raises:
        ValueError: If timestamp format is not supported
    """
    try:
        # Try parsing with milliseconds first
        return dt.datetime.strptime(timestamp_str, DATETIME_FORMAT)
    except ValueError:
        # Fallback to parsing without milliseconds
        return dt.datetime.strptime(timestamp_str, DATETIME_FORMAT_NO_MS)


def generate_last_run(events: list["AwsSecurityFindingTypeDef"], previous_last_run: dict | None = None) -> dict[str, Any]:
    """
    Generate the last run object using events data.

    Args:
        events (list[dict]): List of events to generate the last run object from.
        previous_last_run (dict | None, optional): Previous last run data for smart ignore list accumulation.

    Note:
        Since the time filters seem to be equal or greater than (which results in duplicate from the last run),
        we add findings that are equal to 'last_finding_update_time' and filter them out in the next fetch.
        Smart accumulation: If timestamp hasn't changed, accumulate ignore list instead of replacing it.

    Returns:
        dict: Last run object.
    """
    last_update_date = events[-1].get(TIME_FIELD)

    # Smart ignore list accumulation logic
    if previous_last_run and previous_last_run.get("last_update_date") == last_update_date:
        # Same timestamp - ACCUMULATE ignore list from previous run
        ignore_list: list[str] = previous_last_run.get("last_update_date_finding_ids", []).copy()
        demisto.debug(
            f"Same timestamp detected ({last_update_date}). Accumulating ignore list from {len(ignore_list)} existing IDs."
        )
    else:
        # New timestamp - START FRESH ignore list
        ignore_list = []
        if previous_last_run:
            demisto.debug(
                f"Timestamp changed from {previous_last_run.get('last_update_date')} to {last_update_date}. "
                f"Starting fresh ignore list."
            )

    # Since the "_time" key is added to each event, the event type changes from "AwsSecurityFindingTypeDef" to just dict
    events = cast(list[dict[str, Any]], events)
    for event in events:
        event["_time"] = event[TIME_FIELD]

        if event[TIME_FIELD] == last_update_date:
            ignore_list.append(event["Id"])

    demisto.debug(f"Generated ignore list with {len(ignore_list)} IDs for timestamp {last_update_date}")

    return {
        "last_update_date": last_update_date,
        "last_update_date_finding_ids": ignore_list,
    }


def get_events(
    client: "SecurityHubClient",
    start_time: dt.datetime | None = None,
    end_time: dt.datetime | None = None,
    id_ignore_list: list[str] | None = None,
    page_size: int = API_MAX_PAGE_SIZE,
    limit: int = 0,
    start_token: str | None = None,
) -> tuple[List["AwsSecurityFindingTypeDef"], str | None]:
    """
    Fetch events from AWS Security Hub.

    Args:
        client (SecurityHubClient): Boto3 client to use.
        start_time (datetime | None, optional): Start time to fetch events from. Required if end_time is set.
        end_time (datetime | None, optional): Time to fetch events until. Defaults to current time.
        id_ignore_list (list[str] | None, optional): List of finding IDs to not include in the results.
            Defaults to None.
        page_size (int, optional): Number of results to fetch per request. Defaults to API_MAX_PAGE_SIZE.
        limit (int, optional): Maximum number of results to fetch. Defaults to 0.
        start_token (str | None, optional): Token to use for pagination. Defaults to None.

    Returns:
        tuple: A 2-tuple containing:
            - Filtered events (list[AwsSecurityFindingTypeDef])
            - NextToken for continuation (str | None, None means no more events)
    """
    kwargs: dict = {"SortCriteria": [{"Field": TIME_FIELD, "SortOrder": "asc"}]}
    filters: dict = {}

    # Start from provided token if available
    if start_token:
        kwargs["NextToken"] = start_token

    if end_time and not start_time:
        raise ValueError("start_time must be set if end_time is used.")

    if start_time:
        filters[TIME_FIELD] = [
            {
                "Start": start_time.strftime(DATETIME_FORMAT),
                "End": end_time.strftime(DATETIME_FORMAT) if end_time else dt.datetime.now().strftime(DATETIME_FORMAT),
            }
        ]

    if id_ignore_list:
        id_ignore_set = set(id_ignore_list)
    else:
        id_ignore_set = set()

    if filters:
        # We send kwargs because passing Filters=None to get_findings() tries to use a None value for filters,
        # which raises an error.
        kwargs["Filters"] = filters

    # Collect all raw events first, then filter duplicates at the end
    all_raw_events: list[AwsSecurityFindingTypeDef] = []
    total_collected_count = 0
    pagination_iteration = 0

    demisto.debug(
        f"Starting get_events pagination with limit={limit}, page_size={page_size}, ignore_set_size={len(id_ignore_set)}"
    )

    while True:
        pagination_iteration += 1

        if limit and total_collected_count + page_size > limit:
            kwargs["MaxResults"] = limit - total_collected_count
        else:
            kwargs["MaxResults"] = page_size

        demisto.debug(f"Iteration {pagination_iteration}: Calling get_findings with kwargs: {kwargs}")
        response = client.get_findings(**kwargs)
        result = response.get("Findings", [])
        has_next_token = "NextToken" in response

        demisto.debug(
            f"Iteration {pagination_iteration}: AWS returned {len(result)} findings, NextToken present: {has_next_token}"
        )
        all_raw_events.extend(result)
        total_collected_count += len(result)

        demisto.debug(f"Iteration {pagination_iteration}: Total raw events collected so far: {len(all_raw_events)}")

        # Check if we should continue pagination
        should_continue = (
            has_next_token
            and (limit == 0 or total_collected_count < limit)
            and len(result) > 0  # Stop if AWS returns empty results
        )

        if should_continue:
            kwargs["NextToken"] = response["NextToken"]
            demisto.debug(
                f"Iteration {pagination_iteration}: Continuing pagination (limit={limit}, collected={total_collected_count})"
            )
        else:
            if not has_next_token:
                demisto.debug(f"Iteration {pagination_iteration}: No more pages available from AWS")
            elif limit > 0 and total_collected_count >= limit:
                demisto.debug(f"Iteration {pagination_iteration}: Reached limit ({total_collected_count}/{limit})")
            elif len(result) == 0:
                demisto.debug(f"Iteration {pagination_iteration}: AWS returned empty results, stopping pagination")
            break

    # Now filter duplicates from ALL collected events
    demisto.debug(f"Filtering duplicates from {len(all_raw_events)} total collected events")
    filtered_events = [event for event in all_raw_events if event["Id"] not in id_ignore_set]
    duplicates_filtered = len(all_raw_events) - len(filtered_events)

    demisto.debug(
        f"Pagination completed after {pagination_iteration} iterations. "
        f"Raw events: {len(all_raw_events)}, After deduplication: {len(filtered_events)} "
        f"({duplicates_filtered} duplicates filtered)"
    )

    # Return filtered events and the final NextToken
    final_next_token = response.get("NextToken")
    return (filtered_events, final_next_token)


def fetch_events(
    client: "SecurityHubClient",
    last_run: dict,
    first_fetch_time: dt.datetime | None,
    page_size: int = API_MAX_PAGE_SIZE,
    limit: int = 0,
) -> tuple[list["AwsSecurityFindingTypeDef"], dict, Exception | None]:
    """
    Fetch events from AWS Security Hub and send them to XSIAM.

    Args:
        client (SecurityHubClient): Boto3 client to use.
        last_run (dict): Dict containing the last fetched event creation time.
        first_fetch_time (datetime | None, optional): In case of first fetch, fetch events from this datetime.
        page_size (int, optional): Number of results to fetch per request. Defaults to API_MAX_PAGE_SIZE.
        limit (int, optional): Maximum number of events to fetch. Defaults to 0 (no limit).
    """
    demisto.debug(f"Fetching events with last_run: {last_run}")
    if last_run.get("last_update_date"):
        start_time = parse_aws_timestamp(last_run["last_update_date"])

    else:
        start_time = cast(dt.datetime, first_fetch_time)

    id_ignore_list: list = last_run.get("last_update_date_finding_ids", [])

    events: list[AwsSecurityFindingTypeDef] = []
    error = None
    current_limit = limit
    max_aws_limit = MAX_AWS_LIMIT

    # Fix end_time for all additional retries to keep NextToken valid
    end_time = dt.datetime.now()

    try:
        continuation_token = None
        while True:
            got_new_events = False

            filtered_events, continuation_token = get_events(
                client=client,
                start_time=start_time,
                end_time=end_time,
                id_ignore_list=id_ignore_list,
                page_size=page_size,
                limit=current_limit,
                start_token=continuation_token,
            )

            if filtered_events:  # If any events returned, they're new (already filtered)
                got_new_events = True
            events.extend(filtered_events)

            # If we got new events, we're done!
            if got_new_events:
                demisto.debug(f"Successfully fetched {len(events)} new events")
                break

            # No new events - should we increase the requested limit in order to exit the loop?
            if continuation_token and current_limit < max_aws_limit:
                # AWS has more events and we haven't hit the max API limit - continue from NextToken with a small amount
                current_limit = RETRY_LIMIT_INCREMENT
                demisto.info(
                    f"Infinite loop prevention: All events were duplicates but AWS has more events. "
                    f"Retrying with limit: {current_limit} (continuing from NextToken)"
                )
                continue

            # Can't or shouldn't retry - explain why and exit
            if not continuation_token:
                demisto.debug("All events were duplicates and AWS has no more events - reached end")
            else:
                demisto.debug(f"All events were duplicates but reached max limit ({max_aws_limit}) - stopping retries")
            break

    except Exception as e:
        demisto.error(f"Error while fetching events. Events fetched so far: {len(events)}. Error: {e}")
        error = e

    # --- Set next_run data ---
    if events:
        demisto.info(f"Fetched {len(events)} findings.")
        next_run = generate_last_run(events, last_run)
        demisto.info(f"Last run data updated to: {next_run}.")

    else:
        demisto.info("No new findings were found.")
        next_run = last_run

    return events, next_run, error


def get_events_command(
    client: "SecurityHubClient",
    should_push_events: bool,
    page_size: int,
    limit: int = 0,
    start_time: dt.datetime | None = None,
    end_time: dt.datetime | None = None,
) -> CommandResults:
    """
    Fetch events from AWS Security Hub.

    Args:
        client (SecurityHubClient): Boto3 client to use.
        should_push_events (bool): Whether to push events to XSIAM.
        page_size (int, optional): Number of results to fetch per request. Defaults to API_MAX_PAGE_SIZE.
        limit (int, optional): Maximum number of events to fetch. Defaults to 0 (no limit).
        start_time (dt.datetime, optional): Start time for filtering events. Defaults to None.
        end_time (dt.datetime, optional): End time for filtering events. Defaults to None.

    Returns:
        CommandResults: CommandResults object containing the events.
    """
    events, _ = get_events(client=client, page_size=page_size, limit=limit, start_time=start_time, end_time=end_time)

    if should_push_events:
        send_events_to_xsiam(events, vendor=VENDOR, product=PRODUCT)

    return CommandResults(
        readable_output=tableToMarkdown("AWS Security Hub Events", events, sort_headers=False),
    )


def main():  # pragma: no cover
    params = demisto.params()
    args = demisto.args()
    command = demisto.command()

    aws_role_arn = params.get("role_arn")
    aws_role_session_name = params.get("role_session_name")
    aws_default_region = params.get("default_region")
    aws_role_session_duration = params.get("role_session_duration")
    aws_access_key_id = demisto.get(params, "credentials.identifier")
    aws_secret_access_key = demisto.get(params, "credentials.password")
    verify_certificate = not params.get("insecure", True)
    timeout = params.get("timeout")
    retries = params.get("retries", 5)

    limit = arg_to_number(params.get("max_fetch")) or DEFAULT_MAX_RESULTS

    if limit <= 0:
        raise ValueError("Max fetch value cannot be lower than 1.")

    # How much time before the first fetch to retrieve events
    first_fetch_time = arg_to_datetime(
        arg=params.get("first_fetch", DEFAULT_FIRST_FETCH), arg_name="First fetch time", required=True
    )

    try:
        validate_params(
            aws_default_region=aws_default_region,
            aws_role_arn=aws_role_arn,
            aws_role_session_name=aws_role_session_name,
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key,
        )

        aws_client = AWSClient(
            aws_default_region=aws_default_region,
            aws_role_arn=aws_role_arn,
            aws_role_session_name=aws_role_session_name,
            aws_role_session_duration=aws_role_session_duration,
            aws_role_policy=None,
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key,
            verify_certificate=verify_certificate,
            timeout=timeout,
            retries=retries,
        )

        client: SecurityHubClient = aws_client.aws_session(
            service="securityhub",
            region=aws_default_region,
            role_arn=aws_role_arn,
            role_session_name=aws_role_session_name,
            role_session_duration=aws_role_session_duration,
        )

        demisto.info(f'Executing "{command}" command...')

        if command == "test-module":
            get_events(client=client, limit=1)[0]  # Get events from tuple
            return_results("ok")

        elif command == "aws-securityhub-get-events":
            should_push_events = argToBoolean(args.get("should_push_events", False))
            page_size = arg_to_number(args.get("page_size", API_MAX_PAGE_SIZE))

            if page_size is None or page_size > 100:
                raise ValueError("Page size cannot be larger than 100 (not supported by the API).")

            limit = arg_to_number(args.get("limit")) or 1

            if limit is None or limit <= 0:
                raise ValueError("Max fetch value cannot be lower than 1.")

            # Parse optional start_time and end_time arguments
            start_time = arg_to_datetime(args.get("start_time")) if args.get("start_time") else None
            end_time = arg_to_datetime(args.get("end_time")) if args.get("end_time") else None

            return_results(
                get_events_command(
                    client=client,
                    should_push_events=should_push_events,
                    page_size=page_size,
                    limit=limit,
                    start_time=start_time,
                    end_time=end_time,
                )
            )

        elif command == "fetch-events":
            events, next_run, error = fetch_events(
                client=client,
                last_run=demisto.getLastRun(),
                first_fetch_time=first_fetch_time,
                limit=limit,
            )

            send_events_to_xsiam(events=events, vendor=VENDOR, product=PRODUCT)
            demisto.setLastRun(next_run)

            if error and events:
                raise Exception(
                    f"An error occurred while running fetch-events. "
                    f"The operation was partially successful, but failed midway.\n"
                    f"A total of {len(events)} events were successfully fetched "
                    f"before the error occurred."
                ) from error

            elif error:
                raise error

        else:
            raise NotImplementedError(f'Command "{command}" is not implemented.')

    # Log exceptions and return errors
    except Exception as e:
        return_error(f"Failed to execute {command} command.\nError:\n{e}")


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