From d02e349f71d0b32a25701091511964344401c889 Mon Sep 17 00:00:00 2001
From: Shay Nehmad <shay.nehmad@guardicore.com>
Date: Wed, 2 Oct 2019 12:18:53 +0300
Subject: [PATCH] Added cached checking of is_monkey to optimise runtime of
 EdgeService

---
 monkey/monkey_island/cc/models/monkey.py      | 11 +++++++++
 monkey/monkey_island/cc/models/test_monkey.py | 23 +++++++++++++++++++
 monkey/monkey_island/cc/services/edge.py      |  6 ++---
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py
index 57dd26578..6644fccf7 100644
--- a/monkey/monkey_island/cc/models/monkey.py
+++ b/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)
diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py
index d0b5efd71..c2d860e84 100644
--- a/monkey/monkey_island/cc/models/test_monkey.py
+++ b/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)
+
diff --git a/monkey/monkey_island/cc/services/edge.py b/monkey/monkey_island/cc/services/edge.py
index 029e282d7..1622d5233 100644
--- a/monkey/monkey_island/cc/services/edge.py
+++ b/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)