FlashpointVulnerabilityDetails

"""FlashpointVulnerabilityDetails Test File."""

import json
from unittest.mock import patch

import demistomock as demisto
import pytest
from FlashpointVulnerabilityDetails import (
    ERROR_MESSAGES,
    MESSAGES,
    SET_INCIDENT_COMMAND,
    VULNERABILITY_DETAILS_FIELD,
    VULNERABILITY_LIST_COMMAND,
    get_vulnerability_details,
    main,
)

""" CONSTANTS """

VULN_IDS = ["111111", "222222", "333333"]

EXPECTED_CVE_IDS = {
    "111111": "CVE-1111-11111, CVE-1111-11112",
    "222222": "CVE-2222-22221",
    "333333": "",
}

""" UTILITY FUNCTIONS """


def util_load_json(path: str):
    """Load a json to python dict."""
    with open(path, encoding="utf-8") as f:
        return json.loads(f.read())


def get_execute_command_mock(mocker, vulnerability_response):
    """
    Mock demisto.executeCommand to return the given response for the vulnerability list command.

    :type mocker: pytest_mock.MockerFixture
    :param mocker: Mocker fixture.

    :type vulnerability_response: Any
    :param vulnerability_response: Response to return for the vulnerability list command.

    :return: The mocked executeCommand object.
    """

    def side_effect(command, args):
        if command == VULNERABILITY_LIST_COMMAND:
            return vulnerability_response
        return [{"Type": 1, "Contents": "", "ContentsFormat": "text"}]

    return mocker.patch.object(demisto, "executeCommand", side_effect=side_effect)


def get_set_incident_rows(execute_command_mock):
    """
    Return the grid rows passed to the setIncident command.

    :type execute_command_mock: unittest.mock.MagicMock
    :param execute_command_mock: The mocked executeCommand object.

    :return: The grid rows of the vulnerability details field, or None when setIncident was not called.
    """
    set_incident_calls = [call for call in execute_command_mock.call_args_list if call.args[0] == SET_INCIDENT_COMMAND]

    if not set_incident_calls:
        return None

    assert len(set_incident_calls) == 1
    return set_incident_calls[0].args[1][VULNERABILITY_DETAILS_FIELD]


""" TEST CASES """


def test_get_vulnerability_details_success(mocker):
    """
    Test case scenario for successful execution of get_vulnerability_details.

    Given:
       - an incident with the vulnerability details grid populated
    When:
       - Calling `get_vulnerability_details` function
    Then:
       - Requests all the vulnerability IDs in a single call and sets the CVE IDs on the grid without losing the
         existing column values.
    """
    vulnerability_response = util_load_json("test_data/vulnerability_list_success.json")
    incident = util_load_json("test_data/incident_with_vulnerabilities.json")

    # The rows are updated in place, so load the expected rows separately to compare them against the updated rows.
    expected_rows = util_load_json("test_data/incident_with_vulnerabilities.json")["CustomFields"][VULNERABILITY_DETAILS_FIELD]

    mocker.patch.object(demisto, "incident", return_value=incident)
    execute_command_mock = get_execute_command_mock(mocker, [vulnerability_response])

    result = get_vulnerability_details()

    assert execute_command_mock.call_args_list[0].args == (
        VULNERABILITY_LIST_COMMAND,
        {"vulnerability_ids": ",".join(VULN_IDS), "size": len(VULN_IDS)},
    )

    rows = get_set_incident_rows(execute_command_mock)
    assert len(rows) == len(VULN_IDS)
    for row, expected_row in zip(rows, expected_rows):
        assert row["cve_ids"] == EXPECTED_CVE_IDS[row["vuln_id"]]
        for column, expected_value in expected_row.items():
            assert row[column] == expected_value

    assert result == vulnerability_response


@pytest.mark.parametrize(
    "incident",
    [
        {},
        {"CustomFields": {}},
        {"CustomFields": {VULNERABILITY_DETAILS_FIELD: []}},
        {"CustomFields": {VULNERABILITY_DETAILS_FIELD: [{"title": "No ID available"}]}},
    ],
)
def test_get_vulnerability_details_when_no_vulnerability_ids(mocker, incident):
    """
    Test case scenario for execution of get_vulnerability_details when the incident has no vulnerability IDs.

    Given:
       - an incident without the vulnerability details grid or without any vulnerability ID in it
    When:
       - Calling `get_vulnerability_details` function
    Then:
       - Returns a no records found message without executing any command.
    """
    mocker.patch.object(demisto, "incident", return_value=incident)
    execute_command_mock = mocker.patch.object(demisto, "executeCommand")

    result = get_vulnerability_details()

    assert result.readable_output == MESSAGES["NO_RECORDS_FOUND"]
    assert execute_command_mock.call_count == 0


def test_get_vulnerability_details_when_command_fails(mocker):
    """
    Test case scenario for execution of get_vulnerability_details when the vulnerability list command fails.

    Given:
       - an incident with the vulnerability details grid and an error response from the vulnerability list command
    When:
       - Calling `get_vulnerability_details` function
    Then:
       - Raises a valid error message without updating the grid.
    """
    error = "Error in API call [404] - Not Found"
    incident = util_load_json("test_data/incident_with_vulnerabilities.json")

    mocker.patch.object(demisto, "incident", return_value=incident)
    execute_command_mock = get_execute_command_mock(mocker, {"Type": 4, "Contents": error, "ContentsFormat": "text"})

    with pytest.raises(ValueError) as raised_error:
        get_vulnerability_details()

    assert str(raised_error.value) == ERROR_MESSAGES["FAILED_COMMAND"].format(VULNERABILITY_LIST_COMMAND, error)
    assert get_set_incident_rows(execute_command_mock) is None


@patch("FlashpointVulnerabilityDetails.return_results")
def test_main_success(mock_return_results, mocker):
    """
    Test case scenario for successful execution of the script through the main function.

    Given:
       - an incident with the vulnerability details grid populated
    When:
       - Calling `main` function
    Then:
       - Returns the CVE IDs of the vulnerabilities.
    """
    vulnerability_response = util_load_json("test_data/vulnerability_list_success.json")
    incident = util_load_json("test_data/incident_with_vulnerabilities.json")

    mocker.patch.object(demisto, "incident", return_value=incident)
    execute_command_mock = get_execute_command_mock(mocker, [vulnerability_response])

    main()

    assert execute_command_mock.call_args_list[0].args == (
        VULNERABILITY_LIST_COMMAND,
        {"vulnerability_ids": ",".join(VULN_IDS), "size": len(VULN_IDS)},
    )
    assert mock_return_results.call_args.args[0] == vulnerability_response


@patch("FlashpointVulnerabilityDetails.return_error")
def test_main_calls_return_error_on_exception(mock_return_error, mocker):
    """
    Test case scenario for execution of the script through the main function when an exception is raised.

    Given:
       - an incident with the vulnerability details grid and an error response from the vulnerability list command
    When:
       - Calling `main` function
    Then:
       - Returns a valid error message.
    """
    error = "Error in API call [404] - Not Found"
    incident = util_load_json("test_data/incident_with_vulnerabilities.json")

    mocker.patch.object(demisto, "incident", return_value=incident)
    mocker.patch.object(demisto, "error")
    get_execute_command_mock(mocker, {"Type": 4, "Contents": error, "ContentsFormat": "text"})

    main()

    assert mock_return_error.call_args.args[0] == (
        "Failed to execute FlashpointVulnerabilityDetails. Error: "
        f"{ERROR_MESSAGES['FAILED_COMMAND'].format(VULNERABILITY_LIST_COMMAND, error)}"
    )