AWS Security Hub Event Collector

from pathlib import Path
import datetime as dt

import pytest
from freezegun import freeze_time
from AWSSecurityHubEventCollector import *

# Import test findings data from dedicated test data file
from test_data.test_constants import SAMPLE_FINDINGS, IGNORED_FINDINGS

# Test configuration constants - moved to top for parametrized test
FIXED_END_TIME = dt.datetime(2023, 1, 1, 13, 0, 0)
FIRST_FETCH_TIME = dt.datetime(2023, 1, 1, 10, 0, 0)
TEST_START_TIME = dt.datetime(2023, 1, 1, 10, 0, 0)
TEST_END_TIME = dt.datetime(2023, 1, 1, 14, 0, 0)

# Common ignore list patterns used in multiple tests
STANDARD_IGNORE_LIST = ["ignore-me", "also-ignore"]
DUPLICATE_IGNORE_LIST = ["ignore-1", "ignore-2"]


def load_test_data(folder: str, file_name: str) -> dict:
    """
    A function for loading and returning data from json files within the "test_data" folder.

    Args:
        folder (str): Name of the parent folder of the file within `test_data`.
        file_name (str): Name of a json file to load data from.

    Returns:
        dict: Dictionary data loaded from the json file.
    """
    with open(Path("test_data") / folder / f"{file_name}.json") as f:
        return json.load(f)


class MockClient:
    """
    A client mocking Boto3.client
    """

    findings_data = load_test_data("api_mock", "get_findings_10")
    calls_count = 0
    last_index = 0  # Used to mock pagination

    def __init__(self, *args, **kwargs):
        pass

    def get_findings(self, **kwargs) -> dict:
        """
        A function for mocking the get_findings function.

        Args:
            kwargs: Keyword arguments that would be passed to the get_findings function.
        """
        max_index = self.last_index + kwargs["MaxResults"]

        return_events = self.findings_data[self.last_index : max_index]
        self.last_index += len(return_events)
        result = {"Findings": return_events}

        if self.last_index < len(self.findings_data):
            result["NextToken"] = "next_token"

        self.calls_count += 1
        return result

    def reset(self):
        """
        A function for resetting the mock client.
        """
        self.calls_count = 0
        self.last_index = 0


@pytest.fixture
def client():
    """
    A fixture for creating a mock client.
    """
    return MockClient()


@pytest.mark.parametrize(
    "page_size, limit, expected_api_calls_count, expected_output_file",
    [
        (100, 1000, 1, "fetch_events_expected_results_0"),
        (2, 1000, 5, "fetch_events_expected_results_0"),
        (3, 1000, 4, "fetch_events_expected_results_0"),
        (2, 5, 3, "fetch_events_expected_results_1"),
        (100, 5, 1, "fetch_events_expected_results_1"),
    ],
)
def test_fetch(client, page_size: int, limit: int, expected_api_calls_count: int, expected_output_file: str):
    """
    Given: A page size parameter for the fetch events function.
    When: Fetching events from the API.
    Then: Assert the returned events are valid, and the number of API calls is as expected.

    Note: This is a test for both 'fetch_events' and 'get_events_command' functions.
    """
    expected_output = load_test_data("expected_results", expected_output_file)

    first_fetch_time = dt.datetime(2021, 1, 1)
    events, _, _ = fetch_events(client=client, last_run={}, first_fetch_time=first_fetch_time, page_size=page_size, limit=limit)

    assert client.calls_count == expected_api_calls_count
    assert len(events) == len(expected_output)
    assert events == expected_output

    client.reset()

    result = get_events_command(
        client=client, should_push_events=False, page_size=page_size, limit=limit, start_time=None, end_time=None
    )

    assert client.calls_count == expected_api_calls_count
    assert result.readable_output == tableToMarkdown("AWS Security Hub Events", expected_output, sort_headers=False)


