import pytest from DHS_Feed import * def compare(object_a, object_b): if isinstance(object_a, List): return compare_list(object_a, object_b) elif isinstance(object_a, Dict): return compare_dict(object_a, object_b) else: return object_a == object_b def compare_list(list_a, list_b): try: list_a = sorted(list_a) list_b = sorted(list_b) except TypeError: pass if len(list_a) != len(list_b): return False return all(compare(a_obj, b_obj) for a_obj, b_obj in zip(list_a, list_b)) def compare_dict(dict_a, dict_b): keys = dict_a.keys() if not compare_list(keys, dict_b.keys()): return False return all(compare(dict_a[key], dict_b[key]) for key in keys) class TestTempFile: def test_create_file(self): data = "test" temp_file = TempFile(data) with open(temp_file.path) as _file: assert _file.read() == data, "temp file content failed" def test_removing_file(self): temp_file_name = TempFile("test").path assert not os.path.isfile(temp_file_name), "file was not removed at the ens of live of the TempFile objet." @pytest.mark.parametrize("suffix", ["test", "py", "pem", "crt"]) def test_suffix(self, suffix): temp_file_name = TempFile("test", suffix=suffix).path assert temp_file_name.endswith(suffix), "file suffix dos not working as expected" class TestHelpers: data_test_fix_rsa_data = [ ("test_data/rsa/test_fix_rsa_data.txt", 4), ("test_data/rsa/test_fix_rsa_data2.txt", 2), ] @pytest.mark.parametrize("path, count", data_test_fix_rsa_data) def test_fix_rsa_data(self, path, count): with open(path) as _file: data = _file.read() demisto_data = data.replace("\n", " ") fixed_data = fix_rsa_data(demisto_data, count) assert data == fixed_data, "failed to parse the data from demisto params to RSA file" data_test_insert_id = [ ("te{ID}st", "tetestst"), ("{ID}", "test"), ] @pytest.mark.parametrize("input_str, expected_output", data_test_insert_id) def test_insert_id(self, input_str, expected_output, mocker): mocker.patch.object(uuid, "uuid4", return_value="test") output = insert_id(input_str) assert output == expected_output, "failed to insert uuid" data_test_ssl_files_checker = [("test_data/rsa/2048b-rsa-example-keypair.pem", "test_data/rsa/2048b-rsa-example-cert.pem")] @pytest.mark.parametrize("input_key, input_public", data_test_ssl_files_checker) def test_ssl_files_checker(self, input_key, input_public): with open(input_key) as input_key, open(input_public) as input_public: ssl_files_checker(input_public.read(), input_key.read()) @pytest.mark.parametrize("input_key, input_public", data_test_ssl_files_checker) def test_ssl_files_checker_with_invalid_files(self, input_key, input_public): with open(input_key) as input_key: input_key = input_key.read() with open(input_public) as input_public: input_public = input_public.read() with pytest.raises(ValueError): temp_input_public = input_public.split("\n")[:-6] temp_input_public.extend(input_public.split("\n")[-7:]) ssl_files_checker("\n".join(temp_input_public), input_key) with pytest.raises(ValueError): temp_input_private = input_key.split("\n")[:-6] temp_input_private.extend(input_key.split("\n")[-7:]) ssl_files_checker(input_public, "\n".join(temp_input_private)) class TestSafeDataGet: data_test_safe_data_get = [("without_list", "one_level_get"), (["standard_list", "test"], "multi_level_get")] @pytest.mark.parametrize("get_list, expected_output", data_test_safe_data_get) def test_safe_data_get(self, get_list, expected_output): dict_data = {"without_list": "one_level_get", "standard_list": {"test": "multi_level_get"}} output = safe_data_get(dict_data, get_list) assert output == expected_output, "failed to get the relevant data" def test_without_list(self): dict_data = {"without_list": "one_level_get"} output = safe_data_get(dict_data, "without_list") assert output == "one_level_get" def test_with_list(self): dict_data = {"standard_list": {"test": "multi_level_get"}} output = safe_data_get(dict_data, ["standard_list", "test"]) assert output == "multi_level_get" def test_one_level_with_prefix(self): dict_data = {"TEST:without_list": "one_level_get"} output = safe_data_get(dict_data, "without_list", prefix="TEST") assert output == "one_level_get" def test_multi_level_with_prefix(self): dict_data = {"TEST:standard_list": {"TEST:test": "multi_level_get"}} output = safe_data_get(dict_data, ["standard_list", "test"], prefix="TEST") assert output == "multi_level_get" data_test_get_none_existing_path_with_prefix = [ ({"TEST:standard_list": {"TEST:test": "multi_level_get"}}, ["standard_list1", "test"]), ({"TEST:standard_list": {"TEST:test": "multi_level_get"}}, ["standard_list", "test1"]), ({"without_list": "one_level_get"}, "without_list1"), ] @pytest.mark.parametrize("dict_data, path", data_test_get_none_existing_path_with_prefix) def test_get_none_existing_path_with_prefix(self, dict_data, path): output = safe_data_get(dict_data, path, prefix="TEST") assert output is None data_test_get_none_existing_path_with_prefix2 = [ ({"standard_list": {"test": "multi_level_get"}}, ["standard_list1", "test"]), ({"standard_list": {"test": "multi_level_get"}}, ["standard_list", "test1"]), ({"without_list": "one_level_get"}, "without_list1"), ] @pytest.mark.parametrize("dict_data, path", data_test_get_none_existing_path_with_prefix2) def test_get_none_existing_path_without_prefix(self, dict_data, path): output = safe_data_get(dict_data, path) assert output is None data_test_with_default_value = [ ({}, "test", "test"), ({}, None, None), ({"test": None}, "something", None), ({"something": "nothing"}, "something", "something"), ] @pytest.mark.parametrize("dict_data, default, expected_output", data_test_with_default_value) def test_with_default_value(self, dict_data, default, expected_output): output = safe_data_get(dict_data, "test", default=default) assert output == expected_output class TestIndicators: @staticmethod def read_json(path): with open(path) as json_file: json_file = json_file.read() return json.loads(json_file) @staticmethod def get_stix_header(block): return block.get("stix:STIX_Header", {}) File = "File" IP = "IP" Domain = "Domain" URL = "URL" Email = "Email" data_types = [IP, Domain, URL, Email, File] @pytest.mark.parametrize("data_type", data_types) def test_data_to_blocks(self, data_type): data = self.read_json(f"test_data/data_from_DHS/{data_type}_data.json") test_blocks = Indicators._blocks(data) blocks = self.read_json(f"test_data/blocks/blocks_from_{data_type}_data.json") assert test_blocks == blocks @pytest.mark.parametrize("data_type", data_types) def test_blocks_to_indicators(self, data_type): blocks = self.read_json(f"test_data/blocks/blocks_from_{data_type}_data.json") test_indicators = [] for block in blocks: test_indicators.extend(list(Indicators._indicators(block))) indicators = self.read_json(f"test_data/indicators/indicators_from_{data_type}_data.json") assert test_indicators == indicators @pytest.mark.parametrize("data_type", data_types) def test_indicators_to_indicator_data(self, data_type): indicators = self.read_json(f"test_data/indicators/indicators_from_{data_type}_data.json") test_data_indicators = [Indicators._indicator_data(x, "source", "color", ["tag"]) for x in indicators] data_indicators = self.read_json(f"test_data/data_indicators/{data_type}_data_indicators.json") assert test_data_indicators == data_indicators @pytest.mark.parametrize("data_type", data_types) def test_indicators_to_context_indicators(self, data_type): indicators = self.read_json(f"test_data/data_indicators/{data_type}_data_indicators.json") test_context_indicators = list(map(indicator_to_context, indicators)) context_indicators = self.read_json(f"test_data/context_indicators/context_from_{data_type}_data.json") assert test_context_indicators == context_indicators data_test_tlp_color_from_header = [ (IP, ["WHITE", "WHITE", "WHITE"]), (Domain, ["GREEN"]), (URL, ["AMBER"]), (Email, ["AMBER"]), (File, ["GREEN", "GREEN", "GREEN"]), ] @pytest.mark.parametrize("data_type, tlp_colors", data_test_tlp_color_from_header) def test_tlp_color_from_header(self, data_type, tlp_colors): blocks_headers = list(map(self.get_stix_header, self.read_json(f"test_data/blocks/blocks_from_{data_type}_data.json"))) test_tlp_colors = list(map(Indicators._tlp_color_from_header, blocks_headers)) assert test_tlp_colors == tlp_colors data_test_source_from_header = [(IP, ["Infoblox Inc"] * 3), (File, ["Reversing Labs"] * 3)] @pytest.mark.parametrize("data_type, sources", data_test_source_from_header) def test_source_from_header(self, data_type, sources): sources_headers = list(map(self.get_stix_header, self.read_json(f"test_data/blocks/blocks_from_{data_type}_data.json"))) test_sources = list(map(Indicators._source_from_header, sources_headers)) assert test_sources == sources class TestCommandTestModule: def nothing(self, *args, **kwargs): return lambda: self.data def setup_class(self): self.client = TaxiiClient("", "", "") self.data = None self.ssl_files_checker = self.get_first_fetch = self.discovery_request = None def mock_data(self, mocker): self.ssl_files_checker = mocker.patch("DHS_Feed.ssl_files_checker", new=self.nothing) self.get_first_fetch = mocker.patch("DHS_Feed.get_first_fetch", new=self.nothing) self.discovery_request = mocker.patch.object(self.client, "discovery_request", new_callable=self.nothing) def test_command_test_module(self, mocker): self.mock_data(mocker) self.data = {"taxii_11:Discovery_Response": {"taxii_11:Service_Instance": ["somthing"]}} assert command_test_module(self.client, "", "", "") == "ok" def test_command_test_module_with_invalid_credential(self, mocker): self.mock_data(mocker) self.data = {"taxii_11:Status_Message": {"@status_type": "UNAUTHORIZED"}} with pytest.raises(DemistoException, match="invalid credential."): command_test_module(self.client, "", "", "") def test_command_test_module_with_unknown_error(self, mocker): self.mock_data(mocker) self.data = {} with pytest.raises(DemistoException, match="unknown error."): command_test_module(self.client, "", "", "")