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.

This commit is contained in:
Shay Nehmad 2019-10-02 11:08:49 +03:00
parent 264e7407fb
commit e3b93f1506
2 changed files with 26 additions and 25 deletions

View File

@ -2,6 +2,7 @@ from bson import ObjectId
from monkey_island.cc.database import mongo from monkey_island.cc.database import mongo
import monkey_island.cc.services.node import monkey_island.cc.services.node
from monkey_island.cc.models import Monkey
__author__ = "itay.mizeretz" __author__ = "itay.mizeretz"
@ -141,15 +142,18 @@ class EdgeService:
@staticmethod @staticmethod
def get_edge_label(edge): def get_edge_label(edge):
NodeService = monkey_island.cc.services.node.NodeService NodeService = monkey_island.cc.services.node.NodeService
from_label = NodeService.get_monkey_label(NodeService.get_monkey_by_id(edge["from"])) from_id = edge["from"]
if edge["to"] == ObjectId("000000000000000000000000"): to_id = edge["to"]
from_label = Monkey.get_label_by_id(from_id)
if to_id == ObjectId("000000000000000000000000"):
to_label = 'MonkeyIsland' to_label = 'MonkeyIsland'
else: else:
to_id = NodeService.get_monkey_by_id(edge["to"]) if NodeService.get_monkey_by_id(to_id) is None:
if to_id is None: to_label = NodeService.get_node_label(NodeService.get_node_by_id(to_id))
to_label = NodeService.get_node_label(NodeService.get_node_by_id(edge["to"]))
else: else:
to_label = NodeService.get_monkey_label(to_id) to_label = Monkey.get_label_by_id(to_id)
RIGHT_ARROW = u"\u2192" RIGHT_ARROW = u"\u2192"
return "%s %s %s" % (from_label, RIGHT_ARROW, to_label) return "%s %s %s" % (from_label, RIGHT_ARROW, to_label)

View File

@ -118,17 +118,10 @@ class ReportService:
@staticmethod @staticmethod
def get_scanned(): def get_scanned():
formatted_nodes = [] 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() nodes = ReportService.get_all_displayed_nodes()
print("2")
# for each node (n*...
for node in nodes: for node in nodes:
nodes_that_can_access_current_node = node['accessible_from_nodes_hostnames'] nodes_that_can_access_current_node = node['accessible_from_nodes_hostnames']
formatted_nodes.append( formatted_nodes.append(
@ -156,23 +149,26 @@ class ReportService:
@staticmethod @staticmethod
def get_exploited(): def get_exploited():
exploited = \ exploited_with_monkeys = \
[NodeService.get_displayed_node_by_id(monkey['_id'], True) for monkey in [NodeService.get_displayed_node_by_id(monkey['_id'], True) for monkey in
mongo.db.monkey.find({}, {'_id': 1}) mongo.db.monkey.find({}, {'_id': 1}) if
if not NodeService.get_monkey_manual_run(NodeService.get_monkey_by_id(monkey['_id']))] \ 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})] 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 = [ exploited = [
{ {
'label': monkey['label'], 'label': exploited_node['label'],
'ip_addresses': monkey['ip_addresses'], 'ip_addresses': exploited_node['ip_addresses'],
'domain_name': monkey['domain_name'], 'domain_name': exploited_node['domain_name'],
'exploits': list(set( 'exploits': list(set(
[ReportService.EXPLOIT_DISPLAY_DICT[exploit['exploiter']] for exploit in monkey['exploits'] if [ReportService.EXPLOIT_DISPLAY_DICT[exploit['exploiter']] for exploit in exploited_node['exploits']
exploit['result']])) if exploit['result']]))
} }
for monkey in exploited] for exploited_node in exploited]
logger.info('Exploited nodes generated for reporting') logger.info('Exploited nodes generated for reporting')
@ -711,6 +707,7 @@ class ReportService:
monkey_latest_modify_time = Monkey.get_latest_modifytime() monkey_latest_modify_time = Monkey.get_latest_modifytime()
scanned_nodes = ReportService.get_scanned() scanned_nodes = ReportService.get_scanned()
exploited_nodes = ReportService.get_exploited()
report = \ report = \
{ {
'overview': 'overview':
@ -730,7 +727,7 @@ class ReportService:
'glance': 'glance':
{ {
'scanned': scanned_nodes, 'scanned': scanned_nodes,
'exploited': ReportService.get_exploited(), 'exploited': exploited_nodes,
'stolen_creds': ReportService.get_stolen_creds(), 'stolen_creds': ReportService.get_stolen_creds(),
'azure_passwords': ReportService.get_azure_creds(), 'azure_passwords': ReportService.get_azure_creds(),
'ssh_keys': ReportService.get_ssh_keys(), 'ssh_keys': ReportService.get_ssh_keys(),