DarktraceAIA

import json


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


"""*****COMMAND FUNCTIONS****"""


def test_fetch_incidents(requests_mock):
    """Tests the fetch-incidents command function.

    Configures requests_mock instance to generate the appropriate
    get_alert API response, loaded from a local JSON file. Checks
    the output of the command function with the expected output.
    """
    from DarktraceAIA import Client, fetch_incidents

    # GIVEN an integration is configured and fetch incidents
    mock_response = util_load_json("test_data/incident_fetch.json")
    requests_mock.get(
        "https://usw1-51965-01.cloud.darktrace.com/aianalyst/" + "incidentevents?mingroupscore=0&starttime=1598932817000",
        json=mock_response,
    )

    client = Client(base_url="https://usw1-51965-01.cloud.darktrace.com/", verify=False, auth=("examplepub", "examplepri"))

    # WHEN the most recent call was made on Mon, Aug 31, 2020 9 PM Pacific
    last_run = {
        "last_fetch": 1598932817000  # Mon, Aug 31, 2020 9 PM Pacific
    }

    _, integration_response = fetch_incidents(
        client=client, max_alerts=20, last_run=last_run, first_fetch_time="1 day ago", min_score=0
    )

    # THEN the relevant information will be fetched and pulled
    expected_response = util_load_json("test_data/formatted_incident_fetch.json")

    assert integration_response == expected_response
    assert len(integration_response) == 2


def test_get_ai_analyst_incident_event(requests_mock):
    """Tests get_ai_analyst_incident_event command function.

    Configures requests_mock instance to generate the appropriate
    get_alert API response, loaded from a local JSON file. Checks
    the output of the command function with the expected output.
    """
    from DarktraceAIA import Client, get_ai_analyst_incident_event_command

    # GIVEN an integration is configured and fetch incidents
    eventId = "bc64f242-ce29-4f35-bc94-230991116564"
    mock_api_response = util_load_json("test_data/ai_analyst_incident.json")
    requests_mock.get("https://mock.darktrace.com/aianalyst/incidentevents?uuid=" + eventId, json=mock_api_response)

    client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri"))

    args = {
        "eventId": eventId,
    }
    integration_response = get_ai_analyst_incident_event_command(client, args)
    expected_response = util_load_json("test_data/formatted_ai_analyst_incident.json")

    assert integration_response.outputs == expected_response
    assert integration_response.outputs_prefix == "Darktrace.AIAnalyst"


def test_get_comments_for_ai_analyst_incident_event_command(requests_mock):
    """Tests get_comments_for_ai_analyst_incident_event_command command function.

    Configures requests_mock instance to generate the appropriate
    get_alert API response, loaded from a local JSON file. Checks
    the output of the command function with the expected output.
    """
    from DarktraceAIA import Client, get_comments_for_ai_analyst_incident_event_command

    # GIVEN an integration is configured and fetch incidents
    eventId = "bc64f242-ce29-4f35-bc94-230991116564"
    mock_api_response = util_load_json("test_data/get_comment_response.json")
    requests_mock.get("https://mock.darktrace.com/aianalyst/incident/comments?incident_id=" + eventId, json=mock_api_response)

    client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri"))

    args = {
        "eventId": eventId,
    }
    integration_response = get_comments_for_ai_analyst_incident_event_command(client, args)
    expected_response = util_load_json("test_data/formatted_get_comment_response.json")

    assert integration_response.outputs == expected_response
    assert integration_response.outputs_prefix == "Darktrace.AIAnalyst"


def test_post_comment_to_ai_analyst_incident_event(requests_mock):
    """Tests post_comments_for_ai_analyst_incident_event_command command function.

    Configures requests_mock instance to generate the appropriate
    get_alert API response, loaded from a local JSON file. Checks
    the output of the command function with the expected output.
    """
    from DarktraceAIA import Client, post_comment_to_ai_analyst_incident_event_command

    # GIVEN an integration is configured and fetch incidents
    eventId = "bc64f242-ce29-4f35-bc94-230991116564"
    mock_api_response = util_load_json("test_data/post_comment.json")
    requests_mock.post("https://mock.darktrace.com/aianalyst/incident/comments", json=mock_api_response)

    client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri"))

    args = {"eventId": eventId, "comment": "test comment"}
    integration_response = post_comment_to_ai_analyst_incident_event_command(client, args)
    expected_response = util_load_json("test_data/formatted_post_comment.json")

    assert integration_response.outputs == expected_response
    assert integration_response.outputs_prefix == "Darktrace.AIAnalyst"


def test_acknowledge_ai_analyst_incident_event(requests_mock):
    """Tests acknowledge_ai_analyst_incident_event_command command function.

    Configures requests_mock instance to generate the appropriate
    get_alert API response, loaded from a local JSON file. Checks
    the output of the command function with the expected output.
    """
    from DarktraceAIA import Client, acknowledge_ai_analyst_incident_event_command

    # GIVEN an integration is configured and fetch incidents
    eventId = "bc64f242-ce29-4f35-bc94-230991116564"
    mock_api_response = util_load_json("test_data/ack_response.json")
    requests_mock.post("https://mock.darktrace.com/aianalyst/acknowledge", json=mock_api_response)

    client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri"))

    args = {
        "eventId": eventId,
    }
    integration_response = acknowledge_ai_analyst_incident_event_command(client, args)
    expected_response = util_load_json("test_data/formatted_ack.json")

    assert integration_response.outputs == expected_response
    assert integration_response.outputs_prefix == "Darktrace.AIAnalyst"


def test_unacknowledge_ai_analyst_incident_event_command(requests_mock):
    """Tests acknowledge_ai_analyst_incident_event_command command function.

    Configures requests_mock instance to generate the appropriate
    get_alert API response, loaded from a local JSON file. Checks
    the output of the command function with the expected output.
    """
    from DarktraceAIA import Client, unacknowledge_ai_analyst_incident_event_command

    # GIVEN an integration is configured and fetch incidents
    eventId = "bc64f242-ce29-4f35-bc94-230991116564"
    mock_api_response = util_load_json("test_data/unack_response.json")
    requests_mock.post("https://mock.darktrace.com/aianalyst/unacknowledge", json=mock_api_response)

    client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri"))

    args = {
        "eventId": eventId,
    }
    integration_response = unacknowledge_ai_analyst_incident_event_command(client, args)
    expected_response = util_load_json("test_data/formatted_unack.json")

    assert integration_response.outputs == expected_response
    assert integration_response.outputs_prefix == "Darktrace.AIAnalyst"


def test_get_ai_analyst_incident_group_from_eventId(requests_mock):
    from DarktraceAIA import Client, get__ai_analyst_incident_group_from_eventId_command

    eventId = "bc64f242-ce29-4f35-bc94-230991116564"
    mock_api_response = util_load_json("test_data/group_response.json")
    requests_mock.get("https://mock.darktrace.com/aianalyst/groups?uuid=" + eventId, json=mock_api_response)

    client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri"))

    args = {
        "eventId": eventId,
    }

    integration_response = get__ai_analyst_incident_group_from_eventId_command(client, args)
    expected_response = util_load_json("test_data/formatted_group_response.json")

    assert integration_response.outputs == expected_response
    assert integration_response.outputs_prefix == "Darktrace.AIAnalyst"
    assert integration_response.outputs_key_field == "groupId"