Abnormal Security Event Collector

from AbnormalSecurityEventCollector import get_events
from CommonServerPython import *
from freezegun import freeze_time


def util_load_json(path):
    with open(path, encoding="utf-8") as f:
        return json.loads(f.read())


def generate_mock_response(num_chunks: int):
    mock_responses = []
    for i in range(1, num_chunks + 1):
        chunk_threats = [{"threatId": f"123456789-{i}{j}"} for j in range(1000)]
        mock_responses.append({"threats": chunk_threats, "nextPageNumber": i + 1})
    return mock_responses


class Client(BaseClient):
    def list_threats(self, params):
        return util_load_json("test_data/test_get_list_threats.json")

    def get_threat(self, threat):
        return util_load_json("test_data/test_get_threat.json").get(threat)


"""
    Command Unit Tests
"""


@freeze_time("2022-09-14")
def test_get_events():
    """
    When:
        - running the get_events function
    Then
        - Assert the returned messages contains only messages in the specific time range (.
        - Assert the returned messages are ordered by datetime.
        - Assert the returned "toAddresses" field in the messages returned as a list.

    """
    client = Client(base_url="url")
    messages, last_run = get_events(client, after="2022-05-02T18:44:38Z", before="", next_page_number=1)

    assert messages == [
        {"abxMessageId": 3, "receivedTime": "2022-06-01T18:44:38Z", "threatId": "123456789-1", "toAddresses": []},
        {"abxMessageId": 3, "receivedTime": "2022-06-02T18:44:38Z", "threatId": "123456789-2"},
        {"abxMessageId": 3, "receivedTime": "2022-06-03T18:44:38Z", "threatId": "123456789-3"},
        {"abxMessageId": 2, "receivedTime": "2022-08-01T18:44:38Z", "threatId": "123456789-1", "toAddresses": ["test1", "test2"]},
        {"abxMessageId": 2, "receivedTime": "2022-08-02T18:44:38Z", "threatId": "123456789-2", "toAddresses": ["test1", "test2"]},
        {"abxMessageId": 2, "receivedTime": "2022-08-03T18:44:38Z", "threatId": "123456789-3"},
    ]


@freeze_time("2022-09-14")
def test_fetch_events_without_nextTrigger(mocker):
    """
    Given:
        - The last fetch did not exceed the fetch limit.
    When:
        - Running the fetch-events command.
    Then:
        - The events list should contain the fetched threats, and the next page number should be 1.
    """
    client = Client(base_url="url")
    mock_response = generate_mock_response(num_chunks=1)
    mock_response[0]["nextPageNumber"] = None
    mocker.patch.object(Client, "list_threats", return_value=mock_response[0])
    mocker.patch("AbnormalSecurityEventCollector.format_messages", return_value=mock_response[0]["threats"][0])
    mocker.patch("AbnormalSecurityEventCollector.get_messages_by_datetime", return_value=mock_response[0]["threats"][0])
    mocker.patch("AbnormalSecurityEventCollector.sorted", return_value=mock_response[0]["threats"])
    threats, last_run = get_events(client, after="2022-05-02T18:44:38Z", before="", next_page_number=1)
    assert threats == mock_response[0]["threats"]
    assert last_run.get("before") == "2022-09-14T00:00:00Z"
    assert last_run.get("next_page_number") == 1


def test_fetch_events_with_nextTrigger(mocker):
    """
    Given:
        - The last fetch did exceed the fetch limit.
    When:
        - Running the fetch-events command.
    Then:
        - The events list should contain the fetched threats, and the next page number should be 10.
    """
    client = Client(base_url="url")
    mock_responses = generate_mock_response(num_chunks=10)

    def mock_list_threats(params, mock_responses=mock_responses):
        return mock_responses.pop(0)

    mocker.patch.object(Client, "list_threats", side_effect=mock_list_threats)
    mocker.patch("AbnormalSecurityEventCollector.format_messages", return_value=mock_responses[0]["threats"][0])
    mocker.patch("AbnormalSecurityEventCollector.get_messages_by_datetime", return_value=mock_responses[0]["threats"][0])
    mocker.patch(
        "AbnormalSecurityEventCollector.sorted",
        return_value=[threat for response in mock_responses[:9] for threat in response["threats"]],
    )
    threats, last_run = get_events(client, after="2022-05-02T18:44:38Z", before="", next_page_number=1)
    assert len(threats) == 9000
    assert last_run.get("next_page_number") == 10