import json import pytest from CommonServerPython import arg_to_datetime def test_dedup_elements(): from CohesityHeliosEventCollector import ALERT_TIME_FIELD, AUDIT_LOGS_TIME_FIELD, adjust_and_dedup_elements """ Case 1: Given a list of 3 elements where all IDs appear in the ID list. We expect the result list to have no elements at all and that the final list length was not changed. """ new_elements = [ {"id": "1", "latestTimestampUsecs": 1704096000000000}, {"id": "2", "latestTimestampUsecs": 1704182400000000}, {"id": "3", "latestTimestampUsecs": 1704268800000000}, ] existing_element_ids = ["1", "2", "3"] deduped_elements = adjust_and_dedup_elements( new_elements=new_elements, existing_element_ids=existing_element_ids, time_field_name="" ) assert deduped_elements == [] assert len(new_elements) == 3 """ Case 2: Given a list of 2 elements where all elements appear in the existing ID list We expect the result list to have no elements at all and that the final list length was not changed. """ new_elements = [{"id": "2", "latestTimestampUsecs": 1704182400000000}, {"id": "3", "latestTimestampUsecs": 1704268800000000}] existing_element_ids = ["1", "2", "3"] deduped_elements = adjust_and_dedup_elements( new_elements=new_elements, existing_element_ids=existing_element_ids, time_field_name="" ) assert deduped_elements == [] assert len(new_elements) == 2 """ Case 3: Given a list of 3 elements where the first element appear in the existing ID list. We expect the result list to have the other two elements and that the final list length was not changed. """ new_elements = [ {"id": "1", "latestTimestampUsecs": 1704096000000000}, {"id": "2", "latestTimestampUsecs": 1704182400000000}, {"id": "3", "latestTimestampUsecs": 1704268800000000}, ] existing_element_ids = ["1"] deduped_elements = adjust_and_dedup_elements( new_elements=new_elements, existing_element_ids=existing_element_ids, time_field_name=ALERT_TIME_FIELD ) assert deduped_elements == [ {"id": "2", "latestTimestampUsecs": 1704182400000000, "_time": "2024-01-02T08:00:00.000Z"}, {"id": "3", "latestTimestampUsecs": 1704268800000000, "_time": "2024-01-03T08:00:00.000Z"}, ] assert len(new_elements) == 3 """ Case 4: Given a list of 3 elements while the existing ID list is empty. We expect the result list to have all elements and that the final list length was not changed. """ new_elements = [ {"id": "1", "timestampUsecs": 1704096000000000}, {"id": "2", "timestampUsecs": 1704182400000000}, {"id": "3", "timestampUsecs": 1704268800000000}, ] existing_element_ids = [] deduped_elements = adjust_and_dedup_elements( new_elements=new_elements, existing_element_ids=existing_element_ids, time_field_name=AUDIT_LOGS_TIME_FIELD ) assert deduped_elements == [ {"id": "1", "timestampUsecs": 1704096000000000, "_time": "2024-01-01T08:00:00.000Z"}, {"id": "2", "timestampUsecs": 1704182400000000, "_time": "2024-01-02T08:00:00.000Z"}, {"id": "3", "timestampUsecs": 1704268800000000, "_time": "2024-01-03T08:00:00.000Z"}, ] assert len(new_elements) == 3 """ Case 5: Given an empty list elements. We expect the result list to have no elements at all and that the final list length was not changed. """ new_elements = [] existing_element_ids = ["1", "2", "3"] deduped_elements = adjust_and_dedup_elements( new_elements=new_elements, existing_element_ids=existing_element_ids, time_field_name="" ) assert deduped_elements == [] assert len(new_elements) == 0 def test_get_earliest_event_ids_with_the_same_time(): from CohesityHeliosEventCollector import ALERT_TIME_FIELD, AUDIT_LOGS_TIME_FIELD, get_earliest_event_ids_with_the_same_time time_field = ALERT_TIME_FIELD """ Case 1: Given a list of Alert events where there is only one event that has the earliest timestamp Ensure only the ID of the earliest Alert is returned """ events = [ {"latestTimestampUsecs": "3", "id": "c"}, {"latestTimestampUsecs": "2", "id": "b"}, {"latestTimestampUsecs": "1", "id": "a"}, ] earliest_event_fetched_ids = get_earliest_event_ids_with_the_same_time(events=events, time_field=time_field) assert earliest_event_fetched_ids == ["a"] """ Case 2: Given a list of Alert events where there are two "earliest" events Ensure the ID of the TWO earliest Alerts is returned """ events = [ {"latestTimestampUsecs": "3", "id": "d"}, {"latestTimestampUsecs": "2", "id": "c"}, {"latestTimestampUsecs": "1", "id": "b"}, {"latestTimestampUsecs": "1", "id": "a"}, ] earliest_event_fetched_ids = get_earliest_event_ids_with_the_same_time(events=events, time_field=time_field) assert earliest_event_fetched_ids == ["a", "b"] time_field = AUDIT_LOGS_TIME_FIELD """ Case 3: Given a list of Audit Log events where there is only one event that has the earliest timestamp Ensure only the ID of the earliest event is returned """ events = [{"timestampUsecs": "3", "id": "c"}, {"timestampUsecs": "2", "id": "b"}, {"timestampUsecs": "1", "id": "a"}] earliest_event_fetched_ids = get_earliest_event_ids_with_the_same_time(events=events, time_field=time_field) assert earliest_event_fetched_ids == ["a"] """ Case 4: Given a list of Audit Log events where there are two "earliest" events Ensure the ID of the TWO earliest Audit logs is returned """ events = [ {"timestampUsecs": "3", "id": "d"}, {"timestampUsecs": "2", "id": "c"}, {"timestampUsecs": "1", "id": "b"}, {"timestampUsecs": "1", "id": "a"}, ] earliest_event_fetched_ids = get_earliest_event_ids_with_the_same_time(events=events, time_field=time_field) assert earliest_event_fetched_ids == ["a", "b"] def test_hash_fields_to_create_id(): """ Given dummy audit log event with the relevant fields Ensure the id is created correctly """ from CohesityHeliosEventCollector import hash_fields_to_create_id event = { "details": "dummy_details", "username": "dummy_username", "domain": "dummy_domain", "sourceType": "dummy_sourceType", "entityName": "dummy_entityName", "entityType": "dummy_entityType", "action": "dummy_action", "timestampUsecs": "dummy_timestampUsecs", "ip": "dummy_ip", "isImpersonation": "dummy_isImpersonation", "tenantId": "dummy_tenantId", "originalTenantId": "dummy_originalTenantId", "serviceContext": "dummy_serviceContext", } _id = hash_fields_to_create_id(event) assert _id == "8bb89cb674035796b755e9e1db5022dc750e904f520eb290d18e134b12656bf2" class TestFetchEventsCommand: """ Class to test the different Fetch events flow. Fetch events has test 3 case: 1: There are fewer events than page_size on first request 2: There are more than page_size events but there are less than max_fetch events 3: There are more than max_fetch events """ base_url = "https://test.com" audit_logs_endpoint = "mcm/audit-logs" alerts_endpoint = "mcm/alerts" mock_time = "2024-01-01 10:00:00" mock_fixed_time_unix = int(arg_to_datetime(mock_time).timestamp() * 1000000) @staticmethod def load_response(event_type) -> dict: from CohesityHeliosEventCollector import EventType filename = ( "test_data/CohesityHeliosEventCollector-AuditLogList.json" if event_type == EventType.audit_log else "test_data/CohesityHeliosEventCollector-AlertList.json" ) with open(filename) as f: return json.loads(f.read()) @pytest.fixture() def audit_logs_mock_res(self): import CohesityHeliosEventCollector return self.load_response(CohesityHeliosEventCollector.EventType.audit_log) @pytest.fixture() def alerts_mock_res(self): import CohesityHeliosEventCollector return self.load_response(CohesityHeliosEventCollector.EventType.alert) def test_fetch_events_command_case_1(self, requests_mock, mocker, audit_logs_mock_res, alerts_mock_res): """ Case 1 is when where are fewer events (4) than page_size (10,000) on the first request. We expect: - Each event type API call to be called once - To have only 4 events returned - Audit logs next start time for the next fetch to be set to the latest pulled event timestamp plus 1 (170691857331523) - No list of ids_for_dedup and no latest_event_fetched_timestamp for audit logs - Alerts next start time for the next fetch to be set to the latest pulled event timestamp plus 1 (1708175775539274) - No list of ids_for_dedup and no latest_event_fetched_timestamp for alerts """ from CohesityHeliosEventCollector import Client, fetch_events_command # mockers mocker.patch("CohesityHeliosEventCollector.arg_to_datetime", return_value=arg_to_datetime(self.mock_time)) audit_logs_call = requests_mock.get(f"{self.base_url}/{self.audit_logs_endpoint}", json=audit_logs_mock_res[0]) alerts_call = requests_mock.get(f"{self.base_url}/{self.alerts_endpoint}", json=alerts_mock_res[0]) client = Client(base_url=self.base_url) events, last_run = fetch_events_command(client=client, last_run={}, max_fetch=1000) assert audit_logs_call.call_count == alerts_call.call_count == 1 assert len(events) == 4 assert last_run["audit_cache"]["next_start_timestamp"] == 170691857331523 assert not last_run["audit_cache"]["ids_for_dedup"] assert not last_run["audit_cache"]["latest_event_fetched_timestamp"] assert last_run["alert_cache"]["next_start_timestamp"] == 1708175775539274 assert not last_run["alert_cache"]["ids_for_dedup"] assert not last_run["audit_cache"]["latest_event_fetched_timestamp"] def test_fetch_events_command_case_2(self, requests_mock, mocker, audit_logs_mock_res, alerts_mock_res): """ Case 2 is when there are more events (3) from each type than the page_size (2), but there are not more than max_fetch (1000). We expect: - Each event type API call to be called twice - That the endtimeusecs in the 2dn API call for audit logs will be the same as the time of the earliest event fetched timestamp - That the enddateusecs in the 2dn API call for alerts will be the same as the time of the earliest event fetched timestamp - To have 6 events returned - Audit logs next start time for the next fetch to be set to the latest pulled event timestamp plus 1 (170691857331523) - No list of ids_for_dedup and no latest_event_fetched_timestamp for audit logs - Alerts next start time for the next fetch to be set to the latest pulled event timestamp plus 1 (1708175775539274) - No list of ids_for_dedup and no latest_event_fetched_timestamp for alerts """ import CohesityHeliosEventCollector from CohesityHeliosEventCollector import Client, fetch_events_command # mockers mocker.patch.object(CohesityHeliosEventCollector, "PAGE_SIZE", 2) mocker.patch("CohesityHeliosEventCollector.arg_to_datetime", return_value=arg_to_datetime(self.mock_time)) audit_logs_call = requests_mock.get( f"{self.base_url}/{self.audit_logs_endpoint}", [{"json": audit_logs_mock_res[0]}, {"json": audit_logs_mock_res[2]}] ) alerts_call = requests_mock.get( f"{self.base_url}/{self.alerts_endpoint}", [{"json": alerts_mock_res[0]}, {"json": alerts_mock_res[2]}] ) audit_logs_expected_end_time = audit_logs_mock_res[0].get("auditLogs")[1].get("timestampUsecs") alerts_expected_end_time = alerts_mock_res[0].get("alertsList")[1].get("latestTimestampUsecs") client = Client(base_url=self.base_url) events, last_run = fetch_events_command(client=client, last_run={}, max_fetch=1000) assert audit_logs_call.call_count == alerts_call.call_count == 2 assert audit_logs_call.request_history[1].qs["endtimeusecs"][0] == str(audit_logs_expected_end_time) assert alerts_call.request_history[1].qs["enddateusecs"][0] == str(alerts_expected_end_time) assert len(events) == 6 assert last_run["audit_cache"]["next_start_timestamp"] == 170691857331523 assert not last_run["audit_cache"]["ids_for_dedup"] assert not last_run["audit_cache"]["latest_event_fetched_timestamp"] assert last_run["alert_cache"]["next_start_timestamp"] == 1708175775539274 assert not last_run["alert_cache"]["ids_for_dedup"] assert not last_run["alert_cache"]["latest_event_fetched_timestamp"] def test_fetch_events_command_case_3(self, requests_mock, mocker, audit_logs_mock_res, alerts_mock_res): """ Case 3 is when there are more events than max_fetch events. We expect: - Each event type API call to be called twice - That the endtimeusecs in the 2dn API call for audit logs will be the same as the time of the earliest event fetched timestamp - That the enddateusecs in the 2dn API call for alerts will be the same as the time of the earliest event fetched timestamp - To have 8 events returned - Audit logs next start time for the next fetch to be set to the same initial start time - ids_for_dedup in the audit_log cache has the ID of the earliest audit log event - latest_event_fetched_timestamp in the audit_log cache holds the latest audit log event timestamp plus 1 sec - Alerts next start time for the next fetch to be set to the same initial start time - No list of ids_for_dedup and no latest_event_fetched_timestamp for alerts - ids_for_dedup in the alerts cache has the ID of the earliest alert event - latest_event_fetched_timestamp in the alerts cache holds the latest alert event timestamp plus 1 sec """ import CohesityHeliosEventCollector from CohesityHeliosEventCollector import Client, fetch_events_command # mockers mocker.patch.object(CohesityHeliosEventCollector, "PAGE_SIZE", 2) mocker.patch("CohesityHeliosEventCollector.arg_to_datetime", return_value=arg_to_datetime(self.mock_time)) audit_logs_call = requests_mock.get( f"{self.base_url}/{self.audit_logs_endpoint}", [{"json": audit_logs_mock_res[0]}, {"json": audit_logs_mock_res[1]}] ) alerts_call = requests_mock.get( f"{self.base_url}/{self.alerts_endpoint}", [{"json": alerts_mock_res[0]}, {"json": alerts_mock_res[1]}] ) audit_logs_expected_end_time = audit_logs_mock_res[0].get("auditLogs")[1].get("timestampUsecs") alerts_expected_end_time = alerts_mock_res[0].get("alertsList")[1].get("latestTimestampUsecs") client = Client(base_url=self.base_url) events, last_run = fetch_events_command(client=client, last_run={}, max_fetch=3) assert audit_logs_call.call_count == alerts_call.call_count == 2 assert audit_logs_call.request_history[1].qs["endtimeusecs"][0] == str(audit_logs_expected_end_time) assert alerts_call.request_history[1].qs["enddateusecs"][0] == str(alerts_expected_end_time) assert len(events) == 8 assert last_run["audit_cache"]["next_start_timestamp"] == self.mock_fixed_time_unix assert last_run["audit_cache"]["ids_for_dedup"] == ["4b2c6b16b2d288ed617ec28f5aca574bd8fb5670c696b726f483057ccc8aff16"] assert last_run["audit_cache"]["latest_event_fetched_timestamp"] == 170691857331523 assert last_run["alert_cache"]["next_start_timestamp"] == self.mock_fixed_time_unix assert last_run["alert_cache"]["ids_for_dedup"] == ["66770"] assert last_run["alert_cache"]["latest_event_fetched_timestamp"] == 1708175775539274