Started improving and researching the performence issues - still in progress...

monkey/monkey_island/cc/models/monkey.py

@@ -3,11 +3,14 @@ Define a Document Schema for the Monkey document.
 """
 from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, ReferenceField, \
     DateTimeField, DynamicField, DoesNotExist
+import ring
 
 from monkey_island.cc.models.monkey_ttl import MonkeyTtl, create_monkey_ttl_document
 from monkey_island.cc.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
 from monkey_island.cc.models.command_control_channel import CommandControlChannel
 
+MAX_MONKEYS_AMOUNT_TO_CACHE = 100
+
 
 class Monkey(Document):
     """
@@ -84,6 +87,28 @@ class Monkey(Document):
             os = "windows"
         return os
 
+    # TODO This is not a field therefore cache shouldn't be here
+    @ring.lru()
+    @staticmethod
+    def get_label_by_id(object_id):
+        print("Actually in get_label_by_id - not cached for {}".format(str(object_id)))
+        current_monkey = Monkey.get_single_monkey_by_id(object_id)
+        return Monkey.get_hostname_by_id(object_id) + " : " + current_monkey.ip_addresses[0]
+
+
+    @ring.lru()
+    @staticmethod
+    def get_hostname_by_id(object_id):
+        return Monkey.get_single_monkey_by_id(object_id).hostname
+
+    def set_hostname(self, hostname):
+        """
+        Need this to clear the cache
+        """
+        self.hostname = hostname
+        self.save()
+        Monkey.get_hostname_by_id.delete(self.id)
+
     def get_network_info(self):
         """
         Formats network info from monkey's model

monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py

@@ -16,7 +16,8 @@ def parse_creds(attempt):
         if attempt[key]:
             return '%s ; %s : %s' % (username,
                                      cred['type'],
-                                     cred['output'])
+                                     # TODO Figure out why this is causing an exception with Vakaris
+                                     "cred['output']")
 
 
 def censor_password(password, plain_chars=3, secret_chars=5):

monkey/monkey_island/cc/services/node.py

@@ -24,6 +24,7 @@ class NodeService:
 
         edges = EdgeService.get_displayed_edges_by_to(node_id, for_report)
         accessible_from_nodes = []
+        accessible_from_nodes_hostnames = []
         exploits = []
 
         new_node = {"id": node_id}
@@ -47,15 +48,21 @@ class NodeService:
             new_node["domain_name"] = node["domain_name"]
 
         for edge in edges:
-            accessible_from_nodes.append(NodeService.get_monkey_label(NodeService.get_monkey_by_id(edge["from"])))
+            from_node_id = edge["from"]
+            # from_node_label = NodeService.get_monkey_label(NodeService.get_monkey_by_id(from_node_id))
+            from_node_label = Monkey.get_label_by_id(from_node_id)
+            from_node_hostname = Monkey.get_hostname_by_id(from_node_id)
+            accessible_from_nodes.append(from_node_label)
+            accessible_from_nodes_hostnames.append(from_node_hostname)
             for exploit in edge["exploits"]:
-                exploit["origin"] = NodeService.get_monkey_label(NodeService.get_monkey_by_id(edge["from"]))
+                exploit["origin"] = from_node_label
                 exploits.append(exploit)
 
         exploits.sort(cmp=NodeService._cmp_exploits_by_timestamp)
 
         new_node["exploits"] = exploits
         new_node["accessible_from_nodes"] = accessible_from_nodes
+        new_node["accessible_from_nodes_hostnames"] = accessible_from_nodes_hostnames
         if len(edges) > 0:
             new_node["services"] = edges[-1]["services"]
         else:
@@ -112,6 +119,7 @@ class NodeService:
 
     @staticmethod
     def get_monkey_label(monkey):
+        # todo
         label = monkey["hostname"] + " : " + monkey["ip_addresses"][0]
         ip_addresses = local_ip_addresses()
         if len(set(monkey["ip_addresses"]).intersection(ip_addresses)) > 0:
@@ -349,3 +357,7 @@ class NodeService:
     @staticmethod
     def get_hostname_by_id(node_id):
         return NodeService.get_node_hostname(mongo.db.monkey.find_one({'_id': node_id}, {'hostname': 1}))
+
+    @staticmethod
+    def extract_hostname_from_monkey_label(monkey_label):
+        return monkey_label.split(" ")[0]

monkey/monkey_island/cc/services/reporting/report.py

@@ -23,7 +23,6 @@ from common.network.network_range import NetworkRange
 
 __author__ = "itay.mizeretz"
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -123,19 +122,20 @@ class ReportService:
         formatted_nodes = []
 
         # TODO Figure out and improve
-        nodes = \
-            [NodeService.get_displayed_node_by_id(node['_id'], True) for node in mongo.db.node.find({}, {'_id': 1})] \
-            + [NodeService.get_displayed_node_by_id(monkey['_id'], True) for monkey in
-               mongo.db.monkey.find({}, {'_id': 1})]
+        # 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(
                 {
                     'label': node['label'],
                     'ip_addresses': node['ip_addresses'],
-                    'accessible_from_nodes':
-                        list((x['hostname'] for x in
-                         (NodeService.get_displayed_node_by_id(edge['from'], True)
-                          for edge in EdgeService.get_displayed_edges_by_to(node['id'], True)))),
+                    'accessible_from_nodes': nodes_that_can_access_current_node,
                     'services': node['services'],
                     'domain_name': node['domain_name'],
                     'pba_results': node['pba_results'] if 'pba_results' in node else 'None'
@@ -145,6 +145,15 @@ class ReportService:
 
         return formatted_nodes
 
+    @staticmethod
+    def get_all_displayed_nodes():
+        nodes_without_monkeys = [NodeService.get_displayed_node_by_id(node['_id'], True) for node in
+                                 mongo.db.node.find({}, {'_id': 1})]
+        nodes_with_monkeys = [NodeService.get_displayed_node_by_id(monkey['_id'], True) for monkey in
+                              mongo.db.monkey.find({}, {'_id': 1})]
+        nodes = nodes_without_monkeys + nodes_with_monkeys
+        return nodes
+
     @staticmethod
     def get_exploited():
         exploited = \
@@ -210,8 +219,9 @@ class ReportService:
                 # Pick out all ssh keys not yet included in creds
                 ssh_keys = [{'username': key_pair['name'], 'type': 'Clear SSH private key',
                              'origin': origin} for key_pair in telem['data']['ssh_info']
-                            if key_pair['private_key'] and {'username': key_pair['name'], 'type': 'Clear SSH private key',
-                            'origin': origin} not in creds]
+                            if
+                            key_pair['private_key'] and {'username': key_pair['name'], 'type': 'Clear SSH private key',
+                                                         'origin': origin} not in creds]
                 creds.extend(ssh_keys)
         return creds
 
@@ -700,6 +710,7 @@ class ReportService:
         cross_segment_issues = ReportService.get_cross_segment_issues()
         monkey_latest_modify_time = Monkey.get_latest_modifytime()
 
+        scanned_nodes = ReportService.get_scanned()
         report = \
             {
                 'overview':
@@ -718,7 +729,7 @@ class ReportService:
                     },
                 'glance':
                     {
-                        'scanned': ReportService.get_scanned(),
+                        'scanned': scanned_nodes,
                         'exploited': ReportService.get_exploited(),
                         'stolen_creds': ReportService.get_stolen_creds(),
                         'azure_passwords': ReportService.get_azure_creds(),
@@ -752,7 +763,6 @@ class ReportService:
         report_as_json = json_util.dumps(report_dict).replace('.', ',,,')
         return json_util.loads(report_as_json)
 
-
     @staticmethod
     def is_latest_report_exists():
         """

