import json import demistomock as demisto import pytest from CommonServerPython import * from FeedThreatConnect import ( build_url_with_query_params, create_indicator_fields, create_or_query, create_types_query, get_indicator_value, get_updated_last_run, parse_indicator, set_fields_query, set_tql_query, should_send_request, ) from pytest_mock import MockerFixture def load_json_file(path): with open(path) as _json_file: return json.load(_json_file) @pytest.mark.parametrize( argnames="threatconnect_score, dbot_score", argvalues=[(1000, 3), (830, 3), (664, 2), (498, 2), (332, 1), (166, 1), (0, 0)] ) def test_calculate_dbot_score(threatconnect_score, dbot_score): from FeedThreatConnect import calculate_dbot_score assert calculate_dbot_score(threatconnect_score) == dbot_score def test_parse_indicator(mocker): mocker.patch.object(demisto, "params", return_value={"createRelationships": True, "tlpcolor": None}) data_dir = { "parsed_indicator.json": "./test_data/parsed_indicator.json", # type: ignore # noqa "indicators.json": "./test_data/indicators.json", } # type: ignore # noqa indicator = parse_indicator(load_json_file(data_dir["indicators.json"])) assert load_json_file(data_dir["parsed_indicator.json"]) == indicator def test_create_or_query(): assert create_or_query("test", "1,2,3,4,5") == 'test="1" OR test="2" OR test="3" OR test="4" OR test="5" ' @pytest.mark.parametrize( "params, expected_result, endpoint", [ ( { "indicator_active": False, "indicator_type": ["All"], "createRelationships": False, "confidence": 0, "threat_assess_score": 0, }, 'typeName IN ("EmailAddress","File","Host","URL","ASN","CIDR","Hashtag","Mutex","Registry Key",' '"User Agent","Address")', "indicators", ), # noqa: E501 ( { "indicator_active": True, "group_type": ["File"], "createRelationships": False, "confidence": 0, "threat_assess_score": 0, }, 'typeName IN ("File")', "groups", ), ( { "indicator_active": False, "group_type": ["Tool"], "createRelationships": False, "confidence": 50, "threat_assess_score": 80, }, 'typeName IN ("Tool")', "groups", ), ], ) def test_set_tql_query(params, expected_result, endpoint): """ Given: - an empty from_date value and demisto params Case 1: expecting no tql query Case 2: expecting a specific group type, and only active indicators When: - running set_tql_query command Then: - validate the tql output """ from_date = "" output = set_tql_query(from_date, params, endpoint) assert output == expected_result @pytest.mark.parametrize( "params, expected_result, endpoint", [ ( {"group_type": ["All"], "indicator_type": []}, 'typeName IN ("Attack Pattern","Campaign",' '"Course of Action","Intrusion Set","Malware","Report","Tool","Vulnerability")', "groups", ), ({"group_type": ["File"], "indicator_type": []}, 'typeName IN ("File")', "groups"), ({"group_type": ["File"], "indicator_type": ["All"]}, 'typeName IN ("File")', "groups"), ], ) def test_create_types_query(params, expected_result, endpoint): """ Given: - demisto params and an endpoint When: - running create_types_query command Then: - validate the output """ output = create_types_query(params, endpoint) assert output == expected_result @pytest.mark.parametrize( "params, expected_result, endpoint", [ ({"group_type": ["All"], "indicator_type": []}, False, "indicators"), ({"group_type": [], "indicator_type": ["All"]}, True, "indicators"), ], ) def test_should_send_request(params, expected_result, endpoint): """ Given: - demisto params and an endpoint When: - running should_send_request command Then: - validate the result """ output = should_send_request(params, endpoint) assert output == expected_result @pytest.mark.parametrize( "params, expected_result, endpoint", [ pytest.param( { "indicator_type": ["All"], "indicator_query": "", "group_query": "", "use_indicator_query_for_group": True, "createRelationships": False, }, "/api/v3/indicators?tql=indicatorActive%20EQ%20True&fields=tags&fields=threatAssess&resultStart" "=0&resultLimit=100&sorting=dateAdded%20ASC", "indicators", id="indicators with indicator query", ), pytest.param( { "group_type": ["All"], "indicator_query": "indicatorActive EQ False", "group_query": "", "use_indicator_query_for_group": True, "createRelationships": True, }, "/api/v3/groups?tql=indicatorActive%20EQ%20False&fields=tags&fields=associatedGroups" "&fields=associatedIndicators&resultStart=0&resultLimit=100&sorting=dateAdded%20ASC", "groups", id="groups with indicator query", ), pytest.param( { "group_type": ["All"], "indicator_query": "", "group_query": "Test", "use_indicator_query_for_group": True, "createRelationships": False, }, "/api/v3/groups?tql=indicatorActive%20EQ%20True&fields=tags&resultStart=0&resultLimit=100&sorting=dateAdded%20ASC", "groups", id="groups with default query", ), pytest.param( { "group_type": ["All"], "indicator_query": "indicatorActive EQ False", "group_query": "Test", "use_indicator_query_for_group": False, "createRelationships": True, }, "/api/v3/groups?tql=Test&fields=tags&fields=associatedGroups" "&fields=associatedIndicators&resultStart=0&resultLimit=100&sorting=dateAdded%20ASC", "groups", id="groups with group query", ), ], ) def test_build_url_with_query_params(mocker: MockerFixture, params: dict, expected_result: str, endpoint: str): """ Given: - demisto params and an endpoint When: - running build_url_with_query_params command Then: - validate the result """ mocker.patch("FeedThreatConnect.set_tql_query", return_value="indicatorActive EQ True") # Determine the tql parameter based on endpoint and params if endpoint == "indicators": tql = params.get("indicator_query") or None else: # groups tql = params.get("group_query") if not params.get("use_indicator_query_for_group") else params.get("indicator_query") output = build_url_with_query_params(params, endpoint, {}, tql) assert output == expected_result @pytest.mark.parametrize( "params, expected_result, endpoint", [ ({"createRelationships": False}, "&fields=tags&fields=threatAssess", "indicators"), ({"createRelationships": True}, "&fields=tags&fields=associatedGroups&fields=associatedIndicators", "groups"), ], ) def test_set_fields_query(params, expected_result, endpoint): """ Given: - demisto params and an endpoint When: - running set_fields_query command Then: - validate the result """ output = set_fields_query(params, endpoint) assert output == expected_result @pytest.mark.parametrize( "indicators, groups, previous_run, expected_result", [ ( [{"dateAdded": "dateAdded"}], [{"dateAdded": "dateAdded"}], {}, {"indicators": {"from_date": "dateAdded"}, "groups": {"from_date": "dateAdded"}}, ), ( ( [{"dateAdded": "dateAdded"}], [], {"groups": {"from_date": "from_date"}}, {"indicators": {"from_date": "dateAdded"}, "groups": {"from_date": "from_date"}}, ) ), ], ) def test_get_updated_last_run(indicators, groups, previous_run, expected_result): """ Given: - list of indicators, list of groups, and a previouse run When: - running get_updated_last_run command Then: - validate the result """ output = get_updated_last_run(indicators, groups, previous_run) assert output == expected_result def test_create_indicator_fields_registry_key(): """ Given: - lindicator from type Registry Key When: - running create_indicator_fields command Then: - validate the result contains the 'Key Value' key and the expected data """ indicator = { "Key Name": "key name", "Value Name": "value name", "Key Type": "key type", "dateAdded": "firstseenbysource", "lastModified": "updateddate", "threatAssessRating": "verdict", "threatAssessConfidence": "confidence", "description": "description", "summary": "name", } result = create_indicator_fields(indicator, "Registry Key") assert "Key Value" in result assert "name" in result.get("Key Value")[0] assert result.get("Key Value")[0].get("name") == "key name" def test_get_indicator_value_for_file(): """ Given: An indicator dictionary with file hashes. When: The indicator type is 'File'. Then: It should return the sha256 hash if present, else sha1, else md5. """ indicator = {"sha256": "sha256_hash", "sha1": "sha1_hash", "md5": "md5_hash"} indicator_type = FeedIndicatorType.File indicator_value = get_indicator_value(indicator, indicator_type) assert indicator_value == "sha256_hash" # Test when sha256 is not present del indicator["sha256"] indicator_value = get_indicator_value(indicator, indicator_type) assert indicator_value == "sha1_hash" # Test when sha256 and sha1 are not present del indicator["sha1"] indicator_value = get_indicator_value(indicator, indicator_type) assert indicator_value == "md5_hash" def test_get_indicator_value_for_non_file(mocker: MockerFixture): """ Given: An indicator dictionary without file hashes. When: The indicator type is not 'File'. Then: It should return the summary if present, else name. """ indicator = {"summary": "indicator_summary", "name": "indicator_name"} indicator_type = "IP" indicator_value = get_indicator_value(indicator, indicator_type) assert indicator_value == "indicator_summary" # Test when summary is not present mocker.patch.dict(indicator, {"summary": None}) indicator_value = get_indicator_value(indicator, indicator_type) assert indicator_value == "indicator_name"