From a4ebfb0c3dff47b5b39d3d1d3a353750a78b5d05 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Tue, 15 Oct 2019 11:12:49 +0300 Subject: [PATCH] Fixed ring bugs --- monkey/monkey_island/cc/models/monkey.py | 19 ++++++++++++++----- monkey/monkey_island/cc/models/test_monkey.py | 12 ++++++------ 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index a8a7da2ec..fe0745496 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/monkey/monkey_island/cc/models/monkey.py @@ -45,20 +45,29 @@ class Monkey(Document): aws_instance_id = StringField(required=False) # This field only exists when the monkey is running on an AWS # instance. See https://github.com/guardicore/monkey/issues/426. + @staticmethod + def __ring_key__(): + """ + Cash key representation + https://ring-cache.readthedocs.io/en/stable/quickstart.html#method-classmethod-staticmethod-property + :return: + """ + return Monkey.guid + # LOGIC @staticmethod def get_single_monkey_by_id(db_id): try: return Monkey.objects.get(id=db_id) except DoesNotExist as ex: - raise MonkeyNotFoundError("info: {0} | id: {1}".format(ex.message, str(db_id))) + raise MonkeyNotFoundError("info: {0} | id: {1}".format(ex, str(db_id))) @staticmethod def get_single_monkey_by_guid(monkey_guid): try: return Monkey.objects.get(guid=monkey_guid) except DoesNotExist as ex: - raise MonkeyNotFoundError("info: {0} | guid: {1}".format(ex.message, str(monkey_guid))) + raise MonkeyNotFoundError("info: {0} | guid: {1}".format(ex, str(monkey_guid))) @staticmethod def get_latest_modifytime(): @@ -88,8 +97,8 @@ class Monkey(Document): os = "windows" return os - @staticmethod @ring.lru() + @staticmethod def get_label_by_id(object_id): current_monkey = Monkey.get_single_monkey_by_id(object_id) label = Monkey.get_hostname_by_id(object_id) + " : " + current_monkey.ip_addresses[0] @@ -97,8 +106,8 @@ class Monkey(Document): label = "MonkeyIsland - " + label return label - @staticmethod @ring.lru() + @staticmethod def get_hostname_by_id(object_id): """ :param object_id: the object ID of a Monkey in the database. @@ -124,10 +133,10 @@ class Monkey(Document): """ return {'ips': self.ip_addresses, 'hostname': self.hostname} - @staticmethod @ring.lru( expire=1 # data has TTL of 1 second. This is useful for rapid calls for report generation. ) + @staticmethod def is_monkey(object_id): try: _ = Monkey.get_single_monkey_by_id(object_id) diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py index 7a920409c..1bb805e76 100644 --- a/monkey/monkey_island/cc/models/test_monkey.py +++ b/monkey/monkey_island/cc/models/test_monkey.py @@ -126,7 +126,7 @@ class TestMonkey(IslandTestCase): linux_monkey.save() cache_info_before_query = Monkey.get_label_by_id.storage.backend.cache_info() - self.assertEquals(cache_info_before_query.hits, 0) + self.assertEqual(cache_info_before_query.hits, 0) # not cached label = Monkey.get_label_by_id(linux_monkey.id) @@ -138,7 +138,7 @@ class TestMonkey(IslandTestCase): # should be cached _ = Monkey.get_label_by_id(linux_monkey.id) cache_info_after_query = Monkey.get_label_by_id.storage.backend.cache_info() - self.assertEquals(cache_info_after_query.hits, 1) + self.assertEqual(cache_info_after_query.hits, 1) linux_monkey.set_hostname("Another hostname") @@ -146,8 +146,8 @@ class TestMonkey(IslandTestCase): label = Monkey.get_label_by_id(linux_monkey.id) cache_info_after_second_query = Monkey.get_label_by_id.storage.backend.cache_info() # still 1 hit only - self.assertEquals(cache_info_after_second_query.hits, 1) - self.assertEquals(cache_info_after_second_query.misses, 2) + self.assertEqual(cache_info_after_second_query.hits, 1) + self.assertEqual(cache_info_after_second_query.misses, 2) def test_is_monkey(self): self.fail_if_not_testing_env() @@ -157,7 +157,7 @@ class TestMonkey(IslandTestCase): a_monkey.save() cache_info_before_query = Monkey.is_monkey.storage.backend.cache_info() - self.assertEquals(cache_info_before_query.hits, 0) + self.assertEqual(cache_info_before_query.hits, 0) # not cached self.assertTrue(Monkey.is_monkey(a_monkey.id)) @@ -169,5 +169,5 @@ class TestMonkey(IslandTestCase): self.assertFalse(Monkey.is_monkey(fake_id)) cache_info_after_query = Monkey.is_monkey.storage.backend.cache_info() - self.assertEquals(cache_info_after_query.hits, 2) + self.assertEqual(cache_info_after_query.hits, 2)