Added TTL creation to post as well

monkey/monkey_island/cc/models/errors.py

@@ -1,2 +1,2 @@
 class MonkeyNotFoundError(Exception):
-    pass
+    pass

monkey/monkey_island/cc/resources/monkey.py

@@ -17,6 +17,16 @@ __author__ = 'Barak'
 # TODO: separate logic from interface
 
 
+def create_monkey_ttl():
+    # The TTL data uses the new `models` module which depends on mongoengine.
+    # Using UTC to make the mongodb TTL feature work. See
+    # https://stackoverflow.com/questions/55994379/mongodb-ttl-index-doesnt-delete-expired-documents.
+    current_ttl = MonkeyTtl(expire_at=datetime.utcnow() + timedelta(seconds=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.
@@ -50,13 +60,8 @@ class Monkey(flask_restful.Resource):
             tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "")
             NodeService.set_monkey_tunnel(monkey["_id"], tunnel_host_ip)
 
-        # The TTL data uses the new `models` module which depends on mongoengine.
+        ttlid = create_monkey_ttl()
-        # Using UTC to make the mongodb TTL feature work. See
+        update['$set']['ttl_ref'] = ttlid
-        # https://stackoverflow.com/questions/55994379/mongodb-ttl-index-doesnt-delete-expired-documents.
-        current_ttl = MonkeyTtl(expire_at=datetime.utcnow() + timedelta(seconds=MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS))
-        current_ttl.save()
-
-        update['$set']['ttl_ref'] = current_ttl.id
 
         return mongo.db.monkey.update({"_id": monkey["_id"]}, update, upsert=False)
 
@@ -117,6 +122,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()
+
         mongo.db.monkey.update({"guid": monkey_json["guid"]},
                                {"$set": monkey_json},
                                upsert=True)