class TestParseAwsTimestamp:
    """Test the new flexible timestamp parsing functionality."""

    @pytest.mark.parametrize(
        "timestamp_str, expected_result",
        [
            # Valid timestamp with milliseconds
            ("2023-01-01T12:30:45.123Z", dt.datetime(2023, 1, 1, 12, 30, 45, 123000)),
            # Valid timestamp without milliseconds
            ("2023-01-01T12:30:45Z", dt.datetime(2023, 1, 1, 12, 30, 45)),
        ],
    )
    def test_parse_timestamp_valid_formats(self, timestamp_str, expected_result):
        """
        Given: A valid timestamp string from AWS API (with or without milliseconds).
        When: Parsing the timestamp using parse_aws_timestamp.
        Then: Should return the correct datetime object.
        """
        result = parse_aws_timestamp(timestamp_str)
        assert result == expected_result

    def test_parse_timestamp_invalid_format(self):
        """
        Given: An invalid timestamp string.
        When: Parsing the timestamp using parse_aws_timestamp.
        Then: Should raise ValueError.
        """
        invalid_timestamp = "2023/01/01 12:30:45"
        with pytest.raises(ValueError):
            parse_aws_timestamp(invalid_timestamp)


class TestGenerateLastRunWithSmartAccumulation:
    """Test the smart ignore list accumulation functionality."""

    def test_generate_last_run_same_timestamp_accumulates_ignore_list(self):
        """
        Given: Events with same timestamp as previous run.
        When: Generating last run with previous_last_run parameter.
        Then: Should accumulate ignore list instead of replacing it.
        """
        events = [
            {"CreatedAt": "2023-01-01T12:00:00.000Z", "Id": "finding-3", "_time": "2023-01-01T12:00:00.000Z"},
            {"CreatedAt": "2023-01-01T12:00:00.000Z", "Id": "finding-4", "_time": "2023-01-01T12:00:00.000Z"},
        ]
        previous_last_run = {
            "last_update_date": "2023-01-01T12:00:00.000Z",
            "last_update_date_finding_ids": ["finding-1", "finding-2"],
        }

        result = generate_last_run(events, previous_last_run)

        assert result["last_update_date"] == "2023-01-01T12:00:00.000Z"
        assert len(result["last_update_date_finding_ids"]) == 4
        assert "finding-1" in result["last_update_date_finding_ids"]
        assert "finding-2" in result["last_update_date_finding_ids"]
        assert "finding-3" in result["last_update_date_finding_ids"]
        assert "finding-4" in result["last_update_date_finding_ids"]

    def test_generate_last_run_new_timestamp_starts_fresh(self):
        """
        Given: Events with new timestamp different from previous run.
        When: Generating last run with previous_last_run parameter.
        Then: Should start fresh ignore list instead of accumulating.
        """
        events = [
            {"CreatedAt": "2023-01-01T13:00:00.000Z", "Id": "finding-3", "_time": "2023-01-01T13:00:00.000Z"},
            {"CreatedAt": "2023-01-01T13:00:00.000Z", "Id": "finding-4", "_time": "2023-01-01T13:00:00.000Z"},
        ]
        previous_last_run = {
            "last_update_date": "2023-01-01T12:00:00.000Z",
            "last_update_date_finding_ids": ["finding-1", "finding-2"],
        }

        result = generate_last_run(events, previous_last_run)
        assert result["last_update_date"] == "2023-01-01T13:00:00.000Z"
        assert len(result["last_update_date_finding_ids"]) == 2
        assert "finding-3" in result["last_update_date_finding_ids"]
        assert "finding-4" in result["last_update_date_finding_ids"]
        assert "finding-1" not in result["last_update_date_finding_ids"]
        assert "finding-2" not in result["last_update_date_finding_ids"]

    def test_generate_last_run_no_previous_run(self):
        """
        Given: Events and no previous run data
        When: Generating last run
        Then: Should create fresh ignore list
        """
        # Given
        events = [
            {"CreatedAt": "2023-01-01T12:00:00.000Z", "Id": "finding-1", "_time": "2023-01-01T12:00:00.000Z"},
            {"CreatedAt": "2023-01-01T12:00:00.000Z", "Id": "finding-2", "_time": "2023-01-01T12:00:00.000Z"},
        ]

        # When
        result = generate_last_run(events, None)

        # Then
        assert result["last_update_date"] == "2023-01-01T12:00:00.000Z"
        assert len(result["last_update_date_finding_ids"]) == 2
        assert "finding-1" in result["last_update_date_finding_ids"]
        assert "finding-2" in result["last_update_date_finding_ids"]


