import functools import demistomock as demisto import pytest from CarbonBlackLiveResponseCloud import * from cbc_sdk.live_response_api import LiveResponseMemdump, LiveResponseSessionManager CREDENTIALS = {"url": "https://test.test", "ssl_verify": False, "token": "test_key/test_id", "org_key": "test_org_key"} api_client = CBCloudAPI(**CREDENTIALS) commands_with_args = { command_test_module: {"api_client": api_client}, put_file_command: { "api_client": api_client, "device_id": "test_device_id", "destination_path": "test_destination_path", "file_id": "test_file_id", }, get_file_command: { "api_client": api_client, "device_id": "test_device_id", "source_path": "test_source_path", "timeout": "900", "delay": "900", }, delete_file_command: {"api_client": api_client, "device_id": "test_device_id", "source_path": "test_source_path"}, create_reg_key_command: {"api_client": api_client, "device_id": "test_device_id", "reg_path": "test_reg_path"}, set_reg_value_command: { "api_client": api_client, "device_id": "test_device_id", "reg_path": "test_reg_path", "value_data": "test_value_data", "value_type": "test_value_type", "overwrite": False, }, delete_reg_key_command: { "api_client": api_client, "device_id": "test_device_id", "reg_path": "test_reg_path", }, delete_reg_value_command: { "api_client": api_client, "device_id": "test_device_id", "reg_path": "test_reg_path", }, list_directory_command: { "api_client": api_client, "device_id": "test_device_id", "directory_path": "test_directory_path", "limit": 5, }, list_reg_sub_keys_command: { "api_client": api_client, "device_id": "test_device_id", "reg_path": "test_reg_path", "limit": "5", }, get_reg_values_command: {"api_client": api_client, "device_id": "test_device_id", "reg_path": "test_reg_path", "limit": "5"}, list_processes_command: {"api_client": api_client, "device_id": "test_device_id", "limit": 5}, kill_process_command: {"api_client": api_client, "device_id": "test_device_id", "pid": "100"}, create_process_command: { "api_client": api_client, "device_id": "test_device_id", "command_string": "test_cmd_line_path", "wait_for_output": "True", "wait_for_completion": True, }, memdump_command: {"api_client": api_client, "device_id": "test_device_id", "target_path": "test_target_path"}, } TEST_REG_VALUES = [ {"registry_type": f"test_pbREG_SZ_{i}", "registry_name": f"test_val_{i}", "registry_data": f"value_data_{i}"} for i in range(10) ] TEST_DIR_LIST = [ { "size": 25600, "attributes": ["TEST_ARCHIVE"], "create_time": 123, "last_access_time": 123, "last_write_time": "1970-01-02T03:46:40.000Z", "filename": f"test_{i}.xls", "alternate_name": "test_$9EE1B~1.XLS", } for i in range(10) ] TEST_PROCESSES = [ { "process_path": f"test_path_{i}", "process_pid": i, "process_cmdline": f"test_command_line_{i}", "process_username": f"test_user_{i}", } for i in range(10) ] SUB_KEY_LEN = 10 TEST_SUB_KEYS = {"sub_keys": [f"sub_{i}" for i in range(SUB_KEY_LEN)]} HAPPY_PATH_ARGS = [ ( "put_file", put_file_command, "File: test_file_id is successfully put to the remote destination test_destination_path", {}, None, ), ( "delete_file", delete_file_command, "The file: test_source_path was deleted successfully.", {"filename": "test_source_path"}, None, ), ("create_registry_key", create_reg_key_command, "Reg key: test_reg_path, was created successfully.", {}, None), ("set_registry_value", set_reg_value_command, "Value was set to the reg key: test_reg_path successfully.", {}, None), ("delete_registry_key", delete_reg_key_command, "Registry key: test_reg_path was deleted successfully.", {}, None), ("delete_registry_value", delete_reg_value_command, "Registry value: test_reg_path was deleted successfully.", {}, None), ( "kill_process", kill_process_command, "The process: 100 was killed successfully.", {}, "The process: 100 was killed successfully.", ), ] WRONG_ARGS = [ ("set_registry_value", set_reg_value_command, {"overwrite": "wrong_val"}), ("create_process", create_process_command, {"wait_for_completion": "wrong_val"}), ("create_process", create_process_command, {"wait_for_output": "wrong_val"}), ("create_process", create_process_command, {"wait_timeout": "wrong_val"}), ("get_file", get_file_command, {"timeout": "wrong_val"}), ("delete_registry_key", delete_reg_key_command, {"force": "wrong_val"}), ("list_directory", list_directory_command, {"limit": "wrong_val"}), ] DIR_LIST_EXPECTED_OUTPUT = ( "Name|Type|Date Modified|Size|\n|---|---|---|---|\n| test_0.xls | File | 1970-01-02T03:46:40.000Z | 25600" ) # HAPPY_PATH_ARGS_FOR_COMMAND_RESULTS = [ ("list_registry_keys_and_values", list_reg_sub_keys_command, "sub_1", (), TEST_SUB_KEYS), ("list_registry_values", get_reg_values_command, "test_pbREG_SZ_1", (), TEST_REG_VALUES), ("list_processes", list_processes_command, "test_path_1 | 1 | test_command_line_1", (), TEST_PROCESSES), ("list_directory", list_directory_command, DIR_LIST_EXPECTED_OUTPUT, (), TEST_DIR_LIST), ( "create_process", create_process_command, "פלט בעברית", {"command_string": "test_cmd_line_path", "wait_timeout": 30, "wait_for_output": True, "wait_for_completion": True}, "פלט בעברית".encode(), ), ] class MockedLRObject: def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): pass def lr_session(self): return self MOCKED_LR_SESSION = MockedLRObject() def raise_exception(exception_to_raise, *args): raise exception_to_raise def mock_method_in_lr_session(mocker, method_name, mocked_results=None): """ mocked the required method in CBCloudAPI and CbLRSessionBase :type mocker: ``Mock`` :param mocker: the pytest mock object :type method_name: ``str`` :param method_name: The name of the method in CbLRSessionBase to be mocked e.g put_file :type mocked_results: ``Any`` :param mocked_results: The result to return in call to the mocked method :rtype: ``Mock`` :return: the mocked object """ test_mocker = mocker.patch("cbc_sdk.CBCloudAPI.select", return_value=MOCKED_LR_SESSION) def mock_method(): lambda *args, **kwargs: None # noqa: B018 setattr(MOCKED_LR_SESSION, method_name, mock_method) # Just defined method in the object test_mocker = mocker.patch.object(MOCKED_LR_SESSION, method_name, return_value=mocked_results) mocker.patch.object(demisto, "getFilePath", return_value={"path": "test_data/test.txt", "name": "test"}) return test_mocker class TestCommands: @pytest.mark.parametrize( "api_method_to_be_mocked, tested_command, expected_result, expected_args, mocked_results", HAPPY_PATH_ARGS ) def test_commands_happy_path( self, mocker, api_method_to_be_mocked, tested_command, expected_result, expected_args, mocked_results ): """ Given: - Args for command to send via Live Response API (from commands_with_args) When: - Run the command Then: - Returns the command output """ mocked_obj = mock_method_in_lr_session(mocker=mocker, method_name=api_method_to_be_mocked, mocked_results=mocked_results) kwargs = commands_with_args[tested_command] res = tested_command(**kwargs) # assert result assert res == expected_result if expected_args: assert mocked_obj.call_args[1] == expected_args limit = arg_to_number(kwargs.get("limit")) if limit: assert len(res) <= limit @pytest.mark.parametrize("api_method_to_be_mocked, tested_command, wrong_args", WRONG_ARGS) def test_commands_wrong_args(self, mocker, api_method_to_be_mocked, tested_command, wrong_args): """ Given: - Args for command to send via Live Response API (from commands_with_args) When: - Run the command Then: - Returns the command output """ mock_method_in_lr_session( mocker=mocker, method_name=api_method_to_be_mocked, mocked_results=[], ) kwargs = commands_with_args[tested_command].copy() kwargs.update(wrong_args) with pytest.raises(ValueError): tested_command(**kwargs) @pytest.mark.parametrize( "api_method_to_be_mocked, tested_command, expected_result, expected_args, mocked_results", HAPPY_PATH_ARGS_FOR_COMMAND_RESULTS, ) def test_readable_outputs( self, mocker, api_method_to_be_mocked, tested_command, expected_result, expected_args, mocked_results ): # prepare mock_method_in_lr_session(mocker=mocker, method_name=api_method_to_be_mocked, mocked_results=mocked_results) # run kwargs = commands_with_args[tested_command] res = tested_command(**kwargs) # assert assert expected_result in res.readable_output limit = arg_to_number(kwargs.get("limit")) if limit: assert len(res.raw_response) <= limit def test_memdump_command(self, mocker): """ Given: - Args for memdump command to send via Live Response API (from commands_with_args) When: - Run the memdump command Then: - Returns the command output """ # prepare mocker.patch("cbc_sdk.live_response_api.LiveResponseMemdump.wait") mock_method_in_lr_session( mocker=mocker, method_name="start_memdump", mocked_results=LiveResponseMemdump(None, None, None), ) # run kwargs = commands_with_args[memdump_command] res = memdump_command(**kwargs) # assert assert res == "Memory was successfully dumped to test_target_path." @pytest.mark.parametrize("force, expected_call_count", [(False, 1), (True, 11)]) def test_delete_reg_key_force(self, mocker, force, expected_call_count): """ Given: - Reg key contain sub keys When: - run the delete reg key command with force arg true Then: - validate all sub keys was deleted """ # prepare def mock_sub_keys(reg_key): return {"sub_keys": ["sub_key"] if reg_key.count("\\") < SUB_KEY_LEN else []} mock_method_in_lr_session(mocker=mocker, method_name="list_registry_keys_and_values") mocker.patch.object(MOCKED_LR_SESSION, "list_registry_keys_and_values", side_effect=mock_sub_keys) mocked_method = mock_method_in_lr_session(mocker=mocker, method_name="delete_registry_key") # run kwargs = commands_with_args[delete_reg_key_command] kwargs["force"] = force delete_reg_key_command(**kwargs) assert mocked_method.call_count == expected_call_count def test_command_test_happy_path(self, mocker): """ Given: - Args for test_command When: - Run the test_command Then: - Validate the results """ # prepare kwargs = commands_with_args[command_test_module] mocker.patch.object(LiveResponseSessionManager, "request_session") # run res = command_test_module(**kwargs) # assert assert res == "ok" RAISE_EXCEPTION_PARAMS = [ (errors.UnauthorizedError(""), AUTHORIZATION_ERROR_MSG), (errors.ConnectionError(""), CONNECTION_ERROR_MSG), (errors.ObjectNotFoundError("wrong org_id"), ORG_ID_ERROR_MSG), ] @pytest.mark.parametrize("exception_to_raise, expected_res", RAISE_EXCEPTION_PARAMS) def test_command_test_raise_exception(self, mocker, exception_to_raise, expected_res): """ Given: - Args for test_command When: - Run the test_command Then: - raise exceptions and validate them """ # prepare mocker.patch.object( LiveResponseSessionManager, "request_session", side_effect=functools.partial(raise_exception, exception_to_raise) ) # run kwargs = commands_with_args[command_test_module] with pytest.raises(DemistoException) as exc_info: command_test_module(**kwargs) # validate assert expected_res in exc_info.value.args[0] def test_not_implemented_command(self, mocker): """ Given - When - Try to run not implemented command Then - Validate NotImplementedError was occurred """ # prepare not_implemented_command = "not_implemented" mocker.patch.object(demisto, "command", return_value=not_implemented_command) mocker.patch.object( demisto, "params", return_value={"custom_key": "unit test", "custom_id": "unit test", "org_key": "unit test"} ) # run with pytest.raises(NotImplementedError) as exc_info: main() assert exc_info.value.args[0] == f"Command: {not_implemented_command} not implemented"