p17860254
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Now saving after setting new TTL ref on renewal 

This changed the schema since the parent field is nullable.
This commit is contained in:
Shay Nehmad 2019-07-22 19:17:29 +03:00
parent 1201a3d138
commit f7eee5175a
1 changed files with 10 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
monkey/monkey_island/cc/models/monkey.py
View File

@ -1,17 +1,13 @@
""" """
Define a Document Schema for the Monkey document. Define a Document Schema for the Monkey document.
""" """
import mongoengine
import logging
from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, ReferenceField, \ from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, ReferenceField, \
DateTimeField DateTimeField, DynamicField, DoesNotExist
from monkey_island.cc.models.monkey_ttl import MonkeyTtl, create_monkey_ttl_document from monkey_island.cc.models.monkey_ttl import MonkeyTtl, create_monkey_ttl_document
from monkey_island.cc.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS from monkey_island.cc.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
logger = logging.getLogger(__name__)
class Monkey(Document): class Monkey(Document):
""" """
This class has 2 main section: This class has 2 main section:
@ -30,8 +26,11 @@ class Monkey(Document):
ip_addresses = ListField(StringField()) ip_addresses = ListField(StringField())
keepalive = DateTimeField() keepalive = DateTimeField()
modifytime = DateTimeField() modifytime = DateTimeField()
# TODO change this to an embedded document as well - RN it's an unnamed tuple which is confusing. # TODO make "parent" an embedded document, so this can be removed and the schema explained (and validated) verbosly.
parent = ListField(ListField(StringField())) # This is a temporary fix, since mongoengine doesn't allow for lists of strings to be null
# (even with required=False of null=True).
# See relevant issue: https://github.com/MongoEngine/mongoengine/issues/1904
parent = ListField(ListField(DynamicField()))
config_error = BooleanField() config_error = BooleanField()
critical_services = ListField(StringField()) critical_services = ListField(StringField())
pba_results = ListField() pba_results = ListField()
@ -43,14 +42,14 @@ class Monkey(Document):
def get_single_monkey_by_id(db_id): def get_single_monkey_by_id(db_id):
try: try:
return Monkey.objects.get(id=db_id) return Monkey.objects.get(id=db_id)
except mongoengine.DoesNotExist as ex: except DoesNotExist as ex:
raise MonkeyNotFoundError("info: {0} | id: {1}".format(ex.message, str(db_id))) raise MonkeyNotFoundError("info: {0} | id: {1}".format(ex.message, str(db_id)))
@staticmethod @staticmethod
def get_single_monkey_by_guid(monkey_guid): def get_single_monkey_by_guid(monkey_guid):
try: try:
return Monkey.objects.get(guid=monkey_guid) return Monkey.objects.get(guid=monkey_guid)
except mongoengine.DoesNotExist as ex: except DoesNotExist as ex:
raise MonkeyNotFoundError("info: {0} | guid: {1}".format(ex.message, str(monkey_guid))) raise MonkeyNotFoundError("info: {0} | guid: {1}".format(ex.message, str(monkey_guid)))
@staticmethod @staticmethod
@ -68,7 +67,7 @@ class Monkey(Document):
if MonkeyTtl.objects(id=self.ttl_ref.id).count() == 0: if MonkeyTtl.objects(id=self.ttl_ref.id).count() == 0:
# No TTLs - monkey has timed out. The monkey is MIA. # No TTLs - monkey has timed out. The monkey is MIA.
monkey_is_dead = True monkey_is_dead = True
except (mongoengine.DoesNotExist, AttributeError): except (DoesNotExist, AttributeError):
# Trying to dereference unknown document - the monkey is MIA. # Trying to dereference unknown document - the monkey is MIA.
monkey_is_dead = True monkey_is_dead = True
return monkey_is_dead return monkey_is_dead
@ -83,6 +82,7 @@ class Monkey(Document):
def renew_ttl(self, duration=DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS): def renew_ttl(self, duration=DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS):
self.ttl_ref = create_monkey_ttl_document(duration) self.ttl_ref = create_monkey_ttl_document(duration)
self.save()
class MonkeyNotFoundError(Exception): class MonkeyNotFoundError(Exception):