forked from p15670423/monkey
Improved mokey event fetching and added unit tests
This commit is contained in:
parent
a836ab7e1d
commit
ba9e8c22b4
|
@ -4,9 +4,10 @@ from bson import ObjectId
|
|||
|
||||
from monkey_island.cc.models.zero_trust.monkey_finding_details import MonkeyFindingDetails
|
||||
|
||||
|
||||
# How many events of a single finding to return to UI.
|
||||
# 50 will return 50 latest and 50 oldest events from a finding
|
||||
EVENT_FETCH_CNT = 50
|
||||
# 100 will return 50 latest and 50 oldest events from a finding
|
||||
MAX_EVENT_FETCH_CNT = 100
|
||||
|
||||
|
||||
class MonkeyZTDetailsService:
|
||||
|
@ -14,25 +15,29 @@ class MonkeyZTDetailsService:
|
|||
@staticmethod
|
||||
def fetch_details_for_display(finding_id: ObjectId) -> dict:
|
||||
pipeline = [{'$match': {'_id': finding_id}},
|
||||
{'$addFields': {'oldest_events': {'$slice': ['$events', EVENT_FETCH_CNT]},
|
||||
'latest_events': {'$slice': ['$events', -1 * EVENT_FETCH_CNT]},
|
||||
{'$addFields': {'oldest_events': {'$slice': ['$events', int(MAX_EVENT_FETCH_CNT / 2)]},
|
||||
'latest_events': {'$slice': ['$events', int(-1 * MAX_EVENT_FETCH_CNT / 2)]},
|
||||
'event_count': {'$size': '$events'}}},
|
||||
{'$unset': ['events', 'checked_subnet_pairs']}]
|
||||
{'$unset': ['events']}]
|
||||
details = list(MonkeyFindingDetails.objects.aggregate(*pipeline))
|
||||
if details:
|
||||
details = details[0]
|
||||
details['latest_events'] = MonkeyZTDetailsService._get_events_without_overlap(details['event_count'],
|
||||
details['latest_events'])
|
||||
details['latest_events'] = MonkeyZTDetailsService._remove_redundant_events(details['event_count'],
|
||||
details['latest_events'])
|
||||
return details
|
||||
else:
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def _get_events_without_overlap(event_count: int, events: List[object]) -> List[object]:
|
||||
overlap_count = event_count - EVENT_FETCH_CNT
|
||||
if overlap_count >= EVENT_FETCH_CNT:
|
||||
return events
|
||||
def _remove_redundant_events(fetched_event_count: int, latest_events: List[object]) -> List[object]:
|
||||
overlap_count = fetched_event_count - int(MAX_EVENT_FETCH_CNT/2)
|
||||
# None of 'latest_events' are in 'oldest_events'
|
||||
if overlap_count >= MAX_EVENT_FETCH_CNT:
|
||||
return latest_events
|
||||
# All 'latest_events' are already in 'oldest_events'
|
||||
elif overlap_count <= 0:
|
||||
return []
|
||||
# Some of 'latest_events' are already in 'oldest_events'.
|
||||
# Return only those that are not
|
||||
else:
|
||||
return events[-1 * overlap_count:]
|
||||
return latest_events[-1 * overlap_count:]
|
||||
|
|
|
@ -0,0 +1,31 @@
|
|||
from monkey_island.cc.services.zero_trust.monkey_findings import monkey_zt_details_service
|
||||
from monkey_island.cc.services.zero_trust.monkey_findings.monkey_zt_details_service import MonkeyZTDetailsService
|
||||
|
||||
|
||||
def test__remove_redundant_events(monkeypatch):
|
||||
monkeypatch.setattr(monkey_zt_details_service, 'MAX_EVENT_FETCH_CNT', 6)
|
||||
|
||||
# No events are redundant, 8 events in the database, but we display only 6 (3 latest and 3 oldest)
|
||||
latest_events = ['6', '7', '8']
|
||||
_do_redundant_event_removal_test(latest_events, 8, ['6', '7', '8'])
|
||||
|
||||
# All latest events are redundant (only 3 events in db and we fetched them twice)
|
||||
latest_events = ['1', '2', '3']
|
||||
_do_redundant_event_removal_test(latest_events, 3, [])
|
||||
|
||||
# Some latest events are redundant (5 events in db and we fetched 3 oldest and 3 latest)
|
||||
latest_events = ['3', '4', '5']
|
||||
_do_redundant_event_removal_test(latest_events, 5, ['4', '5'])
|
||||
|
||||
# None of the events are redundant (6 events in db and we fetched 3 oldest and 3 latest)
|
||||
latest_events = ['4', '5', '6']
|
||||
_do_redundant_event_removal_test(latest_events, 6, ['4', '5', '6'])
|
||||
|
||||
# No events fetched, should return empty array also
|
||||
latest_events = []
|
||||
_do_redundant_event_removal_test(latest_events, 0, [])
|
||||
|
||||
|
||||
def _do_redundant_event_removal_test(input_list, fetched_event_cnt, expected_output):
|
||||
output_events = MonkeyZTDetailsService._remove_redundant_events(fetched_event_cnt, input_list)
|
||||
assert output_events == expected_output
|
Loading…
Reference in New Issue