monkey/monkey_island/cc/services/telemetry/processing/system_info.py

@@ -1,4 +1,5 @@
 from monkey_island.cc.database import mongo
+from monkey_island.cc.models import Monkey
 from monkey_island.cc.services import mimikatz_utils
 from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.services.config import ConfigService
@@ -12,6 +13,7 @@ def process_system_info_telemetry(telemetry_json):
     process_credential_info(telemetry_json)
     process_mimikatz_and_wmi_info(telemetry_json)
     process_aws_data(telemetry_json)
+    update_db_with_new_hostname(telemetry_json)
     test_antivirus_existence(telemetry_json)
 
 
@@ -97,3 +99,7 @@ def process_aws_data(telemetry_json):
             monkey_id = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid']).get('_id')
             mongo.db.monkey.update_one({'_id': monkey_id},
                                        {'$set': {'aws_instance_id': telemetry_json['data']['aws']['instance_id']}})
+
+
+def update_db_with_new_hostname(telemetry_json):
+    Monkey.get_single_monkey_by_id(telemetry_json['_id']).set_hostname(telemetry_json['data']['hostname'])

monkey/monkey_island/requirements.txt

@@ -26,4 +26,5 @@ mongoengine
 mongomock
 requests
 dpath
-backports
+backports.functools-lru-cache
+ring