from IsolateEndpoint import * import pytest from unittest.mock import patch def test_get_all_values_returns_expected_list(): """ Ensure get_all_values returns all brand values in the correct order. """ expected_list = [ "FireEyeHX v2", "CrowdstrikeFalcon", "Cortex Core - IR", "Builtin", "Microsoft Defender Advanced Threat Protection", ] assert Brands.get_all_values() == expected_list @pytest.mark.parametrize( "brand_enum, expected_value", [ (Brands.FIREEYE_HX_V2, "FireEyeHX v2"), (Brands.CROWDSTRIKE_FALCON, "CrowdstrikeFalcon"), (Brands.CORTEX_CORE_IR, "Cortex Core - IR"), (Brands.MICROSOFT_DEFENDER_ADVANCED_THREAT_PROTECTION, "Microsoft Defender Advanced Threat Protection"), ], ) def test_individual_enum_values(brand_enum, expected_value): """ Ensure each enum member returns the correct value (parametrized). """ assert brand_enum.value == expected_value @pytest.mark.parametrize( "endpoint_data, expected_output", [ ( {"Hostname": "host123", "ID": "endpoint", "IPAddress": "8.8.1.1", "Brand": "brand", "Message": "Fail"}, { "endpoint_id": "endpoint", "endpoint_hostname": "host123", "endpoint_ip": "8.8.1.1", "endpoint_brand": "brand", "endpoint_message": "Fail", }, ), ], ) def test_get_args_from_endpoint_data(endpoint_data, expected_output): """ Given: Endpoint data where values are either dictionaries or lists of dictionaries. When: The get_args_from_endpoint_data function is called. Then: It extracts and returns the correct values in a structured dictionary. """ result = get_args_from_endpoint_data(endpoint_data) assert result == expected_output def test_structure_endpoints_data(): """ Given: Various formats of `get_endpoint_data_results`, including a dict, a list with multiple elements, and None values. When: The structure_endpoints_data function is called. Then: It returns a properly structured list with the expected values. """ input_data = {"key": "value"} expected_output = [{"key": "value"}] assert structure_endpoints_data(input_data) == expected_output input_data = [None, {"key2": "value2"}] expected_output = [{"key2": "value2"}] assert structure_endpoints_data(input_data) == expected_output input_data = None expected_output = [] assert structure_endpoints_data(input_data) == expected_output @patch("IsolateEndpoint.create_message_to_context_and_hr") def test_check_missing_executed_args_in_output(mock_create_message): """ Given: - Different cases where `zipped_args` contain endpoint details that may or may not be in `valid_args`. When: - The `check_which_args_missing_in_output` function is called. Then: - It should call `create_message_to_context_and_hr` when an endpoint is missing. - It should not call `create_message_to_context_and_hr` when an endpoint is found. """ executed_args = [{"endpoint_id": "123", "endpoint_ip": "192"}, {"endpoint_id": "789", "endpoint_ip": "193"}] outputs = [] zipped_args = [ {"endpoint_id": "", "endpoint_ip": "194"}, {"endpoint_id": "555", "endpoint_ip": "195"}, {"endpoint_id": "123", "endpoint_ip": ""}, {"endpoint_id": "", "endpoint_ip": "192"}, {"endpoint_id": "", "endpoint_ip": ""}, {"endpoint_id": "", "endpoint_ip": "192"}, {"endpoint_id": "456", "endpoint_ip": ""}, ] check_missing_executed_args_in_output(zipped_args, executed_args, outputs) assert mock_create_message.call_count == 4 def test_map_zipped_args(): """ Given: Three lists of endpoint_ids, endpoint_ips, and endpoint_hostnames with varying lengths. When: The map_zipped_args function is called. Then: It correctly maps the elements into a list of dictionaries, filling missing values with empty strings. """ endpoint_ids = ["123", "456"] endpoint_ips = ["192.168.1.1", "192.168.1.2"] expected_output = [ {"endpoint_id": "123", "endpoint_ip": "192.168.1.1"}, {"endpoint_id": "456", "endpoint_ip": "192.168.1.2"}, ] assert map_zipped_args(endpoint_ids, endpoint_ips) == expected_output endpoint_ids = ["123"] endpoint_ips = ["192.168.1.1", "192.168.1.2"] expected_output = [ {"endpoint_id": "123", "endpoint_ip": "192.168.1.1"}, {"endpoint_id": "", "endpoint_ip": "192.168.1.2"}, ] assert map_zipped_args(endpoint_ids, endpoint_ips) == expected_output def test_map_args(): """ Given: - A Command object with `arg_mapping` defining how to map keys in `args`. - Optional hard-coded arguments that should be included in the output. When: - The `map_args` function is called. Then: - It correctly maps the values from `args` based on `arg_mapping`. - It includes hard-coded arguments in the output. - It returns an empty string for missing keys instead of raising an error. """ base_command = Command(brand="test_brand", name="test_command", arg_mapping={}) base_command.arg_mapping = {"new_key1": "old_key1", "new_key2": "old_key2"} args = {"old_key1": "value1", "old_key2": "value2"} expected_output = {"new_key1": "value1", "new_key2": "value2"} assert map_args(base_command, args) == expected_output base_command.arg_mapping = {"new_key1": "old_key1", "new_key2": "missing_key"} args = {"old_key1": "value1"} expected_output = {"new_key1": "value1", "new_key2": ""} assert map_args(base_command, args) == expected_output base_command.arg_mapping = {} assert map_args(base_command, {}) == {} base_command.arg_mapping = {"new_key": "old_key"} assert map_args(base_command, {}) == {"new_key": ""} base_command.arg_mapping = {"new_key1": "old_key1"} base_command.hard_coded_args = {"fixed_key": "fixed_value"} args = {"old_key1": "value1"} expected_output = {"new_key1": "value1", "fixed_key": "fixed_value"} assert map_args(base_command, args) == expected_output def test_are_there_missing_args(): """ Given: - A Command object with arg_mapping defining expected argument keys. When: - The function checks if all mapped arguments are missing. Then: - It correctly identifies when arguments are missing or present. """ base_command = Command(brand="test_brand", name="test_command", arg_mapping={}) base_command.arg_mapping = {"new_key1": "old_key1", "new_key2": "old_key2"} args = {"old_key1": "value1", "old_key2": "value2"} assert are_there_missing_args(command=base_command, endpoint_args=args, endpoint_output={}) is False base_command.arg_mapping = {"new_key1": "old_key1", "new_key2": "missing_key"} args = {"old_key1": "value1"} assert are_there_missing_args(command=base_command, endpoint_args=args, endpoint_output={}) is False base_command.arg_mapping = {"new_key1": "old_key1", "new_key2": "old_key2"} endpoint_output = {} assert are_there_missing_args(command=base_command, endpoint_args={}, endpoint_output=endpoint_output) is True assert endpoint_output.get("Result") == "Fail" assert endpoint_output.get("Message") == "Missing args for test_command." base_command.arg_mapping = {} assert are_there_missing_args(command=base_command, endpoint_args={}, endpoint_output={}) is False def test_is_endpoint_already_isolated(): """ Given: - Various endpoint data scenarios. When: - Checking if the endpoint can be isolated. Then: - Return the correct boolean value and message based on the conditions. """ endpoint_data = {"IsIsolated": "No"} assert is_endpoint_already_isolated(endpoint_data, endpoint_args={}, endpoint_output={}) is False endpoint_data["IsIsolated"] = "Yes" assert is_endpoint_already_isolated(endpoint_data, endpoint_args={}, endpoint_output={}) is True @pytest.mark.parametrize( "is_error, get_error_msg, raw_response, expected_is_isolated, expected_result, expected_message", [ # Error case ( True, "Some error occurred", {"status": "error"}, False, "Fail", "Failed to isolate 1234 with command TestCommand.Error:Some error occurred", ), # Success case ( False, "", {"status": "ok"}, True, "Success", "1234 was isolated successfully with command TestCommand.", ), ], ) @patch("IsolateEndpoint.is_error") @patch("IsolateEndpoint.get_error") @patch("IsolateEndpoint.create_message_to_context_and_hr") def test_handle_raw_response_results( mock_create_message, mock_get_error, mock_is_error, is_error, get_error_msg, raw_response, expected_is_isolated, expected_result, expected_message, ): """ Given: - Different combinations of raw response data and error flags. - Case 1: is_error=True with an error message. - Case 2: is_error=False for a success scenario. When: - Calling handle_raw_response_results with a Command object, endpoint args, and mocked helper functions. Then: - create_message_to_context_and_hr should be called exactly once with the expected parameters, including correct isolation status, result, and message. """ command = Command(brand="BrandA", name="TestCommand", arg_mapping={}) args = {"endpoint_id": "1234"} outputs = {} mock_is_error.return_value = is_error mock_get_error.return_value = get_error_msg handle_raw_response_results(command=command, raw_response=raw_response, endpoint_args=args, endpoint_output=outputs) mock_create_message.assert_called_once_with( is_isolated=expected_is_isolated, endpoint_args=args, result=expected_result, message=expected_message, endpoint_output=outputs, ) def test_initialize_commands(): """ Given: - The initialize_commands function is called to initialize a list of command objects. When: - Running the test_initialize_commands function to validate the list of command names and their associated brands. Then: - Ensure the actual command names match the expected set of command names. - Ensure each command has the correct brand associated with its name. """ commands = initialize_commands() expected_commands = { "core-isolate-endpoint": "Cortex Core - IR", "cs-falcon-contain-host": "CrowdstrikeFalcon", "fireeye-hx-host-containment": "FireEyeHX v2", "microsoft-atp-isolate-machine": "Microsoft Defender ATP", } # Validate names actual_command_names = {cmd.name for cmd in commands} assert set(expected_commands.keys()) == actual_command_names # Validate brands for cmd in commands: expected_brand = expected_commands.get(cmd.name) assert cmd.brand == expected_brand def test_find_command_by_brand(): """ Given: - A list of Command objects with different brand names. When: - Calling the find_command_by_brand function with the brand 'BrandB'. Then: - Ensure the function returns the Command object with brand 'BrandB'. """ command_a = Command(brand="BrandA", name="command-a", arg_mapping={}) command_b = Command(brand="BrandB", name="command-b", arg_mapping={}) result = find_command_by_brand(commands=[command_a, command_b], brand="BrandB") assert result == command_b @pytest.mark.parametrize( "endpoint_args,is_isolated,result,message,expected_source,expected_isolated", [ ( {"endpoint_id": "1234", "endpoint_brand": "SomeBrand"}, True, "Success", "Test Message", "SomeBrand", "Yes", ), # Microsoft Defender ATP case ( {"endpoint_id": "5678", "endpoint_brand": "Microsoft Defender ATP"}, False, "Fail", "Converted Brand Test", Brands.MICROSOFT_DEFENDER_ADVANCED_THREAT_PROTECTION, "No", ), ], ) def test_create_message_to_context_and_hr(endpoint_args, is_isolated, result, message, expected_source, expected_isolated): """ Given: - Endpoint arguments with various endpoint_brand values (including 'Microsoft Defender ATP' for conversion). When: - Calling create_message_to_context_and_hr with different is_isolated flags and results. Then: - The endpoint_output should contain the expected Endpoint, Result, Source (converted if needed), Message, and Isolated status. """ endpoint_output = {} create_message_to_context_and_hr( is_isolated=is_isolated, endpoint_args=endpoint_args, result=result, message=message, endpoint_output=endpoint_output, ) expected = { "Endpoint": endpoint_args["endpoint_id"], "Result": result, "Source": expected_source, "Message": message, "Isolated": expected_isolated, } assert endpoint_output == expected @patch("IsolateEndpoint.demisto.executeCommand") # @patch("IsolateEndpoint.handle_raw_response_results") def test_run_commands_for_endpoint_executes_command(mock_execute): """ Given: A list containing a single Command object with a specific brand, name, and argument mapping. When: Calling run_commands_for_endpoint with the mock command, endpoint arguments, and an empty results list. Then: - executeCommand should be called exactly once with the mapped arguments. - handle_raw_response_results should be called exactly once with the execution result. - The processed command result ("CommandResult") should be appended to the results list. """ command = Command(brand="Brand", name="test-command", arg_mapping={"endpoint_id": "endpoint_id"}) mock_commands = [command] mock_execute.return_value = [{"Type": 1, "Contents": "Done"}] endpoint_args = {"endpoint_id": "id1", "endpoint_brand": "Brand"} endpoint_output = {} run_commands_for_endpoint(commands=mock_commands, endpoint_args=endpoint_args, endpoint_output=endpoint_output) mock_execute.assert_called_once() assert endpoint_output["Endpoint"] == "id1" assert endpoint_output["Result"] == "Success" assert endpoint_output["Source"] == "Brand" @pytest.mark.parametrize( "endpoint_data_results, is_already_isolated, has_fail, expected_context_len", [ # Case 1: endpoint fails -> skipped entirely ( [{"id": "ep1", "properties": {"message": "Failing"}}], False, True, 0, ), # Case 2: endpoint already isolated -> included in context but no run_commands ( [{"id": "ep2", "properties": {"message": "ok"}}], True, False, 1, ), # Case 3: normal endpoint -> should run commands ( [{"id": "ep3", "properties": {"message": "ok"}}], False, False, 1, ), ], ) @patch("IsolateEndpoint.get_args_from_endpoint_data") @patch("IsolateEndpoint.is_endpoint_already_isolated") @patch("IsolateEndpoint.run_commands_for_endpoint") def test_process_endpoints( mock_run_commands, mock_is_already_isolated, mock_get_args, endpoint_data_results, is_already_isolated, has_fail, expected_context_len, ): """ Given: - A list of endpoint data results with different conditions: 1. Endpoint marked as failing. 2. Endpoint already isolated. 3. Normal endpoint to run commands. When: - Calling process_endpoints with mocked helper functions. Then: - Endpoints with "fail" in their message are skipped. - Already isolated endpoints are added to args/context but do not run commands. - Normal endpoints trigger run_commands_for_endpoint and are added to args/context. """ mock_args = {"endpoint_id": "1234", "endpoint_message": "Fail" if has_fail else "ok"} mock_get_args.return_value = mock_args mock_is_already_isolated.return_value = is_already_isolated commands = [Command(brand="BrandA", name="TestCommand", arg_mapping={})] results, context_outputs, args_from_endpoint_data = process_endpoints(endpoint_data_results, commands) assert len(context_outputs) == expected_context_len if not has_fail and not is_already_isolated: mock_run_commands.assert_called_once_with(commands, mock_args, context_outputs[0]) else: mock_run_commands.assert_not_called() @pytest.mark.parametrize( "args, brands_values, expect_error, expected_brands", [ # Case 1: endpoint_id with explicit brands -> keep brands as-is ( {"endpoint_id": ["1234"], "endpoint_ip": [], "brands": ["Cortex XDR"]}, ["BrandA", "BrandB"], False, ["Cortex XDR"], ), # Case 2: endpoint_id but no brands -> use Brands.get_all_values() ( {"endpoint_id": ["5678"], "endpoint_ip": [], "brands": []}, ["BrandA", "BrandB"], False, ["BrandA", "BrandB"], ), ], ) @patch("IsolateEndpoint.map_zipped_args") @patch("IsolateEndpoint.Brands.get_all_values") @patch("IsolateEndpoint.demisto.args") def test_prepare_args( mock_args, mock_get_all_values, mock_map_zipped, args, brands_values, expect_error, expected_brands, ): """ Given: - Different combinations of endpoint arguments. 1. No endpoint_id or endpoint_ip (invalid). 2. endpoint_id with brands provided. 3. endpoint_id but no brands (default brands applied). When: - Calling prepare_args with mocked demisto.args, Brands.get_all_values, and map_zipped_args. Then: - A ValueError is raised if no endpoint_id or endpoint_ip is supplied. - Brands remain unchanged if explicitly provided. - Brands default to Brands.get_all_values() when not provided. """ mock_args.return_value = args.copy() mock_get_all_values.return_value = brands_values mock_map_zipped.return_value = [("endpoint_id", "endpoint_ip")] if expect_error: with pytest.raises(ValueError, match="At least one of the following arguments must be specified"): prepare_args() else: endpoint_args, zipped_args = prepare_args() assert endpoint_args["brands"] == expected_brands assert zipped_args == [("endpoint_id", "endpoint_ip")]