diff --git a/monkey/monkey_island/cc/models/monkey_ttl.py b/monkey/monkey_island/cc/models/monkey_ttl.py index 853edb9b0..20d3d3eb0 100644 --- a/monkey/monkey_island/cc/models/monkey_ttl.py +++ b/monkey/monkey_island/cc/models/monkey_ttl.py @@ -1,3 +1,5 @@ +from datetime import datetime, timedelta + from mongoengine import Document, DateTimeField @@ -6,11 +8,21 @@ class MonkeyTtl(Document): This model represents the monkey's TTL, and is referenced by the main Monkey 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. - - When initializing this object, do it like so: - t = MonkeyTtl(expire_at=datetime.utcnow() + timedelta(seconds=XXX)) + for more information about how TTL indexing works and why this class is set up the way it is. """ + + def __init__(self, expiry_in_seconds, *args, **values): + """ + Initializes a TTL object which will expire in expire_in_seconds seconds from when created. + Remember to call .save() on the object after creation. + :param expiry_in_seconds: How long should the TTL be in the DB, in seconds. Please take into consideration + that the cleanup thread of mongo might take extra time to delete the TTL from the DB. + """ + # 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) + meta = { 'indexes': [ { diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py index ad5fc17de..8f54635c0 100644 --- a/monkey/monkey_island/cc/models/test_monkey.py +++ b/monkey/monkey_island/cc/models/test_monkey.py @@ -21,7 +21,7 @@ class TestMonkey(TestCase): """ def test_is_dead(self): # Arrange - alive_monkey_ttl = MonkeyTtl(expire_at=datetime.now() + timedelta(seconds=30)) + alive_monkey_ttl = MonkeyTtl(30) alive_monkey_ttl.save() alive_monkey = Monkey( guid=str(uuid.uuid4()), @@ -30,7 +30,7 @@ class TestMonkey(TestCase): alive_monkey.save() # MIA stands for Missing In Action - mia_monkey_ttl = MonkeyTtl(expire_at=datetime.now() + timedelta(seconds=30)) + mia_monkey_ttl = MonkeyTtl(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 e7744f007..a5544c5ce 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/monkey/monkey_island/cc/resources/monkey.py @@ -19,9 +19,7 @@ __author__ = 'Barak' 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 = MonkeyTtl(MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS) current_ttl.save() ttlid = current_ttl.id return ttlid