From 1060c004bd2799bbcddb79fed73cb45d5ade0c6c Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Sun, 29 Sep 2019 15:54:24 +0300 Subject: [PATCH] Started improving and researching the performence issues - still in progress... --- monkey/monkey_island/cc/models/monkey.py | 25 +++++++++++++ .../technique_report_tools.py | 3 +- monkey/monkey_island/cc/services/node.py | 16 +++++++-- .../cc/services/reporting/report.py | 36 ++++++++++++------- .../telemetry/processing/system_info.py | 6 ++++ monkey/monkey_island/requirements.txt | 3 +- 6 files changed, 72 insertions(+), 17 deletions(-) diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index c0eeb20b3..4726dfdf6 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/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 diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py b/monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py index 05cef3684..390852d27 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py +++ b/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): diff --git a/monkey/monkey_island/cc/services/node.py b/monkey/monkey_island/cc/services/node.py index 2c75d7187..008516c15 100644 --- a/monkey/monkey_island/cc/services/node.py +++ b/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] diff --git a/monkey/monkey_island/cc/services/reporting/report.py b/monkey/monkey_island/cc/services/reporting/report.py index 84633cc81..d792e7bfd 100644 --- a/monkey/monkey_island/cc/services/reporting/report.py +++ b/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(): """ diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info.py index ebf11c219..a970c0cd4 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/system_info.py +++ b/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']) diff --git a/monkey/monkey_island/requirements.txt b/monkey/monkey_island/requirements.txt index 92a118b07..d9fa69f0a 100644 --- a/monkey/monkey_island/requirements.txt +++ b/monkey/monkey_island/requirements.txt @@ -26,4 +26,5 @@ mongoengine mongomock requests dpath -backports +backports.functools-lru-cache +ring