diff --git a/monkey/monkey_island/cc/models/errors.py b/monkey/monkey_island/cc/models/errors.py index ff9eb13ea..69cb529dd 100644 --- a/monkey/monkey_island/cc/models/errors.py +++ b/monkey/monkey_island/cc/models/errors.py @@ -1,2 +1,2 @@ class MonkeyNotFoundError(Exception): - pass \ No newline at end of file + pass diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py index e1f52f73e..e7744f007 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/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. - # 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() - - update['$set']['ttl_ref'] = current_ttl.id + ttlid = create_monkey_ttl() + update['$set']['ttl_ref'] = ttlid 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)