From 8046cddb8b7e53e1d15694709d6a1c0aec9b2a1a Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Wed, 23 Oct 2019 09:19:03 +0300 Subject: [PATCH] Fixed ring bug with @staticmethod --- monkey/monkey_island/cc/models/monkey.py | 65 ++++++++++--------- monkey/monkey_island/cc/models/test_monkey.py | 26 ++++---- monkey/monkey_island/cc/services/edge.py | 8 +-- monkey/monkey_island/cc/services/node.py | 8 +-- 4 files changed, 54 insertions(+), 53 deletions(-) diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index 324903809..daeb9ea5b 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/monkey/monkey_island/cc/models/monkey.py @@ -97,25 +97,6 @@ class Monkey(Document): os = "windows" return os - @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] - if len(set(current_monkey.ip_addresses).intersection(local_ip_addresses())) > 0: - label = "MonkeyIsland - " + label - return label - - @ring.lru() - @staticmethod - def get_hostname_by_id(object_id): - """ - :param object_id: the object ID of a Monkey in the database. - :return: The hostname of that machine. - :note: Use this and not monkey.hostname for performance - this is lru-cached. - """ - return Monkey.get_single_monkey_by_id(object_id).hostname - def set_hostname(self, hostname): """ Sets a new hostname for a machine and clears the cache for getting it. @@ -123,8 +104,8 @@ class Monkey(Document): """ self.hostname = hostname self.save() - Monkey.get_hostname_by_id.delete(self.id) - Monkey.get_label_by_id.delete(self.id) + get_monkey_hostname_by_id.delete(self.id) + get_monkey_label_by_id.delete(self.id) def get_network_info(self): """ @@ -133,17 +114,6 @@ class Monkey(Document): """ return {'ips': self.ip_addresses, 'hostname': self.hostname} - @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) - return True - except: - return False - @staticmethod def get_tunneled_monkeys(): return Monkey.objects(tunnel__exists=True) @@ -153,5 +123,36 @@ class Monkey(Document): self.save() +# Can't make following methods static under Monkey class due to ring bug +@ring.lru( + expire=1 # data has TTL of 1 second. This is useful for rapid calls for report generation. +) +def is_monkey(object_id): + try: + _ = Monkey.get_single_monkey_by_id(object_id) + return True + except: + return False + + +@ring.lru() +def get_monkey_label_by_id(object_id): + current_monkey = Monkey.get_single_monkey_by_id(object_id) + label = get_monkey_hostname_by_id(object_id) + " : " + current_monkey.ip_addresses[0] + if len(set(current_monkey.ip_addresses).intersection(local_ip_addresses())) > 0: + label = "MonkeyIsland - " + label + return label + + +@ring.lru() +def get_monkey_hostname_by_id(object_id): + """ + :param object_id: the object ID of a Monkey in the database. + :return: The hostname of that machine. + :note: Use this and not monkey.hostname for performance - this is lru-cached. + """ + return Monkey.get_single_monkey_by_id(object_id).hostname + + class MonkeyNotFoundError(Exception): pass diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py index 1bb805e76..fb9b329b1 100644 --- a/monkey/monkey_island/cc/models/test_monkey.py +++ b/monkey/monkey_island/cc/models/test_monkey.py @@ -2,7 +2,7 @@ import uuid from time import sleep from .monkey import Monkey -from monkey_island.cc.models.monkey import MonkeyNotFoundError +from monkey_island.cc.models.monkey import MonkeyNotFoundError, is_monkey, get_monkey_label_by_id from monkey_island.cc.testing.IslandTestCase import IslandTestCase from .monkey_ttl import MonkeyTtl @@ -125,26 +125,26 @@ class TestMonkey(IslandTestCase): ip_addresses=[ip_example]) linux_monkey.save() - cache_info_before_query = Monkey.get_label_by_id.storage.backend.cache_info() + cache_info_before_query = get_monkey_label_by_id.storage.backend.cache_info() self.assertEqual(cache_info_before_query.hits, 0) # not cached - label = Monkey.get_label_by_id(linux_monkey.id) + label = get_monkey_label_by_id(linux_monkey.id) self.assertIsNotNone(label) self.assertIn(hostname_example, label) self.assertIn(ip_example, label) # should be cached - _ = Monkey.get_label_by_id(linux_monkey.id) - cache_info_after_query = Monkey.get_label_by_id.storage.backend.cache_info() + _ = get_monkey_label_by_id(linux_monkey.id) + cache_info_after_query = get_monkey_label_by_id.storage.backend.cache_info() self.assertEqual(cache_info_after_query.hits, 1) linux_monkey.set_hostname("Another hostname") # should be a miss - label = Monkey.get_label_by_id(linux_monkey.id) - cache_info_after_second_query = Monkey.get_label_by_id.storage.backend.cache_info() + label = get_monkey_label_by_id(linux_monkey.id) + cache_info_after_second_query = get_monkey_label_by_id.storage.backend.cache_info() # still 1 hit only self.assertEqual(cache_info_after_second_query.hits, 1) self.assertEqual(cache_info_after_second_query.misses, 2) @@ -156,18 +156,18 @@ class TestMonkey(IslandTestCase): a_monkey = Monkey(guid=str(uuid.uuid4())) a_monkey.save() - cache_info_before_query = Monkey.is_monkey.storage.backend.cache_info() + cache_info_before_query = is_monkey.storage.backend.cache_info() self.assertEqual(cache_info_before_query.hits, 0) # not cached - self.assertTrue(Monkey.is_monkey(a_monkey.id)) + self.assertTrue(is_monkey(a_monkey.id)) fake_id = "123456789012" - self.assertFalse(Monkey.is_monkey(fake_id)) + self.assertFalse(is_monkey(fake_id)) # should be cached - self.assertTrue(Monkey.is_monkey(a_monkey.id)) - self.assertFalse(Monkey.is_monkey(fake_id)) + self.assertTrue(is_monkey(a_monkey.id)) + self.assertFalse(is_monkey(fake_id)) - cache_info_after_query = Monkey.is_monkey.storage.backend.cache_info() + cache_info_after_query = is_monkey.storage.backend.cache_info() self.assertEqual(cache_info_after_query.hits, 2) diff --git a/monkey/monkey_island/cc/services/edge.py b/monkey/monkey_island/cc/services/edge.py index ae3d2a2de..bf9417309 100644 --- a/monkey/monkey_island/cc/services/edge.py +++ b/monkey/monkey_island/cc/services/edge.py @@ -2,7 +2,7 @@ from bson import ObjectId from monkey_island.cc.database import mongo import monkey_island.cc.services.node -from monkey_island.cc.models import Monkey +from monkey_island.cc.models.monkey import get_monkey_label_by_id, is_monkey __author__ = "itay.mizeretz" @@ -145,13 +145,13 @@ class EdgeService: from_id = edge["from"] to_id = edge["to"] - from_label = Monkey.get_label_by_id(from_id) + from_label = get_monkey_label_by_id(from_id) if to_id == ObjectId("000000000000000000000000"): to_label = 'MonkeyIsland' else: - if Monkey.is_monkey(to_id): - to_label = Monkey.get_label_by_id(to_id) + if is_monkey(to_id): + to_label = get_monkey_label_by_id(to_id) else: to_label = NodeService.get_node_label(NodeService.get_node_by_id(to_id)) diff --git a/monkey/monkey_island/cc/services/node.py b/monkey/monkey_island/cc/services/node.py index 27d2d299a..0c0a873e8 100644 --- a/monkey/monkey_island/cc/services/node.py +++ b/monkey/monkey_island/cc/services/node.py @@ -4,7 +4,7 @@ from bson import ObjectId import monkey_island.cc.services.log from monkey_island.cc.database import mongo -from monkey_island.cc.models import Monkey +from monkey_island.cc.models.monkey import Monkey, get_monkey_hostname_by_id, get_monkey_label_by_id from monkey_island.cc.services.edge import EdgeService from monkey_island.cc.utils import local_ip_addresses import socket @@ -50,8 +50,8 @@ class NodeService: for edge in edges: from_node_id = edge["from"] - from_node_label = Monkey.get_label_by_id(from_node_id) - from_node_hostname = Monkey.get_hostname_by_id(from_node_id) + from_node_label = get_monkey_label_by_id(from_node_id) + from_node_hostname = get_monkey_hostname_by_id(from_node_id) accessible_from_nodes.append(from_node_label) accessible_from_nodes_hostnames.append(from_node_hostname) @@ -140,7 +140,7 @@ class NodeService: @staticmethod def monkey_to_net_node(monkey, for_report=False): monkey_id = monkey["_id"] - label = Monkey.get_hostname_by_id(monkey_id) if for_report else Monkey.get_label_by_id(monkey_id) + label = get_monkey_hostname_by_id(monkey_id) if for_report else get_monkey_label_by_id(monkey_id) monkey_group = NodeService.get_monkey_group(monkey) return \ {