diff --git a/monkey/monkey_island/cc/models/errors.py b/monkey/monkey_island/cc/models/errors.py new file mode 100644 index 000000000..ff9eb13ea --- /dev/null +++ b/monkey/monkey_island/cc/models/errors.py @@ -0,0 +1,2 @@ +class MonkeyNotFoundError(Exception): + pass \ No newline at end of file diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index f2a8ae5a4..c0e4e88b2 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/monkey/monkey_island/cc/models/monkey.py @@ -5,6 +5,7 @@ import mongoengine from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, DateField, \ ReferenceField +from monkey_island.cc.models.errors import MonkeyNotFoundError from monkey_island.cc.models.monkey_ttl import MonkeyTtl @@ -32,6 +33,13 @@ class Monkey(Document): pba_results = ListField() ttl_ref = ReferenceField(MonkeyTtl) + @staticmethod + def get_single_monkey_by_id(db_id): + try: + return Monkey.objects(id=db_id)[0] + except IndexError: + raise MonkeyNotFoundError("id: {0}".format(str(db_id))) + def is_dead(self): monkey_is_dead = False if self.dead: @@ -41,7 +49,7 @@ class Monkey(Document): if MonkeyTtl.objects(id=self.ttl_ref.id).count() == 0: # No TTLs - monkey has timed out. The monkey is MIA. monkey_is_dead = True - except mongoengine.DoesNotExist: + except (mongoengine.DoesNotExist, AttributeError): # Trying to dereference unknown document - the monkey is MIA. monkey_is_dead = True return monkey_is_dead diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py index a79bf9bbb..31f3ff511 100644 --- a/monkey/monkey_island/cc/models/test_monkey.py +++ b/monkey/monkey_island/cc/models/test_monkey.py @@ -7,6 +7,7 @@ from unittest import TestCase import mongomock from monkey import Monkey +from monkey_island.cc.models.errors import MonkeyNotFoundError from monkey_ttl import MonkeyTtl @@ -35,3 +36,9 @@ class TestMonkey(TestCase): self.assertTrue(mia_monkey.is_dead()) self.assertFalse(alive_monkey.is_dead()) + def test_get_single_monkey_by_id(self): + a_monkey = Monkey(guid=str(uuid.uuid4())) + a_monkey.save() + + self.assertIsNotNone(Monkey.get_single_monkey_by_id(a_monkey.id)) + self.assertRaises(MonkeyNotFoundError, Monkey.get_single_monkey_by_id, "abcdefabcdefabcdefabcdef")