diff --git a/monkey/monkey_island/cc/models/event.py b/monkey/monkey_island/cc/models/event.py new file mode 100644 index 000000000..dd277c7c9 --- /dev/null +++ b/monkey/monkey_island/cc/models/event.py @@ -0,0 +1,10 @@ +from mongoengine import EmbeddedDocument, DateTimeField, StringField + +EVENT_TYPES = ("monkey_local_action", "monkey_network_action", "island_action") + + +class Event(EmbeddedDocument): + timestamp = DateTimeField(required=True) + title = StringField(required=True) + message = StringField() + event_type = StringField(required=True, choices=EVENT_TYPES) diff --git a/monkey/monkey_island/cc/models/finding.py b/monkey/monkey_island/cc/models/finding.py new file mode 100644 index 000000000..c01e6b955 --- /dev/null +++ b/monkey/monkey_island/cc/models/finding.py @@ -0,0 +1,47 @@ +""" +Define a Document Schema for Zero Trust findings. +""" + +from mongoengine import Document, StringField, ListField, EmbeddedDocumentField + +# 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): + """ + This class has 2 main section: + * The schema section defines the DB fields in the document. This is the data of the object. + * The logic section defines complex questions we can ask about a single document which are asked multiple + times, or complex action we will perform - somewhat like an API. + """ + # SCHEMA + test = StringField(required=True, choices=TESTS) + status = StringField(required=True, choices=TEST_STATUSES) + events = ListField(field=EmbeddedDocumentField('Event')) + + # LOGIC + def get_test_explanation(self): + return TESTS_MAP[self.test][EXPLANATION_KEY] + + def get_pillars(self): + return TESTS_MAP[self.test][PILLARS_KEY] + + # Creation methods + @staticmethod + def save_finding(test, status, events): + finding = Finding( + test=test, + status=status, + events=events) + + finding.save() + + return finding + + +class UnknownTest(Exception): + pass diff --git a/monkey/monkey_island/cc/models/test_finding.py b/monkey/monkey_island/cc/models/test_finding.py new file mode 100644 index 000000000..69a3986c1 --- /dev/null +++ b/monkey/monkey_island/cc/models/test_finding.py @@ -0,0 +1,40 @@ +from datetime import datetime + +from mongoengine import ValidationError + +from common.data.zero_trust_consts import TEST_SEGMENTATION, STATUS_CONCLUSIVE, NETWORKS +from finding import Finding, UnknownTest +from monkey_island.cc.models.event import Event + +from monkey_island.cc.testing.IslandTestCase import IslandTestCase + + +class TestFinding(IslandTestCase): + """ + Make sure to set server environment to `testing` in server.json! Otherwise this will mess up your mongo instance and + won't work. + + Also, the working directory needs to be the working directory from which you usually run the island so the + server.json file is found and loaded. + """ + def test_save_finding_validation(self): + self.fail_if_not_testing_env() + self.clean_monkey_db() + + with self.assertRaises(ValidationError): + _ = Finding.save_finding(test="bla bla", status="Conclusive", events=[]) + + with self.assertRaises(ValidationError): + _ = Finding.save_finding(test=TEST_SEGMENTATION, status="bla bla", events=[]) + + def test_save_finding_sanity(self): + self.fail_if_not_testing_env() + self.clean_monkey_db() + + 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") + Finding.save_finding(test=TEST_SEGMENTATION, status=STATUS_CONCLUSIVE, events=[event_example]) + + self.assertEquals(len(Finding.objects(test=TEST_SEGMENTATION)), 1) + self.assertEquals(len(Finding.objects(status=STATUS_CONCLUSIVE)), 1)