p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add test for `process_post_breach_telemetry()`

This commit is contained in:
Shreya 2020-08-01 14:48:09 +05:30
parent 10b15b3efe
commit 435637ad3a
2 changed files with 93 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
monkey/monkey_island/cc/services/telemetry/processing/post_breach.py
View File

@ -36,12 +36,6 @@ def process_post_breach_telemetry(telemetry_json):
modified_data.append(temp)
return modified_data
def update_data(data):
data = add_message_for_blank_outputs(data)
mongo.db.monkey.update(
{'guid': telemetry_json['monkey_guid']},
{'$push': {'pba_results': data}})
def add_message_for_blank_outputs(data):
if not data['result'][0]:
data['result'][0] = EXECUTION_WITHOUT_OUTPUT
@ -54,4 +48,11 @@ def process_post_breach_telemetry(telemetry_json):
telemetry_json['data'] = convert_telem_data_to_list(telemetry_json['data'])
for pba_data in telemetry_json['data']:
update_data(pba_data)
pba_data = add_message_for_blank_outputs(pba_data)
update_data(telemetry_json, pba_data)
def update_data(telemetry_json, data):
mongo.db.monkey.update(
{'guid': telemetry_json['monkey_guid']},
{'$push': {'pba_results': data}})

85
monkey/monkey_island/cc/services/telemetry/processing/test_post_breach.py Normal file
View File

@ -0,0 +1,85 @@
from unittest.mock import Mock
import monkey_island.cc.services.telemetry.processing.post_breach as post_breach
from .post_breach import EXECUTION_WITHOUT_OUTPUT
original_telem_multiple_results =\
{
'data': {
'command': 'COMMAND',
'hostname': 'HOST',
'ip': '127.0.1.1',
'name': 'PBA NAME',
'result': [
['SUCCESSFUL', True],
['UNSUCCESFUL', False],
['', True]
]
},
'telem_category': 'post_breach'
}
expected_telem_multiple_results =\
{
'data': [
{
'command': 'COMMAND',
'hostname': 'HOST',
'ip': '127.0.1.1',
'name': 'PBA NAME',
'result': ['SUCCESSFUL', True]
},
{
'command': 'COMMAND',
'hostname': 'HOST',
'ip': '127.0.1.1',
'name': 'PBA NAME',
'result': ['UNSUCCESFUL', False]
},
{
'command': 'COMMAND',
'hostname': 'HOST',
'ip': '127.0.1.1',
'name': 'PBA NAME',
'result': [EXECUTION_WITHOUT_OUTPUT, True]
}
],
'telem_category': 'post_breach'
}
original_telem_single_result =\
{
'data': {
'command': 'COMMAND',
'hostname': 'HOST',
'ip': '127.0.1.1',
'name': 'PBA NAME',
'result': ['', True]
},
'telem_category': 'post_breach'
}
expected_telem_single_result =\
{
'data': [
{
'command': 'COMMAND',
'hostname': 'HOST',
'ip': '127.0.1.1',
'name': 'PBA NAME',
'result': [EXECUTION_WITHOUT_OUTPUT, True]
},
],
'telem_category': 'post_breach'
}
def test_process_post_breach_telemetry():
post_breach.update_data = Mock() # actual behavior of update_data() is to access mongodb
# multiple results in PBA
post_breach.process_post_breach_telemetry(original_telem_multiple_results)
assert original_telem_multiple_results == expected_telem_multiple_results
# single result in PBA
post_breach.process_post_breach_telemetry(original_telem_single_result)
assert original_telem_single_result == expected_telem_single_result