@pytest.mark.parametrize(
    (
        "scenario, mock_api_responses, start_time, end_time, id_ignore_list, "
        "page_size, limit, start_token, expected_events_count, expected_next_token, expected_api_calls, expected_debug_messages"
    ),
    [
        (
            "single_page_no_pagination",
            # Mock API returns findings in single call, no NextToken
            [{"Findings": SAMPLE_FINDINGS[:3]}],
            TEST_START_TIME,  # start_time
            TEST_END_TIME,  # end_time
            None,  # id_ignore_list
            50,  # page_size
            0,  # limit (no limit)
            None,  # start_token
            3,  # expected_events_count
            None,  # expected_next_token
            1,  # expected_api_calls (single page)
            [
                "Starting get_events pagination",
                "No more pages available from AWS",
                "Pagination completed after",
            ],  # expected_debug_messages
        ),
        (
            "multiple_pages_with_pagination",
            # Mock API returns findings across multiple calls with NextToken
            [
                {"Findings": SAMPLE_FINDINGS[:2], "NextToken": "page-2-token"},
                {"Findings": SAMPLE_FINDINGS[2:4], "NextToken": "page-3-token"},
                {"Findings": SAMPLE_FINDINGS[4:5]},
            ],
            TEST_START_TIME,  # start_time
            None,  # end_time (use current time)
            None,  # id_ignore_list
            2,  # page_size (forces pagination)
            5,  # limit
            None,  # start_token
            5,  # expected_events_count
            None,  # expected_next_token (no more pages)
            3,  # expected_api_calls (3 pages: 2+2+1 events)
            [  # expected_debug_messages
                "Starting get_events pagination",
                "Continuing pagination",
                "Pagination completed after",
            ],
        ),
        (
            "with_ignore_list_filters_duplicates",
            # Mock API returns findings, some filtered by ignore list
            [{"Findings": [SAMPLE_FINDINGS[0], IGNORED_FINDINGS[0], SAMPLE_FINDINGS[2], IGNORED_FINDINGS[1]]}],
            TEST_START_TIME,  # start_time
            TEST_END_TIME,  # end_time
            STANDARD_IGNORE_LIST,  # id_ignore_list
            50,  # page_size
            10,  # limit
            None,  # start_token
            2,  # expected_events_count (2 not ignored)
            None,  # expected_next_token
            1,  # expected_api_calls (single page, filtering happens after)
            ["Starting get_events pagination", "2 duplicates filtered", "Pagination completed after"],  # expected_debug_messages
        ),
        (
            "limit_stops_pagination_early",
            # Mock API has more pages but limit stops collection
            [
                {"Findings": SAMPLE_FINDINGS[:2], "NextToken": "has-more-pages"},
                {"Findings": SAMPLE_FINDINGS[2:3], "NextToken": "still-more-pages"},
            ],
            None,  # start_time (no time filter)
            None,  # end_time
            None,  # id_ignore_list
            2,  # page_size
            3,  # limit (stops after 3 events)
            None,  # start_token
            3,  # expected_events_count
            "still-more-pages",  # expected_next_token (stopped due to limit)
            2,  # expected_api_calls (2 pages: 2+1 events, stops at limit)
            [  # expected_debug_messages
                "Starting get_events pagination",
                "Reached limit",
                "Pagination completed after",
            ],
        ),
        (
            "continuation_from_start_token",
            # Mock API continues from provided start_token
            [{"Findings": SAMPLE_FINDINGS[3:5]}],
            TEST_START_TIME,  # start_time
            TEST_END_TIME,  # end_time
            None,  # id_ignore_list
            50,  # page_size
            0,  # limit
            "continue-from-here",  # start_token
            2,  # expected_events_count
            None,  # expected_next_token
            1,  # expected_api_calls (single page from token)
            [
                "Starting get_events pagination",
                "No more pages available from AWS",
                "Pagination completed after",
            ],  # expected_debug_messages
        ),
        (
            "empty_response_no_events",
            # Mock API returns no findings
            [{"Findings": []}],
            TEST_START_TIME,  # start_time
            TEST_END_TIME,  # end_time
            None,  # id_ignore_list
            50,  # page_size
            10,  # limit
            None,  # start_token
            0,  # expected_events_count
            None,  # expected_next_token
            1,  # expected_api_calls (single empty response)
            [  # expected_debug_messages
                "Starting get_events pagination",
                "AWS returned 0 findings",
                "Pagination completed after",
            ],
        ),
        (
            "all_events_filtered_by_ignore_list",
            # Mock API returns findings but all are in ignore list
            [{"Findings": IGNORED_FINDINGS[2:4]}],
            None,  # start_time (no time filter)
            None,  # end_time
            DUPLICATE_IGNORE_LIST,  # id_ignore_list
            50,  # page_size
            0,  # limit
            None,  # start_token
            0,  # expected_events_count (all filtered)
            None,  # expected_next_token
            1,  # expected_api_calls (single page, all filtered after)
            ["Starting get_events pagination", "2 duplicates filtered", "Pagination completed after"],  # expected_debug_messages
        ),
    ],
)
@freeze_time("2023-01-01 15:00:00")
def test_get_events__end_to_end(
    mocker,
    scenario,
    mock_api_responses,
    start_time,
    end_time,
    id_ignore_list,
    page_size,
    limit,
    start_token,
    expected_events_count,
    expected_next_token,
    expected_api_calls,
    expected_debug_messages,
):
    """
    Test end-to-end get_events functionality covering all scenarios:

    Case 1 (single_page_no_pagination): Single API call returns all findings, no NextToken pagination
    Case 2 (multiple_pages_with_pagination): Multiple API calls with NextToken pagination until completion
    Case 3 (with_ignore_list_filters_duplicates): Ignore list properly filters out specified finding IDs
    Case 4 (limit_stops_pagination_early): Limit parameter stops pagination before all pages processed
    Case 5 (continuation_from_start_token): Pagination continues from provided start_token
    Case 6 (empty_response_no_events): API returns empty findings list gracefully
    Case 7 (all_events_filtered_by_ignore_list): All findings filtered out by ignore list

    Given:
        - A mocked AWS SecurityHub client that returns predefined API response sequences
        - Various time ranges, ignore lists, page sizes, limits, and pagination tokens
        - Mocked demisto.debug for logging verification
        - Frozen time at 2023-01-01 15:00:00 for consistent datetime.now() calls
        - Expected API call counts and specific debug messages for each scenario
    When:
        - Calling get_events with different parameter combinations and scenarios
    Then:
        - Ensure correct number of events are returned after deduplication and filtering
        - Ensure pagination works correctly with NextToken handling and limit enforcement
        - Ensure ignore list properly filters duplicate finding IDs
        - Ensure time filtering parameters are correctly passed to AWS API
        - Ensure continuation tokens work for resuming pagination from specific points
        - Ensure correct number of API calls are made (pagination iteration validation)
        - Ensure scenario-specific debug messages are logged correctly
        - Ensure events have proper structure (Id and CreatedAt fields)
    """
    # Mock demisto.debug to capture logging
    mock_debug = mocker.patch("AWSSecurityHubEventCollector.demisto.debug")

    # Create mock client with sequential API responses
    mock_client = mocker.Mock()
    mock_client.get_findings.side_effect = mock_api_responses

    # Call get_events
    events, next_token = get_events(
        client=mock_client,
        start_time=start_time,
        end_time=end_time,
        id_ignore_list=id_ignore_list,
        page_size=page_size,
        limit=limit,
        start_token=start_token,
    )

    # Verify events count
    assert len(events) == expected_events_count

    # Verify next_token
    assert next_token == expected_next_token

    # Verify expected number of API calls (pagination iterations)
    assert mock_client.get_findings.call_count == expected_api_calls

    # Verify first API call parameters
    first_call = mock_client.get_findings.call_args_list[0]
    call_kwargs = first_call.kwargs

    # Check sorting
    assert call_kwargs["SortCriteria"] == [{"Field": "CreatedAt", "SortOrder": "asc"}]

    # Check time filters
    expected_filters = {}
    expected_end_time = end_time or dt.datetime(2023, 1, 1, 15, 0, 0)  # frozen time
    expected_start_time_str = start_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ") if start_time else None
    expected_end_time_str = expected_end_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ")

    if start_time:
        expected_filters["CreatedAt"] = [{"Start": expected_start_time_str, "End": expected_end_time_str}]
        assert call_kwargs["Filters"] == expected_filters
    else:
        # No Filters key when no start_time
        assert "Filters" not in call_kwargs

    # Check pagination token
    if start_token:
        assert call_kwargs["NextToken"] == start_token
    else:
        assert "NextToken" not in call_kwargs

    # Check page size (MaxResults)
    expected_max_results = page_size if limit == 0 or page_size <= limit else min(page_size, limit)
    assert call_kwargs["MaxResults"] == expected_max_results

    # Verify ignore list filtering worked correctly
    if id_ignore_list:
        returned_ids = {event["Id"] for event in events}
        ignored_ids = set(id_ignore_list)
        assert returned_ids.isdisjoint(ignored_ids), "Ignored IDs should not appear in results"

    # Verify events have proper structure
    for event in events:
        assert "Id" in event
        assert "CreatedAt" in event

    # Verify expected debug messages are present
    debug_calls = [call.args[0] for call in mock_debug.call_args_list]

    for expected_msg in expected_debug_messages:
        assert any(
            expected_msg in msg for msg in debug_calls
        ), f"Expected debug message '{expected_msg}' not found in: {debug_calls}"


