import json def util_load_json(path): with open(path, encoding="utf-8") as f: return json.loads(f.read()) """*****COMMAND FUNCTIONS****""" def test_get_model_breach(requests_mock): """Tests darktrace-get--model-breach command function. Configures requests_mock instance to generate the appropriate get_alerts API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, get_model_breach_command # GIVEN an integration is configured to Darktrace mock_api_response = util_load_json("test_data/get_breach.json") requests_mock.get("https://mock.darktrace.com/modelbreaches?pbid=95", json=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the desired model breach has id 95 args = { "pbid": "95", } integration_response = get_model_breach_command(client, args) expected_response = util_load_json("test_data/formatted_get_breach.json") # THEN the response should be returned and formatted assert integration_response.outputs == expected_response assert integration_response.outputs_prefix == "Darktrace.ModelBreach" def test_fetch_incidents(requests_mock): """Tests the fetch-incidents command function. Configures requests_mock instance to generate the appropriate get_alert API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, fetch_incidents # GIVEN an integration is configured and fetch incidents mock_response = util_load_json("test_data/fetch_breach.json") requests_mock.get( "https://usw1-51965-01.cloud.darktrace.com/modelbreaches?minscore=0.0&starttime=1598932817000", json=mock_response ) client = Client(base_url="https://usw1-51965-01.cloud.darktrace.com/", verify=False, auth=("examplepub", "examplepri")) # WHEN the most recent call was made on Mon, Aug 31, 2020 9 PM Pacific last_run = { "last_fetch": 1598932817000 # Mon, Aug 31, 2020 9 PM Pacific } _, integration_response = fetch_incidents( client=client, max_alerts=20, last_run=last_run, first_fetch_time="1 day ago", min_score=0 ) # THEN the relevant information will be fetched and pulled expected_response = util_load_json("test_data/formatted_fetch_breach.json") assert integration_response == expected_response assert len(integration_response) == 2 def test_get_model_breach_connections(mocker): """Tests the get-modelbreach-connections command function. Configures requests_mock instance to generate the appropriate get_alert API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, get_model_breach_connections_command # GIVEN an integration is configured and you would like to find similar devices mock_api_response = util_load_json("test_data/breach_details.json") mocker.patch.object(Client, "get_model_breach_connections", return_value=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the specified device id is 1 and there are 2 results max desired args = {"pbid": "123", "count": "2", "endtime": 1629803362, "offset": 0} # THEN the context will be updated and information about similar devices will be fetched and pulled integration_response = get_model_breach_connections_command(client, args) expected_response = util_load_json("test_data/formatted_breach_details.json") assert integration_response.outputs == expected_response assert len(mock_api_response) == len(expected_response) + 1 assert integration_response.outputs_prefix == "Darktrace.ModelBreach" def test_get_model(requests_mock): """Tests the get-model command function. Configures requests_mock instance to generate the appropriate get_alert API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, get_model_command # GIVEN an integration is configured and you would like to find similar devices mock_api_response = util_load_json("test_data/model.json") requests_mock.get("https://mock.darktrace.com/models?uuid=80010119-6d7f-0000-0305-5e0000000325", json=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the specified device id is 1 and there are 2 results max desired args = {"uuid": "80010119-6d7f-0000-0305-5e0000000325"} # THEN the context will be updated and information about similar devices will be fetched and pulled integration_response = get_model_command(client, args) expected_response = util_load_json("test_data/formatted_model.json") assert integration_response.outputs == expected_response assert integration_response.outputs_prefix == "Darktrace.Model" def test_get_model_component(requests_mock): """Tests the get-component command function. Configures requests_mock instance to generate the appropriate get_alert API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, get_model_component_command # GIVEN an integration is configured and you would like to find similar devices mock_api_response = util_load_json("test_data/component.json") requests_mock.get("https://mock.darktrace.com/components?cid=254503", json=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the specified device id is 1 and there are 2 results max desired args = {"cid": "254503"} # THEN the context will be updated and information about similar devices will be fetched and pulled integration_response = get_model_component_command(client, args) expected_response = util_load_json("test_data/formatted_component.json") assert integration_response.outputs == expected_response assert integration_response.outputs_prefix == "Darktrace.Model.Component" def test_get_model_breach_comments(requests_mock): """Tests darktrace-get-model-breach-comments command function. Configures requests_mock instance to generate the appropriate get_alerts API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, get_model_breach_comments_command # GIVEN an integration is configured and comments are desired mock_api_response = util_load_json("test_data/get_comments.json") requests_mock.get("https://mock.darktrace.com/mbcomments?pbid=2507", json=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the desired model breach has id 46 args = { "pbid": "2507", } integration_response = get_model_breach_comments_command(client, args) expected_response = util_load_json("test_data/formatted_get_comments.json") # THEN the comments should be returned and formatted assert integration_response.outputs == expected_response assert integration_response.outputs_prefix == "Darktrace.ModelBreach.Comment" def test_acknowledge_model_breach(requests_mock): """Tests darktrace-acknowledge-model-breach command function. Configures requests_mock instance to generate the appropriate get_alerts API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, acknowledge_model_breach_command # GIVEN an integration is configured and you would like to acknowledge a breach mock_api_response = util_load_json("test_data/ack_success.json") requests_mock.post("https://mock.darktrace.com/modelbreaches/2509/acknowledge", json=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the desired model breach has id 111 args = { "pbid": "2509", } integration_response = acknowledge_model_breach_command(client, args) expected_response = util_load_json("test_data/formatted_ack_success.json") # THEN the breach should be acknowledged, context updated, and message posted assert integration_response.outputs == expected_response assert integration_response.outputs_prefix == "Darktrace.ModelBreach" assert integration_response.outputs_key_field == "pbid" def test_unacknowledge(requests_mock): """Tests darktrace-unacknowledge command function. Configures requests_mock instance to generate the appropriate get_alerts API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, unacknowledge_model_breach_command # GIVEN an integration is configured and you would like to unacknowledge a breach mock_api_response = util_load_json("test_data/ack_success.json") requests_mock.post("https://mock.darktrace.com/modelbreaches/2509/unacknowledge", json=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the desired model breach has id 111 args = { "pbid": "2509", } integration_response = unacknowledge_model_breach_command(client, args) expected_response = util_load_json("test_data/formatted_unack_success.json") # THEN the breach should be acknowledged, context updated, and message posted assert integration_response.outputs == expected_response assert integration_response.outputs_prefix == "Darktrace.ModelBreach" assert integration_response.outputs_key_field == "pbid" def test_post_comment_to_model_breach(requests_mock): from DarktraceMBs import Client, post_comment_to_model_breach_command # GIVEN an integration is configured and you would like to unacknowledge a breach mock_api_response = util_load_json("test_data/comment_post.json") requests_mock.post("https://mock.darktrace.com/modelbreaches/2509/comments", json=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the desired model breach has id 111 args = {"pbid": "2509", "message": "Test comment post"} integration_response = post_comment_to_model_breach_command(client, args) expected_response = util_load_json("test_data/formatted_comment_post.json") # THEN the breach should be acknowledged, context updated, and message posted assert integration_response.outputs == expected_response assert integration_response.outputs_prefix == "Darktrace.ModelBreach" assert integration_response.outputs_key_field == "pbid" def test_acknowledge_model_breach_UV(requests_mock): """Tests darktrace-acknowledge-model-breach command function. Configures requests_mock instance to generate the appropriate get_alerts API response, loaded from a local JSON file. Checks the output of the command function with the expected output. """ from DarktraceMBs import Client, acknowledge_model_breach_command # GIVEN an integration is configured and you would like to acknowledge a breach mock_api_response = util_load_json("test_data/ack_success_UV.json") requests_mock.post("https://mock.darktrace.com/modelbreaches/1000000000001/acknowledge", json=mock_api_response) client = Client(base_url="https://mock.darktrace.com", verify=False, auth=("examplepub", "examplepri")) # WHEN the desired model breach has id 111 args = { "pbid": "1000000000001", } integration_response = acknowledge_model_breach_command(client, args) expected_response = util_load_json("test_data/formatted_ack_success_UV.json") # THEN the breach should be acknowledged, context updated, and message posted assert integration_response.outputs == expected_response assert integration_response.outputs_prefix == "Darktrace.ModelBreach" assert integration_response.outputs_key_field == "pbid" def test_test_module(requests_mock): """ Tests the test module function """ from DarktraceMBs import Client, test_module # GIVEN an integration is configured and you would like to test the configuration mock_api_response = util_load_json("test_data/fetch_breach.json") requests_mock.get( "https://mock.darktrace.com/modelbreaches?minscore=0&starttime=1598932817000&minimal=false&deviceattop=true", json=mock_api_response, ) client = Client(base_url="https://mock.darktrace.com/", verify=False, auth=("examplepub", "examplepri")) first_fetch_time = 1598932817000 integration_response = test_module(client, first_fetch_time) expected_response = "ok" # THEN if there is no error, test_module should return 'ok' assert len(integration_response) == 2 assert isinstance(expected_response, str) def test_helper_stringify_data(): """ Tests the stringify_data helper function. """ from DarktraceMBs import stringify_data # GIVEN you have a params data structure in dictionary format and would like to test the stringify_data helper function data_1 = {"first": 1} data_2 = {"first": 1, "second": 2} data_3 = {} data_4 = {"test_1": "string_1", "test_2": 2} # WHEN calling the helper function integration_response_1 = stringify_data(data_1) integration_response_2 = stringify_data(data_2) integration_response_3 = stringify_data(data_3) integration_response_4 = stringify_data(data_4) # THEN the following outputs are expected assert integration_response_1 == "first=1" assert integration_response_2 == "first=1&second=2" assert integration_response_3 == "" assert integration_response_4 == "test_1=string_1&test_2=2"