Added cached checking of is_monkey to optimise runtime of EdgeService

monkey/monkey_island/cc/models/monkey.py

@@ -123,6 +123,17 @@ 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.
+    )
+    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)

monkey/monkey_island/cc/models/test_monkey.py

@@ -139,3 +139,26 @@ class TestMonkey(IslandTestCase):
         _ = 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)
+
+    def test_is_monkey(self):
+        self.fail_if_not_testing_env()
+        self.clean_monkey_db()
+
+        a_monkey = Monkey(guid=str(uuid.uuid4()))
+        a_monkey.save()
+
+        cache_info_before_query = Monkey.is_monkey.storage.backend.cache_info()
+        self.assertEquals(cache_info_before_query.hits, 0)
+
+        # not cached
+        self.assertTrue(Monkey.is_monkey(a_monkey.id))
+        fake_id = "123456789012"
+        self.assertFalse(Monkey.is_monkey(fake_id))
+
+        # should be cached
+        self.assertTrue(Monkey.is_monkey(a_monkey.id))
+        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)
+

monkey/monkey_island/cc/services/edge.py

@@ -150,10 +150,10 @@ class EdgeService:
         if to_id == ObjectId("000000000000000000000000"):
             to_label = 'MonkeyIsland'
         else:
-            if NodeService.get_monkey_by_id(to_id) is None:
-                to_label = NodeService.get_node_label(NodeService.get_node_by_id(to_id))
-            else:
+            if Monkey.is_monkey(to_id):
                 to_label = Monkey.get_label_by_id(to_id)
+            else:
+                to_label = NodeService.get_node_label(NodeService.get_node_by_id(to_id))
 
         RIGHT_ARROW = u"\u2192"
         return "%s %s %s" % (from_label, RIGHT_ARROW, to_label)