# Test findings data and constants imported at top of file


@pytest.fixture
def fetch_events_mocks(mocker):
    """Common mocks for fetch_events tests."""
    mocks = {
        "debug": mocker.patch("AWSSecurityHubEventCollector.demisto.debug"),
        "info": mocker.patch("AWSSecurityHubEventCollector.demisto.info"),
        "error": mocker.patch("AWSSecurityHubEventCollector.demisto.error"),
        "generate_last_run": mocker.patch("AWSSecurityHubEventCollector.generate_last_run"),
        "parse_timestamp": mocker.patch("AWSSecurityHubEventCollector.parse_aws_timestamp"),
        "datetime_now": mocker.patch("AWSSecurityHubEventCollector.dt.datetime"),
        "get_events": mocker.patch("AWSSecurityHubEventCollector.get_events"),
        "client": mocker.Mock(),
    }

    # Configure common mock behaviors
    mocks["datetime_now"].now.return_value = FIXED_END_TIME
    mocks["parse_timestamp"].side_effect = lambda ts: dt.datetime.strptime(ts, "%Y-%m-%dT%H:%M:%S.%fZ")

    return mocks


def _verify_debug_messages(debug_calls, info_calls, expected_debug=None, expected_info=None):
    """Helper to verify debug/info messages."""
    # Always check for initial fetch message
    assert any("Fetching events with last_run:" in msg for msg in debug_calls)

    if expected_debug:
        for msg in expected_debug:
            assert any(msg in call for call in debug_calls), f"Expected debug message '{msg}' not found"

    if expected_info:
        for msg in expected_info:
            assert any(msg in call for call in info_calls), f"Expected info message '{msg}' not found"


