McAfee Threat Intelligence Exchange V2

import json
from unittest.mock import patch

import pytest


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


def test_safe_get_file_reputation_returned_exception(mocker):
    """
    Given:
        - TIE client and hash parameter
    When:
        - The TIE client returns some exception due to an exception thrown from the API call when running get_file_reputation
    Then:
        - Validate that we print to log and return None object
    """
    from McAfeeTIEV2 import safe_get_file_reputation

    hash_param = {"test": "test"}
    with patch("dxltieclient.TieClient") as mock_tie_client:
        tie_client = mock_tie_client.return_value
        mocker.patch.object(tie_client, "get_file_reputation", side_effect=Exception())
        assert not safe_get_file_reputation(tie_client, hash_param)


def test_safe_get_file_reputation_returned_rep(mocker):
    """
    Given:
        - TIE client and hash parameter
    When:
        - The TIE client returns reputation
    Then:
        - Validate that we return the reputation and not None object
    """
    from McAfeeTIEV2 import safe_get_file_reputation

    hash_param = {"test": "test"}
    with patch("dxltieclient.TieClient") as mock_tie_client:
        tie_client = mock_tie_client.return_value
        mocker.patch.object(tie_client, "get_file_reputation", return_value="test_value")
        assert safe_get_file_reputation(tie_client, hash_param)


def test_set_files_reputation_invalid():
    """
    Given:
        - TIE client, a list of hashes, and an invalid trust_level argument
    When:
        - The function set_file_reputation_command is called to set a new Enterprise reputation for the specified files
    Then:
        - Validate that an exception is thrown in response to the invalid trust_level argument
    """
    from McAfeeTIEV2 import set_files_reputation_command

    with patch("dxltieclient.TieClient") as mock_tie_client:
        tie_client = mock_tie_client.return_value
        with pytest.raises(Exception) as e:
            set_files_reputation_command(
                hashes=["hash1", "hash2"], tie_client=tie_client, trust_level="invalid_trust_level", filename="", comment=""
            )
        assert "Illegal argument trust_level" in str(e)


def test_set_files_reputation_valid(mocker):
    """
    Given:
        - TIE client, a list of hashes, and a valid trust_level argument
    When:
        - The function set_file_reputation_command is called to set a new reputation for the specified files
    Then:
        - Validate the contect of the Command Result that is returned from the function
    """
    from McAfeeTIEV2 import set_files_reputation_command

    mocker.patch("McAfeeTIEV2.get_trust_level_key", return_value="trust_level_key")
    mocker.patch("McAfeeTIEV2.get_hash_type_key", return_value="hash_type_key")
    with patch("dxltieclient.TieClient") as mock_tie_client:
        tie_client = mock_tie_client.return_value
        mocker.patch.object(tie_client, "set_file_reputation", return_value="test_value")
        result = set_files_reputation_command(
            hashes=["hash1", "hash2"], tie_client=tie_client, trust_level="valid_trust_level", filename="", comment=""
        )
        assert "Successfully set files reputation" in result.readable_output


QUERY_LIMIT_CASES = [
    ({"query_limit": -30}),
    (
        {"query_limit": 10000}  # A number larger than the maximum
    ),
    ({"query_limit": 0}),
]


@pytest.mark.parametrize("params", QUERY_LIMIT_CASES)
def test_files_references_invalid_query_limit(params):
    """
    Given:
        - TIE client, a list of hashes, and a query limit
    When:
        - The function file_references_command is called to retrieve the references of the given files
    Then:
        - Validate that an exception is thrown in response to the invalid query_limit value
    """
    from McAfeeTIEV2 import files_references_command

    with patch("dxltieclient.TieClient") as mock_tie_client:
        tie_client = mock_tie_client.return_value
        with pytest.raises(Exception) as e:
            files_references_command(hashes=["hash1", "hash2"], tie_client=tie_client, query_limit=params["query_limit"])
        assert "Query limit must not exceed" in str(e) or "Query limit must not be zero or negative" in str(e)


REFERENCES_HUMAN_READABLE_CASES = [
    (
        [
            {"agentGuid": "0c906be0-224c-45d4-8e6f-bc89da69d268", "date": 1508081600445},
            {"agentGuid": "{3a6f574a-3e6f-436d-acd4-bcde336b054d}", "date": 1475873692},
        ],
        [
            {"Date": "2017-10-15 15:33:20", "AgentGuid": "0c906be0-224c-45d4-8e6f-bc89da69d268"},
            {"Date": "2016-10-07 20:54:52", "AgentGuid": "3a6f574a-3e6f-436d-acd4-bcde336b054d"},
        ],
    ),
]


