Added monkey TTL renewal to monkey model.

In future, we should use only this method to update the monkey's TTL.

monkey/monkey_island/cc/consts.py

@@ -3,3 +3,5 @@ import os
 __author__ = 'itay.mizeretz'
 
 MONKEY_ISLAND_ABS_PATH = os.path.join(os.getcwd(), 'monkey_island')
+# DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS = 60 * 5
+DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS = 6

monkey/monkey_island/cc/models/monkey.py

@@ -5,7 +5,8 @@ import mongoengine
 from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, ReferenceField, \
     DateTimeField
 
-from monkey_island.cc.models.monkey_ttl import MonkeyTtl
+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
 
 
 class Monkey(Document):
@@ -42,6 +43,13 @@ class Monkey(Document):
         except IndexError:
             raise MonkeyNotFoundError("id: {0}".format(str(db_id)))
 
+    @staticmethod
+    def get_single_monkey_by_guid(monkey_guid):
+        try:
+            return Monkey.objects(guid=monkey_guid)[0]
+        except IndexError:
+            raise MonkeyNotFoundError("guid: {0}".format(str(monkey_guid)))
+
     @staticmethod
     def get_latest_modifytime():
         return Monkey.objects.order_by('-modifytime').first().modifytime
@@ -68,6 +76,9 @@ class Monkey(Document):
             os = "windows"
         return os
 
+    def renew_ttl(self, duration=DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS):
+        self.ttl_ref = create_monkey_ttl_document(duration)
+
 
 class MonkeyNotFoundError(Exception):
     pass

monkey/monkey_island/cc/models/monkey_ttl.py

@@ -38,3 +38,16 @@ class MonkeyTtl(Document):
     }
 
     expire_at = DateTimeField()
+
+
+def create_monkey_ttl_document(expiry_duration_in_seconds):
+    """
+    Create a new Monkey TTL document and save it as a document.
+    :param expiry_duration_in_seconds:  How long should the TTL last for. THIS IS A LOWER BOUND - depends on mongodb
+    performance.
+    :return: The TTL document. To get its ID use `.id`.
+    """
+    # The TTL data uses the new `models` module which depends on mongoengine.
+    current_ttl = MonkeyTtl.create_ttl_expire_in(expiry_duration_in_seconds)
+    current_ttl.save()
+    return current_ttl

monkey/monkey_island/cc/resources/monkey.py

@@ -5,26 +5,17 @@ import dateutil.parser
 import flask_restful
 from flask import request
 
+from monkey_island.cc.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
 from monkey_island.cc.database import mongo
-from monkey_island.cc.models.monkey_ttl import MonkeyTtl
+from monkey_island.cc.models.monkey_ttl import create_monkey_ttl_document
 from monkey_island.cc.services.config import ConfigService
 from monkey_island.cc.services.node import NodeService
 
-MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS = 60 * 5
-
 __author__ = 'Barak'
 
 # TODO: separate logic from interface
 
 
-def create_monkey_ttl():
-    # The TTL data uses the new `models` module which depends on mongoengine.
-    current_ttl = MonkeyTtl.create_ttl_expire_in(MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS)
-    current_ttl.save()
-    ttlid = current_ttl.id
-    return ttlid
-
-
 class Monkey(flask_restful.Resource):
 
     # Used by monkey. can't secure.
@@ -58,8 +49,8 @@ class Monkey(flask_restful.Resource):
             tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "")
             NodeService.set_monkey_tunnel(monkey["_id"], tunnel_host_ip)
 
-        ttlid = create_monkey_ttl()
+        ttl = create_monkey_ttl_document(DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS)
-        update['$set']['ttl_ref'] = ttlid
+        update['$set']['ttl_ref'] = ttl.id
 
         return mongo.db.monkey.update({"_id": monkey["_id"]}, update, upsert=False)
 
@@ -120,7 +111,8 @@ class Monkey(flask_restful.Resource):
             tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "")
             monkey_json.pop('tunnel')
 
-        monkey_json['ttl_ref'] = create_monkey_ttl()
+        ttl = create_monkey_ttl_document(DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS)
+        monkey_json['ttl_ref'] = ttl.id
 
         mongo.db.monkey.update({"guid": monkey_json["guid"]},
                                {"$set": monkey_json},

monkey/monkey_island/cc/resources/telemetry.py

@@ -15,10 +15,10 @@ from monkey_island.cc.services.edge import EdgeService
 from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.encryptor import encryptor
 from monkey_island.cc.services.wmi_handler import WMIHandler
+from monkey_island.cc.models.monkey import Monkey
 
 __author__ = 'Barak'
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -49,6 +49,9 @@ class Telemetry(flask_restful.Resource):
         telemetry_json = json.loads(request.data)
         telemetry_json['timestamp'] = datetime.now()
 
+        # Monkey communicated, so it's alive. Update the TTL.
+        Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid']).renew_ttl()
+
         monkey = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid'])
 
         try:
@@ -59,7 +62,7 @@ class Telemetry(flask_restful.Resource):
             else:
                 logger.info('Got unknown type of telemetry: %s' % telem_category)
         except Exception as ex:
-            logger.error("Exception caught while processing telemetry", exc_info=True)
+            logger.error("Exception caught while processing telemetry. Info: {}".format(ex.message), exc_info=True)
 
         telem_id = mongo.db.telemetry.insert(telemetry_json)
         return mongo.db.telemetry.find_one_or_404({"_id": telem_id})