def test_get_built_in_query_results_polling_command(mocker):
"""
Given:
- A user arguments.
When:
- Calling get_built_in_query_results_polling_command function.
Then:
- Ensure start_xql_query_polling_command function was called with the right query and argument.
"""
args = {
'endpoint_id': '123456,654321',
'file_sha256': 'abcde,edcba,p1p2p3',
'extra_fields': 'EXTRA1, EXTRA2',
'limit': '400',
'tenants': "tenantID,tenantID",
'time_frame': '7 days'
}
res = mocker.patch('XQLQueryingEngine.start_xql_query_polling_command')
mocker.patch.object(demisto, 'command', return_value='xdr-xql-file-event-query')
XQLQueryingEngine.get_built_in_query_results_polling_command(CLIENT, args)
assert res.call_args.args[1]['query'] == '''dataset = xdr_data | filter agent_id in ("123456","654321") and event_type = FILE and action_file_sha256
in ("abcde","edcba","p1p2p3")| fields agent_hostname, agent_ip_addresses, agent_id, action_file_path, action_file_sha256,
actor_process_file_create_time, EXTRA1, EXTRA2 | limit 400'''
assert res.call_args.args[1]['tenants'] == ["tenantID", "tenantID"]
assert res.call_args.args[1]['time_frame'] == '7 days'
def test_get_xql_query_results_polling_command_pending(mocker):
"""
Given:
- A query that has a pending status.
When:
- Calling get_xql_query_results_polling_command function.
Then:
- Ensure returned command results are correct and the scheduled_command is set properly.
"""
query = 'MOCK_QUERY'
mock_response = {
'status': 'PENDING',
'execution_id': 'query_id_mock',
'results': None
}
mocker.patch('XQLQueryingEngine.get_xql_query_results',
return_value=(mock_response, None))
mocker.patch('XQLQueryingEngine.is_demisto_version_ge', return_value=True)
mocker.patch.object(demisto,
'command',
return_value='xdr-xql-generic-query')
mocker.patch('XQLQueryingEngine.ScheduledCommand', return_value=None)
command_results = XQLQueryingEngine.get_xql_query_results_polling_command(
CLIENT, {'query': query})
assert command_results.readable_output == 'Query is still running, it may take a little while...'
assert command_results.outputs == {
'status': 'PENDING',
'execution_id': 'query_id_mock',
'results': None,
'query_name': ''
}
def test_get_xql_query_results_success_under_1000(mocker):
"""
Given:
- a query ID which has 1 result.
When:
- Calling get_xql_query_results function.
Then:
- Ensure the results were retrieved properly.
"""
args = {
'query_id': 'query_id_mock',
'time_frame': '1 year'
}
mock_response = {
'status': 'SUCCESS',
'number_of_results': 1,
'query_cost': {
"376699223": 0.0031591666666666665
},
'remaining_quota': 1000.0,
'results': {
'data': [{'x': 'test1'}]
}
}
mocker.patch.object(CLIENT, 'get_xql_query_results', return_value=mock_response)
response, file_data = XQLQueryingEngine.get_xql_query_results(CLIENT, args=args)
assert response == {'status': 'SUCCESS',
'number_of_results': 1,
'query_cost': {'376699223': 0.0031591666666666665},
'remaining_quota': 1000.0,
'results': [{'x': 'test1'}],
'execution_id': 'query_id_mock'}
assert file_data is None
def test_start_xql_query_polling_not_supported(mocker):
"""
Given:
- A query that has a pending status.
When:
- Calling get_xql_query_results_polling_command function but polling is not supported.
Then:
- Ensure returned command results are correct.
"""
query = 'MOCK_QUERY'
mock_response = {'status': 'PENDING',
'execution_id': 'query_id_mock',
'results': None}
mocker.patch.object(CLIENT, 'start_xql_query', return_value='1234')
mocker.patch('XQLQueryingEngine.get_xql_query_results', return_value=(mock_response, None))
mocker.patch('XQLQueryingEngine.is_demisto_version_ge', return_value=False)
mocker.patch.object(demisto, 'command', return_value='xdr-xql-generic-query')
command_results = XQLQueryingEngine.start_xql_query_polling_command(CLIENT, {'query': query, 'query_name': 'mock_name'})
assert command_results.outputs == {'status': 'PENDING',
'execution_id': 'query_id_mock',
'results': None,
'query_name': 'mock_name'}
def test_get_xql_query_results_clear_integration_context_on_success(mocker):
"""
Given:
- A query that has a successful status and the number of results is under 1000.
When:
- Calling get_xql_query_results_polling_command function.
Then:
- Ensure the integration context was cleared.
"""
query = 'MOCK_QUERY'
mock_response = {'status': 'SUCCESS',
'number_of_results': 1,
'query_cost': {'376699223': 0.0031591666666666665},
'remaining_quota': 1000.0,
'results': [{'x': 'test1', 'y': None}],
'execution_id': 'query_id_mock'}
mocker.patch('XQLQueryingEngine.get_xql_query_results', return_value=(mock_response, None))
mocker.patch.object(demisto, 'command', return_value='xdr-xql-generic-query')
command_results = XQLQueryingEngine.get_xql_query_results_polling_command(CLIENT, {'query': query})
assert command_results.outputs == {'status': 'SUCCESS', 'number_of_results': 1, 'query_name': '',
'query_cost': {'376699223': 0.0031591666666666665}, 'remaining_quota': 1000.0,
'execution_id': 'query_id_mock', 'results': [{'x': 'test1'}]}
assert '| query_id_mock | 1 | MOCK_QUERY | 376699223: 0.0031591666666666665 | 1000.0 | SUCCESS |' in \
command_results.readable_output
assert 'y' in command_results.raw_response['results'][0]
def test_start_xql_query_invalid(mocker):
"""
Given:
- An invalid query to search.
When:
- Calling start_xql_query function.
Then:
- Ensure an error is returned.
"""
args = {'query': 'test_query // Some Note', 'time_frame': '1 year'}
mocker.patch.object(CLIENT, 'start_xql_query', return_value='execution_id')
with pytest.raises(Exception) as exc:
XQLQueryingEngine.start_xql_query(CLIENT, args=args)
assert str(exc.value) == 'Please remove notes (//) from query'
def test_get_xql_query_results_polling_command_success_more_than_1000(mocker):
"""
Given:
- A query that has a successful status and the number of results is more than 1000.
When:
- Calling get_xql_query_results_polling_command function.
Then:
- Ensure returned command results are correct.
"""
query = 'MOCK_QUERY'
mock_response = {'status': 'SUCCESS',
'number_of_results': 1500,
'query_cost': {'376699223': 0.0031591666666666665},
'remaining_quota': 1000.0,
'results': {'stream_id': 'test_stream_id'},
'execution_id': 'query_id_mock'}
mocker.patch('XQLQueryingEngine.get_xql_query_results', return_value=(mock_response, 'File Data'))
mocker.patch.object(demisto, 'command', return_value='xdr-xql-generic-query')
mocker.patch('XQLQueryingEngine.fileResult',
return_value={'Contents': '', 'ContentsFormat': 'text', 'Type': 3, 'File': 'results.gz',
'FileID': '12345'})
results = XQLQueryingEngine.get_xql_query_results_polling_command(CLIENT, {'query': query})
assert results[0] == {'Contents': '', 'ContentsFormat': 'text', 'Type': 3, 'File': 'results.gz', 'FileID': '12345'}
command_result = results[1]
assert command_result.outputs == {'status': 'SUCCESS', 'number_of_results': 1500, 'query_name': '',
'query_cost': {'376699223': 0.0031591666666666665}, 'remaining_quota': 1000.0,
'results': {'stream_id': 'test_stream_id'}, 'execution_id': 'query_id_mock'}
def test_get_xql_quota_command(mocker):
"""
Given:
- A client object.
When:
- Calling get_xql_quota_command function.
Then:
- Ensure returned command results are correct.
"""
mock_response = {
"reply": {
"license_quota": 1000,
"additional_purchased_quota": 0,
"used_quota": 0.0
}
}
mocker.patch.object(CLIENT, 'get_xql_quota', return_value=mock_response)
response = XQLQueryingEngine.get_xql_quota_command(CLIENT, {})
assert '|Additional Purchased Quota|License Quota|Used Quota|' in response.readable_output
assert response.outputs == {
'license_quota': 1000,
'additional_purchased_quota': 0,
'used_quota': 0.0
}
def test_format_results_do_not_remove_empty_fields():
"""
Given:
- A list to format with remove_empty_fields flag turned off.
When:
- Calling format_results function.
Then:
- Ensure the list was formatted properly.
"""
list_to_format = [{
'h': 4
}, {
'x': 1,
'e': None,
'y': 'FALSE',
'z': {
'w': 'NULL',
'x': None,
},
's': {
'a': 5,
'b': None,
'c': {
'time': 1629619736000,
'd': 3,
'v': 'TRUE'
}
}
}]
expected = [{
'h': 4
}, {
'x': 1,
'e': None,
'y': False,
'z': {
'w': None,
'x': None,
},
's': {
'a': 5,
'b': None,
'c': {
'time': '2021-08-22T08:08:56.000Z',
'd': 3,
'v': True
}
}
}]
response = XQLQueryingEngine.format_results(list_to_format,
remove_empty_fields=False)
assert expected == response
def test_wrap_list_items_in_double_quotes(input_arg, expected):
"""
Given:
- A string list to format.
When:
- Calling format_arg function.
Then:
- Ensure the returned string is correct.
"""
response = XQLQueryingEngine.wrap_list_items_in_double_quotes(input_arg)
assert response == expected
def test_convert_timeframe_string_to_json(time_to_convert, expected):
"""
Given:
- A relative time or time range to convert.
When:
- Calling convert_timeframe_string_to_json function.
Then:
- Ensure the returned timestamp is correct.
"""
response = XQLQueryingEngine.convert_timeframe_string_to_json(time_to_convert=time_to_convert)
assert response == expected
def test_get_query_result_stream(mocker):
"""
Given:
- a stream_id.
When:
- Calling get_query_result_stream function.
Then:
- Ensure the results were retrieved properly.
"""
stream_id = 'mock_stream_id'
mocker.patch.object(CLIENT, 'get_query_result_stream', return_value='Raw Data')
response = XQLQueryingEngine.get_query_result_stream(CLIENT, stream_id=stream_id)
assert response == 'Raw Data'
def test_start_xql_query_valid(mocker):
"""
Given:
- A valid query to search.
When:
- Calling start_xql_query function.
Then:
- Ensure the returned execution_id is correct.
"""
args = {'query': 'test_query', 'time_frame': '1 year'}
mocker.patch.object(CLIENT, 'start_xql_query', return_value='execution_id')
response = XQLQueryingEngine.start_xql_query(CLIENT, args=args)
assert response == 'execution_id'
def test_get_xql_query_results_polling_command_success_more_than_1000_results_parse_to_context(mocker):
"""
Given:
- A query that has a successful status and the number of results is more than 1000.
When:
- Calling get_xql_query_results_polling_command function with 'parse_result_file_to_context' argument set to True.
Then:
- Ensure returned command results are correct.
- Ensure the results were parsed to context instead of being extracted to a file.
"""
query = 'MOCK_QUERY'
mock_response = {'status': 'SUCCESS',
'number_of_results': 1500,
'query_cost': {'376699223': 0.0031591666666666665},
'remaining_quota': 1000.0,
'results': {'stream_id': 'test_stream_id'},
'execution_id': 'query_id_mock'}
# The results that should be parsed to context instead of being extracted to a file:
expected_results_in_context = [
{"_time": "2021-10-14 03:59:09.793 UTC", "event_id": "123", "_vendor": "PANW", "_product": "XDR agent",
"insert_timestamp": "2021-10-14 04:02:12.883114 UTC"},
{"_time": "2021-10-14 03:59:09.809 UTC", "event_id": "234", "_vendor": "PANW", "_product": "XDR agent",
"insert_timestamp": "2021-10-14 04:02:12.883114 UTC"},
{"_time": "2021-10-14 04:00:27.78 UTC", "event_id": "456", "_vendor": "PANW", "_product": "XDR agent",
"insert_timestamp": "2021-10-14 04:04:34.332563 UTC"},
{"_time": "2021-10-14 04:00:27.797 UTC", "event_id": "567", "_vendor": "PANW", "_product": "XDR agent",
"insert_timestamp": "2021-10-14 04:04:34.332563 UTC"}
]
# Creates the mocked data which returns from 'XQLQueryingEngine.get_xql_query_results' command:
mock_file_data = b''
for item in expected_results_in_context:
mock_file_data += json.dumps(item).encode('utf-8')
mock_file_data += b'\n'
compressed_mock_file_data = gzip.compress(mock_file_data)
mocker.patch('XQLQueryingEngine.get_xql_query_results', return_value=(mock_response, compressed_mock_file_data))
mocker.patch.object(demisto, 'command', return_value='xdr-xql-generic-query')
results = XQLQueryingEngine.get_xql_query_results_polling_command(CLIENT, {'query': query,
'parse_result_file_to_context': True})
assert results.outputs.get('results', []) == expected_results_in_context, \
'There might be a problem in parsing the results into the context'
assert results.outputs == {'status': 'SUCCESS', 'number_of_results': 1500, 'query_name': '',
'query_cost': {'376699223': 0.0031591666666666665}, 'remaining_quota': 1000.0,
'results': expected_results_in_context, 'execution_id': 'query_id_mock'}
def test_convert_relative_time_to_milliseconds(time_to_convert, expected):
"""
Given:
- A relative time to convert.
When:
- Calling convert_relative_time_to_milliseconds function.
Then:
- Ensure the returned timestamp is correct.
"""
response = XQLQueryingEngine.convert_relative_time_to_milliseconds(
time_to_convert=time_to_convert)
assert response == expected
def test_get_registry_query():
"""
Given:
- ENDPOINT_IDS and reg_key_name list (as a string).
When:
- Calling get_registry_query function.
Then:
- Ensure the returned query is correct.
"""
args = {'reg_key_name': 'testARG1,testARG2'}
response = XQLQueryingEngine.get_registry_query(endpoint_ids=ENDPOINT_IDS,
args=args)
assert response == '''dataset = xdr_data | filter agent_id in ("test1","test2") and event_type = REGISTRY and
def test_get_dll_module_query():
"""
Given:
- ENDPOINT_IDS and loaded_module_sha256 list (as a string).
When:
- Calling get_dll_module_query function.
Then:
- Ensure the returned query is correct.
"""
args = {'loaded_module_sha256': 'testSHA1,testSHA2'}
response = XQLQueryingEngine.get_dll_module_query(
endpoint_ids=ENDPOINT_IDS, args=args)
assert response == '''dataset = xdr_data | filter agent_id in ("test1","test2") and event_type = LOAD_IMAGE and
def test_get_event_log_query():
"""
Given:
- ENDPOINT_IDS and get_event_log list (as a string).
When:
- Calling get_event_log_query function.
Then:
- Ensure the returned query is correct.
"""
args = {'event_id': '1234,4321'}
response = XQLQueryingEngine.get_event_log_query(endpoint_ids=ENDPOINT_IDS,
args=args)
assert response == '''dataset = xdr_data | filter agent_id in ("test1","test2") and event_type = EVENT_LOG and
def test_get_dns_query_no_external_domain_arg():
"""
Given:
- ENDPOINT_IDS and dns_query list (as a string).
When:
- Calling get_dns_query function.
Then:
- Ensure the returned query is correct.
"""
args = {
'dns_query': 'testARG3,testARG4',
}
response = XQLQueryingEngine.get_dns_query(endpoint_ids=ENDPOINT_IDS, args=args)
assert response == '''dataset = xdr_data | filter (agent_id in ("test1","test2") and event_type = STORY) and
def test_get_network_connection_query_only_remote_ip():
"""
Given:
- ENDPOINT_IDS and remote_ip_list (as a string).
When:
- Calling get_network_connection_query function.
Then:
- Ensure the returned query is correct.
"""
args = {
'remote_ip': '3.3.3.3,4.4.4.4',
}
response = XQLQueryingEngine.get_network_connection_query(endpoint_ids=ENDPOINT_IDS, args=args)
assert response == '''dataset = xdr_data | filter agent_id in ("test1","test2") and event_type = STORY
def test_get_file_dropper_query_no_file_path_arg():
"""
Given:
- ENDPOINT_IDS and file_sha256 list (as a string).
When:
- Calling get_file_dropper_query function.
Then:
- Ensure the returned query is correct.
"""
args = {
'file_sha256': 'testARG3,testARG4',
}
response = XQLQueryingEngine.get_file_dropper_query(endpoint_ids=ENDPOINT_IDS, args=args)
assert response == '''dataset = xdr_data | filter (agent_id in ("test1","test2") and event_type = FILE and event_sub_type in (
def test_get_process_causality_network_activity_query():
"""
Given:
- ENDPOINT_IDS and process_causality_id list (as a string).
When:
- Calling get_process_causality_network_activity_query function.
Then:
- Ensure the returned query is correct.
"""
args = {
'process_causality_id': 'testARG1,testARG2',
}
response = XQLQueryingEngine.get_process_causality_network_activity_query(endpoint_ids=ENDPOINT_IDS, args=args)
assert response == '''dataset = xdr_data | filter agent_id in ("test1","test2") and event_type = NETWORK
def test_get_xql_query_results_success_more_than_1000(mocker):
"""
Given:
- a query ID which has more than 1000 results.
When:
- Calling get_xql_query_results function.
Then:
- Ensure the results were retrieved properly and a stream ID was returned.
"""
args = {'query_id': 'query_id_mock', 'time_frame': '1 year'}
mock_response = {
'status': 'SUCCESS',
'number_of_results': 1500,
'query_cost': {
"376699223": 0.0031591666666666665
},
'remaining_quota': 1000.0,
'results': {
"stream_id": "test_stream_id"
}
}
mocker.patch.object(CLIENT,
'get_xql_query_results',
return_value=mock_response)
mocker.patch.object(CLIENT,
'get_query_result_stream',
return_value='FILE DATA')
response, file_data = XQLQueryingEngine.get_xql_query_results(CLIENT,
args=args)
assert response == {
'status': 'SUCCESS',
'number_of_results': 1500,
'query_cost': {
'376699223': 0.0031591666666666665
},
'remaining_quota': 1000.0,
'results': {
'stream_id': 'test_stream_id'
},
'execution_id': 'query_id_mock'
}
assert file_data == 'FILE DATA'
def test_start_xql_query_polling_command(mocker):
"""
Given:
- A query that has a successful status and the number of results is under 1000.
When:
- Calling get_xql_query_results_polling_command function.
Then:
- Ensure returned command results are correct and integration_context was cleared.
"""
query = 'MOCK_QUERY'
context = {
'mock_id': {
'query': 'mock_query',
'time_frame': '3 days',
'command_name': 'previous command',
'query_name': 'mock_name',
}
}
set_integration_context(context)
mock_response = {'status': 'SUCCESS',
'number_of_results': 1,
'query_cost': {'376699223': 0.0031591666666666665},
'remaining_quota': 1000.0,
'results': [{'x': 'test1', 'y': None}],
'execution_id': 'query_id_mock'}
mocker.patch.object(CLIENT, 'start_xql_query', return_value='1234')
mocker.patch('XQLQueryingEngine.get_xql_query_results', return_value=(mock_response, None))
mocker.patch.object(demisto, 'command', return_value='xdr-xql-generic-query')
mocker.patch.object(demisto, 'getIntegrationContext', side_effect=get_integration_context)
mocker.patch.object(demisto, 'setIntegrationContext', side_effect=set_integration_context)
command_results = XQLQueryingEngine.start_xql_query_polling_command(CLIENT, {'query': query, 'query_name': 'mock_name'})
assert command_results.outputs == {'status': 'SUCCESS', 'number_of_results': 1, 'query_name': 'mock_name',
'query_cost': {'376699223': 0.0031591666666666665}, 'remaining_quota': 1000.0,
'execution_id': 'query_id_mock', 'results': [{'x': 'test1'}]}
assert '| query_id_mock | 1 | MOCK_QUERY | 376699223: 0.0031591666666666665 | mock_name | 1000.0 | SUCCESS |' in \
command_results.readable_output
assert 'y' in command_results.raw_response['results'][0]
assert get_integration_context() == context
def test_get_xql_query_results_pending(mocker):
"""
Given:
- a query ID which will cause a pending status.
When:
- Calling get_xql_query_results function.
Then:
- Ensure the results were retrieved properly.
"""
args = {'query_id': 'query_id_mock', 'time_frame': '1 year'}
mock_response = {"status": "PENDING"}
mocker.patch.object(CLIENT,
'get_xql_query_results',
return_value=mock_response)
response, _ = XQLQueryingEngine.get_xql_query_results(CLIENT, args=args)
assert response == {
'status': 'PENDING',
'execution_id': 'query_id_mock',
'results': None
}
import gzip
import json
import io
from freezegun import freeze_time
import XQLQueryingEngine
import pytest
from CommonServerPython import *
CLIENT = XQLQueryingEngine.Client(base_url='some_mock_url', verify=False)
ENDPOINT_IDS = '"test1","test2"'
INTEGRATION_CONTEXT = {}
def util_load_json(path):
with io.open(path, mode='r', encoding='utf-8') as f:
return json.loads(f.read())
def get_integration_context():
return INTEGRATION_CONTEXT
def set_integration_context(integration_context):
global INTEGRATION_CONTEXT
INTEGRATION_CONTEXT = integration_context
# =========================================== TEST Built-In Queries helpers ===========================================#
@pytest.mark.parametrize(