@pytest.mark.parametrize("raw_result, expected_parsed_res", REFERENCES_HUMAN_READABLE_CASES)
def test_references_to_human_readable(raw_result, expected_parsed_res):
    """
    Given:
        - An example of raw result from the API that represents the references of a hash
    When:
        - The function references_to_human_readable is called to convert the raw response to human readable data
    Then:
        - Assert that the parsed raw result is of the correct form
    """
    from McAfeeTIEV2 import references_to_human_readable

    result = references_to_human_readable(raw_result)
    assert expected_parsed_res == result


FILE_REFERENCES_CASES = [
    (
        {"hashes": ["hash1", "hash2"]},
        [
            [
                {"agentGuid": "0c906be0-224c-45d4-8e6f-bc89da69d268", "date": 1508081600445},
                {"agentGuid": "{68125cd6-a5d8-11e6-348e-000c29663178}", "date": 1478626172},
            ],
            [
                {"agentGuid": "{3a6f574a-3e6f-436d-acd4-bcde336b054d}", "date": 1475873692},
                {"agentGuid": "70be2ee9-7166-413b-b03e-64a48f6ab6c8", "date": 1508081651295},
            ],
        ],
        [
            {
                "Hash": "hash1",
                "References": [
                    {"Date": "2017-10-15 15:33:20", "AgentGuid": "0c906be0-224c-45d4-8e6f-bc89da69d268"},
                    {"Date": "2016-11-08 17:29:32", "AgentGuid": "68125cd6-a5d8-11e6-348e-000c29663178"},
                ],
            },
            {
                "Hash": "hash2",
                "References": [
                    {"Date": "2016-10-07 20:54:52", "AgentGuid": "3a6f574a-3e6f-436d-acd4-bcde336b054d"},
                    {"Date": "2017-10-15 15:34:11", "AgentGuid": "70be2ee9-7166-413b-b03e-64a48f6ab6c8"},
                ],
            },
        ],
    ),
    ({"hashes": ["hash1"]}, [[]], [None]),
]


@pytest.mark.parametrize("params, raw_result, expected_parsed_results", FILE_REFERENCES_CASES)
def test_files_references(mocker, params, raw_result, expected_parsed_results):
    """
    Given:
        - TIE client and a list of hashes
    When:
        - The function file_references_command is called to retrieve the references of the given files
    Then:
        - Assert that the parsed raw result is of the correct form
    """
    from CommonServerPython import Common
    from McAfeeTIEV2 import MAX_QUERY_LIMIT, files_references_command

    query_limit = MAX_QUERY_LIMIT

    mocker.patch("McAfeeTIEV2.get_hash_type_key", return_value="hash_type_key")
    mocker.patch("McAfeeTIEV2.get_file_instance", return_value=Common.File(dbot_score=None))
    with patch("dxltieclient.TieClient") as mock_tie_client:
        tie_client = mock_tie_client.return_value
        mocker.patch.object(tie_client, "get_file_first_references", side_effect=raw_result)
        results = files_references_command(hashes=params["hashes"], tie_client=tie_client, query_limit=query_limit)
        for exp_parsed_res, result in zip(expected_parsed_results, results):
            assert exp_parsed_res == result.outputs


TEST_DATA_FILES_CASES = [
    ("test_data/files_reputations/raw_result.json", "test_data/files_reputations/parsed_results.json"),
    ("test_data/files_reputations/raw_result_empty.json", "test_data/files_reputations/parsed_results_empty.json"),
]


@pytest.mark.parametrize("raw_result_file, parsed_results_file", TEST_DATA_FILES_CASES)
def test_files_reputations(mocker, raw_result_file, parsed_results_file):
    """
    Given:
        - TIE client and a list of hashes
    When:
        - The function files_reputations_command is called to retrieve the reputations of the given files
    Then:
        - Assert that raw result is parsed to the correct form
    """
    from McAfeeTIEV2 import files_reputations_command

    raw_responses = util_load_json(raw_result_file)
    parsed_results = util_load_json(parsed_results_file)
    mocker.patch("McAfeeTIEV2.safe_get_file_reputation", side_effect=raw_responses["raw_results"])
    expected_command_results = parsed_results["parsed_results"]
    with patch("dxltieclient.TieClient") as mock_tie_client:
        tie_client = mock_tie_client.return_value
        results = files_reputations_command(
            hashes=raw_responses["hashes"], tie_client=tie_client, reliability="C - Fairly reliable"
        )
        for exp_command_res, result in zip(expected_command_results, results):
            to_context = result.to_context()
            assert exp_command_res.get("HumanReadable") == to_context.get("HumanReadable")
            assert exp_command_res.get("EntryContext") == to_context.get("EntryContext")