diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py new file mode 100644 index 000000000..c02d7f0c1 --- /dev/null +++ b/monkey/monkey_island/cc/models/__init__.py @@ -0,0 +1 @@ +from monkey import Monkey diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py new file mode 100644 index 000000000..d54fe1b43 --- /dev/null +++ b/monkey/monkey_island/cc/models/monkey.py @@ -0,0 +1,89 @@ +""" +Define a Document Schema for the Monkey document. +""" +import mongoengine +from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, DateField, \ + EmbeddedDocument, connect, ReferenceField, DateTimeField + +from monkey_island.cc.environment.environment import env + +connect(db=env.mongo_db_name, host=env.mongo_db_host, port=env.mongo_db_port) + + +class Config(EmbeddedDocument): + """ + No need to define this schema here. It will change often and is already is defined in + monkey_island.cc.services.config_schema. + See https://mongoengine-odm.readthedocs.io/apireference.html#mongoengine.FieldDoesNotExist + """ + meta = {'strict': False} + pass + + +class Creds(EmbeddedDocument): + """ + TODO get an example of this data + """ + meta = {'strict': False} + pass + + +class PbaResults(EmbeddedDocument): + ip = StringField() + hostname = StringField() + command = StringField() + name = StringField() + result = ListField() + + +class Ttl(Document): + meta = { + 'indexes': [ + { + 'name': 'TTL_index', + 'fields': ['expire_at'], + 'expireAfterSeconds': 0 + } + ] + } + + expire_at = DateTimeField() + + +class Monkey(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, somewhat like an API. + """ + guid = StringField(required=True) + config = EmbeddedDocumentField('Config') + creds = ListField(EmbeddedDocumentField('Creds')) + dead = BooleanField() + description = StringField() + hostname = StringField() + internet_access = BooleanField() + ip_addresses = ListField(StringField()) + keepalive = DateField() + modifytime = DateField() + # TODO change this to an embedded document as well - RN it's an unnamed tuple which is confusing. + parent = ListField(ListField(StringField())) + config_error = BooleanField() + critical_services = ListField(StringField()) + pba_results = ListField() + ttl_ref = ReferenceField(Ttl) + + def is_dead(self): + monkey_is_dead = False + if self.dead: + monkey_is_dead = True + else: + try: + if Ttl.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: + # Trying to dereference unknown document + monkey_is_dead = True + return monkey_is_dead diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py index 057ebf149..c6e7af908 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/monkey/monkey_island/cc/resources/monkey.py @@ -1,13 +1,13 @@ import json -from datetime import datetime +from datetime import datetime, timedelta import dateutil.parser import flask_restful from flask import request +from monkey_island.cc import models from monkey_island.cc.database import mongo from monkey_island.cc.services.config import ConfigService -from monkey_island.cc.services.monkey_timeout import start_timer_decorator from monkey_island.cc.services.node import NodeService __author__ = 'Barak' @@ -18,7 +18,6 @@ __author__ = 'Barak' class Monkey(flask_restful.Resource): # Used by monkey. can't secure. - @start_timer_decorator def get(self, guid=None, **kw): NodeService.update_dead_monkeys() # refresh monkeys status if not guid: @@ -49,6 +48,11 @@ class Monkey(flask_restful.Resource): tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "") NodeService.set_monkey_tunnel(monkey["_id"], tunnel_host_ip) + current_ttl = models.monkey.Ttl(expire_at=datetime.now() + timedelta(seconds=30)) + current_ttl.save() + + update['$set']['ttl_ref'] = current_ttl.id + return mongo.db.monkey.update({"_id": monkey["_id"]}, update, upsert=False) # Used by monkey. can't secure. diff --git a/monkey/monkey_island/cc/services/node.py b/monkey/monkey_island/cc/services/node.py index fa500aab5..83474a2c3 100644 --- a/monkey/monkey_island/cc/services/node.py +++ b/monkey/monkey_island/cc/services/node.py @@ -7,6 +7,7 @@ from monkey_island.cc.database import mongo from monkey_island.cc.services.edge import EdgeService from monkey_island.cc.utils import local_ip_addresses import socket +from monkey_island.cc import models __author__ = "itay.mizeretz" @@ -141,7 +142,7 @@ class NodeService: "label": label, "group": NodeService.get_monkey_group(monkey), "os": NodeService.get_monkey_os(monkey), - "dead": monkey["dead"], + "dead": models.Monkey.objects(id=monkey["_id"])[0].is_dead(), "domain_name": "", "pba_results": monkey["pba_results"] if "pba_results" in monkey else [] } @@ -293,7 +294,7 @@ class NodeService: @staticmethod def is_any_monkey_alive(): - return mongo.db.monkey.find_one({'dead': False}) is not None + return models.Monkey.objects(dead=False).count() > 0 @staticmethod def is_any_monkey_exists(): diff --git a/monkey/monkey_island/requirements.txt b/monkey/monkey_island/requirements.txt index 5e0faea19..953cad2cb 100644 --- a/monkey/monkey_island/requirements.txt +++ b/monkey/monkey_island/requirements.txt @@ -22,3 +22,4 @@ bson cffi virtualenv wheel +mongoengine