From e3b93f1506ada398b66701e4736ebba0b006d8b1 Mon Sep 17 00:00:00 2001
From: Shay Nehmad <shay.nehmad@guardicore.com>
Date: Wed, 2 Oct 2019 11:08:49 +0300
Subject: [PATCH] get_edge_label is a little quicker - uses cache. Still calls
 the DB on every run to check if something is a node or a monkey.

---
 monkey/monkey_island/cc/services/edge.py      | 16 +++++----
 .../cc/services/reporting/report.py           | 35 +++++++++----------
 2 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/monkey/monkey_island/cc/services/edge.py b/monkey/monkey_island/cc/services/edge.py
index eb23c3901..029e282d7 100644
--- a/monkey/monkey_island/cc/services/edge.py
+++ b/monkey/monkey_island/cc/services/edge.py
@@ -2,6 +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
 
 __author__ = "itay.mizeretz"
 
@@ -141,15 +142,18 @@ class EdgeService:
     @staticmethod
     def get_edge_label(edge):
         NodeService = monkey_island.cc.services.node.NodeService
-        from_label = NodeService.get_monkey_label(NodeService.get_monkey_by_id(edge["from"]))
-        if edge["to"] == ObjectId("000000000000000000000000"):
+        from_id = edge["from"]
+        to_id = edge["to"]
+
+        from_label = Monkey.get_label_by_id(from_id)
+
+        if to_id == ObjectId("000000000000000000000000"):
             to_label = 'MonkeyIsland'
         else:
-            to_id = NodeService.get_monkey_by_id(edge["to"])
-            if to_id is None:
-                to_label = NodeService.get_node_label(NodeService.get_node_by_id(edge["to"]))
+            if NodeService.get_monkey_by_id(to_id) is None:
+                to_label = NodeService.get_node_label(NodeService.get_node_by_id(to_id))
             else:
-                to_label = NodeService.get_monkey_label(to_id)
+                to_label = Monkey.get_label_by_id(to_id)
 
         RIGHT_ARROW = u"\u2192"
         return "%s %s %s" % (from_label, RIGHT_ARROW, to_label)
diff --git a/monkey/monkey_island/cc/services/reporting/report.py b/monkey/monkey_island/cc/services/reporting/report.py
index d792e7bfd..8215209f0 100644
--- a/monkey/monkey_island/cc/services/reporting/report.py
+++ b/monkey/monkey_island/cc/services/reporting/report.py
@@ -118,17 +118,10 @@ class ReportService:
 
     @staticmethod
     def get_scanned():
-
         formatted_nodes = []
 
-        # TODO Figure out and improve
-        # This part collects all the nodes in the DB. 2 accesses to the DB for getting all DB nodes and then
-        # get_displayed_node_by_id on all of them.
         nodes = ReportService.get_all_displayed_nodes()
 
-        print("2")
-
-        # for each node (n*...
         for node in nodes:
             nodes_that_can_access_current_node = node['accessible_from_nodes_hostnames']
             formatted_nodes.append(
@@ -156,23 +149,26 @@ class ReportService:
 
     @staticmethod
     def get_exploited():
-        exploited = \
+        exploited_with_monkeys = \
             [NodeService.get_displayed_node_by_id(monkey['_id'], True) for monkey in
-             mongo.db.monkey.find({}, {'_id': 1})
-             if not NodeService.get_monkey_manual_run(NodeService.get_monkey_by_id(monkey['_id']))] \
-            + [NodeService.get_displayed_node_by_id(node['_id'], True)
-               for node in mongo.db.node.find({'exploited': True}, {'_id': 1})]
+             mongo.db.monkey.find({}, {'_id': 1}) if
+             not NodeService.get_monkey_manual_run(NodeService.get_monkey_by_id(monkey['_id']))]
+
+        exploited_without_monkeys = [NodeService.get_displayed_node_by_id(node['_id'], True) for node in
+                                     mongo.db.node.find({'exploited': True}, {'_id': 1})]
+
+        exploited = exploited_with_monkeys + exploited_without_monkeys
 
         exploited = [
             {
-                'label': monkey['label'],
-                'ip_addresses': monkey['ip_addresses'],
-                'domain_name': monkey['domain_name'],
+                'label': exploited_node['label'],
+                'ip_addresses': exploited_node['ip_addresses'],
+                'domain_name': exploited_node['domain_name'],
                 'exploits': list(set(
-                    [ReportService.EXPLOIT_DISPLAY_DICT[exploit['exploiter']] for exploit in monkey['exploits'] if
-                     exploit['result']]))
+                    [ReportService.EXPLOIT_DISPLAY_DICT[exploit['exploiter']] for exploit in exploited_node['exploits']
+                     if exploit['result']]))
             }
-            for monkey in exploited]
+            for exploited_node in exploited]
 
         logger.info('Exploited nodes generated for reporting')
 
@@ -711,6 +707,7 @@ class ReportService:
         monkey_latest_modify_time = Monkey.get_latest_modifytime()
 
         scanned_nodes = ReportService.get_scanned()
+        exploited_nodes = ReportService.get_exploited()
         report = \
             {
                 'overview':
@@ -730,7 +727,7 @@ class ReportService:
                 'glance':
                     {
                         'scanned': scanned_nodes,
-                        'exploited': ReportService.get_exploited(),
+                        'exploited': exploited_nodes,
                         'stolen_creds': ReportService.get_stolen_creds(),
                         'azure_passwords': ReportService.get_azure_creds(),
                         'ssh_keys': ReportService.get_ssh_keys(),