def test_fetch_events_first_fetch_success(fetch_events_mocks):
    """
    Test fetch_events for first fetch scenario: empty last_run with first_fetch_time.

    Given: Empty last_run and valid first_fetch_time
    When: Calling fetch_events
    Then: Should fetch new events and create proper next_run structure
    """
    mocks = fetch_events_mocks

    # Configure test-specific mocks
    mocks["get_events"].return_value = (SAMPLE_FINDINGS[:2], None)
    next_run = {
        "last_update_date": "2023-01-01T12:01:00.000Z",
        "last_update_date_finding_ids": ["finding-2"],
    }
    mocks["generate_last_run"].return_value = next_run

    # Call fetch_events
    events, actual_next_run, error = fetch_events(
        client=mocks["client"], last_run={}, first_fetch_time=FIRST_FETCH_TIME, page_size=50, limit=0
    )

    # Verify results
    assert len(events) == 2
    assert error is None
    assert "last_update_date" in str(actual_next_run)
    assert "finding-2" in str(actual_next_run)

    # Verify get_events was called correctly
    assert mocks["get_events"].call_count == 1
    call_kwargs = mocks["get_events"].call_args_list[0].kwargs

    assert call_kwargs["client"] == mocks["client"]
    assert call_kwargs["page_size"] == 50
    assert call_kwargs["end_time"] == FIXED_END_TIME  # Fixed end_time consistency
    assert call_kwargs["id_ignore_list"] == []  # Empty ignore list for first fetch

    # Verify generate_last_run was called
    mocks["generate_last_run"].assert_called_once_with(events, {})

    # Verify debug messages
    debug_calls = [call.args[0] for call in mocks["debug"].call_args_list]
    info_calls = [call.args[0] for call in mocks["info"].call_args_list]

    _verify_debug_messages(
        debug_calls, info_calls, expected_debug=["Successfully fetched 2 new events"], expected_info=["Fetched 2 findings."]
    )


