diff --git a/monkey/monkey_island/cc/models/monkey_ttl.py b/monkey/monkey_island/cc/models/monkey_ttl.py index 20d3d3eb0..63c1ad4db 100644 --- a/monkey/monkey_island/cc/models/monkey_ttl.py +++ b/monkey/monkey_island/cc/models/monkey_ttl.py @@ -9,9 +9,14 @@ class MonkeyTtl(Document): See https://docs.mongodb.com/manual/tutorial/expire-data/ and https://stackoverflow.com/questions/55994379/mongodb-ttl-index-doesnt-delete-expired-documents/56021663#56021663 for more information about how TTL indexing works and why this class is set up the way it is. + + If you wish to use this class, you can create it using the create_ttl_expire_in(seconds) function. + If you wish to create an instance of this class directly, see the inner implementation of + create_ttl_expire_in(seconds) to see how to do so. """ - def __init__(self, expiry_in_seconds, *args, **values): + @staticmethod + def create_ttl_expire_in(expiry_in_seconds): """ Initializes a TTL object which will expire in expire_in_seconds seconds from when created. Remember to call .save() on the object after creation. @@ -20,8 +25,7 @@ class MonkeyTtl(Document): """ # Using UTC to make the mongodb TTL feature work. See # https://stackoverflow.com/questions/55994379/mongodb-ttl-index-doesnt-delete-expired-documents. - super(MonkeyTtl, self).__init__( - expire_at=datetime.utcnow() + timedelta(seconds=expiry_in_seconds), *args, **values) + return MonkeyTtl(expire_at=datetime.utcnow() + timedelta(seconds=expiry_in_seconds)) meta = { 'indexes': [ diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py index 8f54635c0..41920d2a0 100644 --- a/monkey/monkey_island/cc/models/test_monkey.py +++ b/monkey/monkey_island/cc/models/test_monkey.py @@ -1,5 +1,4 @@ import uuid -from datetime import timedelta, datetime from time import sleep from unittest import TestCase @@ -21,7 +20,7 @@ class TestMonkey(TestCase): """ def test_is_dead(self): # Arrange - alive_monkey_ttl = MonkeyTtl(30) + alive_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30) alive_monkey_ttl.save() alive_monkey = Monkey( guid=str(uuid.uuid4()), @@ -30,7 +29,7 @@ class TestMonkey(TestCase): alive_monkey.save() # MIA stands for Missing In Action - mia_monkey_ttl = MonkeyTtl(30) + mia_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30) mia_monkey_ttl.save() mia_monkey = Monkey(guid=str(uuid.uuid4()), dead=False, ttl_ref=mia_monkey_ttl) mia_monkey.save() diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py index a5544c5ce..0207709eb 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/monkey/monkey_island/cc/resources/monkey.py @@ -1,12 +1,12 @@ import json -from datetime import datetime, timedelta +from datetime import datetime import dateutil.parser import flask_restful from flask import request -from monkey_island.cc.models.monkey_ttl import MonkeyTtl from monkey_island.cc.database import mongo +from monkey_island.cc.models.monkey_ttl import MonkeyTtl from monkey_island.cc.services.config import ConfigService from monkey_island.cc.services.node import NodeService @@ -19,7 +19,7 @@ __author__ = 'Barak' def create_monkey_ttl(): # The TTL data uses the new `models` module which depends on mongoengine. - current_ttl = MonkeyTtl(MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS) + current_ttl = MonkeyTtl.create_ttl_expire_in(MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS) current_ttl.save() ttlid = current_ttl.id return ttlid