diff --git a/monkey/common/data/zero_trust_consts.py b/monkey/common/data/zero_trust_consts.py index 6d01fbfc4..c37b261fe 100644 --- a/monkey/common/data/zero_trust_consts.py +++ b/monkey/common/data/zero_trust_consts.py @@ -99,3 +99,7 @@ TESTS_MAP = { POSSIBLE_STATUSES_KEY: [STATUS_UNEXECUTED, STATUS_CONCLUSIVE, STATUS_POSITIVE] }, } +EVENT_TYPE_ISLAND = "island" +EVENT_TYPE_MONKEY_NETWORK = "monkey_network" +EVENT_TYPE_MONKEY_LOCAL = "monkey_local" +EVENT_TYPES = (EVENT_TYPE_MONKEY_LOCAL, EVENT_TYPE_MONKEY_NETWORK, EVENT_TYPE_ISLAND) \ No newline at end of file diff --git a/monkey/monkey_island/cc/models/event.py b/monkey/monkey_island/cc/models/event.py index dd277c7c9..0b8386ff8 100644 --- a/monkey/monkey_island/cc/models/event.py +++ b/monkey/monkey_island/cc/models/event.py @@ -1,6 +1,8 @@ +from datetime import datetime + from mongoengine import EmbeddedDocument, DateTimeField, StringField -EVENT_TYPES = ("monkey_local_action", "monkey_network_action", "island_action") +from common.data.zero_trust_consts import EVENT_TYPES class Event(EmbeddedDocument): @@ -8,3 +10,16 @@ class Event(EmbeddedDocument): title = StringField(required=True) message = StringField() event_type = StringField(required=True, choices=EVENT_TYPES) + + @staticmethod + def create_event(title, message, event_type): + event = Event( + timestamp=datetime.now(), + title=title, + message=message, + event_type=event_type + ) + + event.validate(clean=True) + + return event diff --git a/monkey/monkey_island/cc/models/finding.py b/monkey/monkey_island/cc/models/finding.py index c01e6b955..d67b10247 100644 --- a/monkey/monkey_island/cc/models/finding.py +++ b/monkey/monkey_island/cc/models/finding.py @@ -2,14 +2,13 @@ Define a Document Schema for Zero Trust findings. """ -from mongoengine import Document, StringField, ListField, EmbeddedDocumentField +from mongoengine import Document, StringField, EmbeddedDocumentListField +from common.data.zero_trust_consts import TEST_STATUSES, TESTS, TESTS_MAP, EXPLANATION_KEY, PILLARS_KEY # Dummy import for mongoengine. # noinspection PyUnresolvedReferences from event import Event -from common.data.zero_trust_consts import TEST_STATUSES, PILLARS, TESTS, TESTS_MAP, EXPLANATION_KEY, PILLARS_KEY - class Finding(Document): """ @@ -21,7 +20,7 @@ class Finding(Document): # SCHEMA test = StringField(required=True, choices=TESTS) status = StringField(required=True, choices=TEST_STATUSES) - events = ListField(field=EmbeddedDocumentField('Event')) + events = EmbeddedDocumentListField(document_type=Event) # LOGIC def get_test_explanation(self): diff --git a/monkey/monkey_island/cc/models/test_event.py b/monkey/monkey_island/cc/models/test_event.py new file mode 100644 index 000000000..f36f4f562 --- /dev/null +++ b/monkey/monkey_island/cc/models/test_event.py @@ -0,0 +1,31 @@ +from mongoengine import ValidationError + +from common.data.zero_trust_consts import EVENT_TYPE_ISLAND +from monkey_island.cc.models.event import Event +from monkey_island.cc.testing.IslandTestCase import IslandTestCase + + +class TestEvent(IslandTestCase): + def test_create_event(self): + self.fail_if_not_testing_env() + self.clean_monkey_db() + + with self.assertRaises(ValidationError): + Event.create_event( + title=None, # title required + message="bla bla", + event_type=EVENT_TYPE_ISLAND + ) + + with self.assertRaises(ValidationError): + Event.create_event( + title="skjs", + message="bla bla", + event_type="Unknown" # Unknown event type + ) + + _ = Event.create_event( + title="skjs", + message="bla bla", + event_type=EVENT_TYPE_ISLAND # Unknown event type + ) diff --git a/monkey/monkey_island/cc/models/test_finding.py b/monkey/monkey_island/cc/models/test_finding.py index 69a3986c1..2b52553da 100644 --- a/monkey/monkey_island/cc/models/test_finding.py +++ b/monkey/monkey_island/cc/models/test_finding.py @@ -2,7 +2,7 @@ from datetime import datetime from mongoengine import ValidationError -from common.data.zero_trust_consts import TEST_SEGMENTATION, STATUS_CONCLUSIVE, NETWORKS +from common.data.zero_trust_consts import TEST_SEGMENTATION, STATUS_CONCLUSIVE, NETWORKS, EVENT_TYPE_MONKEY_NETWORK from finding import Finding, UnknownTest from monkey_island.cc.models.event import Event @@ -33,7 +33,8 @@ class TestFinding(IslandTestCase): self.assertEquals(len(Finding.objects(test=TEST_SEGMENTATION)), 0) - event_example = Event(timestamp=datetime.now(), title="Event Title", message="event message", event_type="monkey_network_action") + event_example = Event.create_event( + title="Event Title", message="event message", event_type=EVENT_TYPE_MONKEY_NETWORK) Finding.save_finding(test=TEST_SEGMENTATION, status=STATUS_CONCLUSIVE, events=[event_example]) self.assertEquals(len(Finding.objects(test=TEST_SEGMENTATION)), 1)