def test_fetch_events_infinite_loop_prevention(fetch_events_mocks):
    """
    Test fetch_events infinite loop prevention: all duplicates but AWS has more events.

    Given: All events are duplicates but AWS has continuation token
    When: Calling fetch_events
    Then: Should make retry with increased limit, then exit when no more events
    """
    mocks = fetch_events_mocks

    # Configure test-specific mocks for infinite loop scenario
    mocks["get_events"].side_effect = [
        ([], "continuation-token"),  # First call: all duplicates, has next token -> triggers retry
        ([], None),  # Second call: still no events, no more pages -> exits loop
    ]
    mocks["generate_last_run"].return_value = {
        "last_update_date": "2023-01-01T12:00:00.000Z",
        "last_update_date_finding_ids": ["duplicate-1"],
    }

    # Call fetch_events
    events, next_run, error = fetch_events(
        client=mocks["client"],
        last_run={"last_update_date": "2023-01-01T12:00:00.000Z", "last_update_date_finding_ids": ["duplicate-1"]},
        first_fetch_time=None,
        page_size=50,
        limit=10,
    )

    # Verify results
    assert len(events) == 0  # No new events (all duplicates)
    assert error is None

    # Verify loop behavior: should make 2 calls
    assert mocks["get_events"].call_count == 2

    # Verify second call has increased limit (RETRY_LIMIT_INCREMENT = 100)
    second_call = mocks["get_events"].call_args_list[1]
    expected_retry_limit = 100  # RETRY_LIMIT_INCREMENT value from code
    assert second_call.kwargs["limit"] == expected_retry_limit

    # Verify both calls use same fixed end_time
    for call_args in mocks["get_events"].call_args_list:
        assert call_args.kwargs["end_time"] == FIXED_END_TIME

    # Verify debug messages
    debug_calls = [call.args[0] for call in mocks["debug"].call_args_list]
    info_calls = [call.args[0] for call in mocks["info"].call_args_list]

    _verify_debug_messages(
        debug_calls,
        info_calls,
        expected_debug=["All events were duplicates and AWS has no more events"],
        expected_info=["No new findings were found.", "Infinite loop prevention"],
    )


def test_fetch_events_api_error_handling(fetch_events_mocks):
    """
    Test fetch_events error handling: API raises exception.

    Given: get_events raises an exception
    When: Calling fetch_events
    Then: Should capture error and return it without raising
    """
    mocks = fetch_events_mocks

    # Configure test-specific mocks for error scenario
    test_exception = Exception("AWS API Error: Rate limit exceeded")
    mocks["get_events"].side_effect = test_exception
    mocks["generate_last_run"].return_value = {"last_update_date": "2023-01-01T12:00:00.000Z", "last_update_date_finding_ids": []}

    # Call fetch_events
    events, next_run, error = fetch_events(
        client=mocks["client"],
        last_run={"last_update_date": "2023-01-01T12:00:00.000Z", "last_update_date_finding_ids": []},
        first_fetch_time=None,
        page_size=50,
        limit=10,
    )

    # Verify error handling
    assert len(events) == 0  # No events due to error
    assert error is not None
    assert isinstance(error, Exception)
    assert str(error) == "AWS API Error: Rate limit exceeded"

    # Verify error logging
    error_calls = [call.args[0] for call in mocks["error"].call_args_list]
    assert any("Error while fetching events" in msg for msg in error_calls)