Refactored edge from mongo queries to DAL

monkey/monkey_island/cc/models/edge.py

@@ -0,0 +1,16 @@
+from mongoengine import Document, ObjectIdField, ListField, DynamicField, BooleanField, StringField
+
+
+class Edge(Document):
+    # SCHEMA
+    src_node_id = ObjectIdField(required=True)
+    dst_node_id = ObjectIdField(required=True)
+    scans = ListField(DynamicField(), default=[])
+    exploits = ListField(DynamicField(), default=[])
+    tunnel = BooleanField(default=False)
+    exploited = BooleanField(default=False)
+    src_label = StringField()
+    dst_label = StringField()
+    group = StringField()
+    domain_name = StringField()
+    ip_address = StringField()

monkey/monkey_island/cc/resources/edge.py

@@ -1,7 +1,7 @@
-from flask import request
 import flask_restful
+from flask import request
 
-from monkey_island.cc.services.edge import EdgeService
+from monkey_island.cc.services.edge.displayed_edge import DisplayedEdgeService
 
 __author__ = 'Barak'
 
@@ -9,7 +9,8 @@ __author__ = 'Barak'
 class Edge(flask_restful.Resource):
     def get(self):
         edge_id = request.args.get('id')
+        displayed_edge = DisplayedEdgeService.get_displayed_edge_by_id(edge_id)
         if edge_id:
-            return {"edge": EdgeService.get_displayed_edge_by_id(edge_id)}
+            return {"edge": displayed_edge}
 
         return {}

monkey/monkey_island/cc/resources/monkey.py

@@ -3,6 +3,8 @@ from datetime import datetime
 
 import dateutil.parser
 import flask_restful
+
+from monkey_island.cc.models.edge import Edge
 from monkey_island.cc.resources.test.utils.telem_store import TestTelemStore
 from flask import request
 
@@ -129,8 +131,9 @@ class Monkey(flask_restful.Resource):
 
         if existing_node:
             node_id = existing_node["_id"]
-            for edge in mongo.db.edge.find({"to": node_id}):
+            for edge in Edge.objects(dst_node_id=node_id):
-                mongo.db.edge.update({"_id": edge["_id"]}, {"$set": {"to": new_monkey_id}})
+                edge.dst_node_id = new_monkey_id
+                edge.save()
             for creds in existing_node['creds']:
                 NodeService.add_credentials_to_monkey(new_monkey_id, creds)
             mongo.db.node.remove({"_id": node_id})

monkey/monkey_island/cc/resources/netmap.py

@@ -1,9 +1,8 @@
 import flask_restful
 
 from monkey_island.cc.auth import jwt_required
-from monkey_island.cc.services.edge import EdgeService
+from monkey_island.cc.services.netmap.net_edge import NetEdgeService
-from monkey_island.cc.services.node import NodeService
+from monkey_island.cc.services.netmap.net_node import NetNodeService
-from monkey_island.cc.database import mongo
 
 __author__ = 'Barak'
 
@@ -11,19 +10,11 @@ __author__ = 'Barak'
 class NetMap(flask_restful.Resource):
     @jwt_required()
     def get(self, **kw):
-        monkeys = [NodeService.monkey_to_net_node(x) for x in mongo.db.monkey.find({})]
+        net_nodes = NetNodeService.get_all_net_nodes()
-        nodes = [NodeService.node_to_net_node(x) for x in mongo.db.node.find({})]
+        net_edges = NetEdgeService.get_all_net_edges()
-        edges = [EdgeService.edge_to_net_edge(x) for x in mongo.db.edge.find({})]
-
-        if NodeService.get_monkey_island_monkey() is None:
-            monkey_island = [NodeService.get_monkey_island_pseudo_net_node()]
-            edges += EdgeService.get_monkey_island_pseudo_edges()
-        else:
-            monkey_island = []
-            edges += EdgeService.get_infected_monkey_island_pseudo_edges()
 
         return \
             {
-                "nodes": monkeys + nodes + monkey_island,
+                "nodes": net_nodes,
-                "edges": edges
+                "edges": net_edges
             }

monkey/monkey_island/cc/services/edge.py

@@ -1,173 +0,0 @@
-from bson import ObjectId
-
-from monkey_island.cc.database import mongo
-import monkey_island.cc.services.node
-from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError
-
-__author__ = "itay.mizeretz"
-
-
-class EdgeService:
-    def __init__(self):
-        pass
-
-    @staticmethod
-    def get_displayed_edge_by_id(edge_id, for_report=False):
-        edge = mongo.db.edge.find({"_id": ObjectId(edge_id)})[0]
-        return EdgeService.edge_to_displayed_edge(edge, for_report)
-
-    @staticmethod
-    def get_displayed_edges_by_to(to, for_report=False):
-        edges = mongo.db.edge.find({"to": ObjectId(to)})
-        return [EdgeService.edge_to_displayed_edge(edge, for_report) for edge in edges]
-
-    @staticmethod
-    def edge_to_displayed_edge(edge, for_report=False):
-        services = []
-        os = {}
-
-        if len(edge["scans"]) > 0:
-            services = EdgeService.services_to_displayed_services(edge["scans"][-1]["data"]["services"], for_report)
-            os = edge["scans"][-1]["data"]["os"]
-
-        displayed_edge = EdgeService.edge_to_net_edge(edge)
-
-        displayed_edge["ip_address"] = edge["ip_address"]
-        displayed_edge["services"] = services
-        displayed_edge["os"] = os
-        displayed_edge["exploits"] = edge['exploits']
-        displayed_edge["_label"] = EdgeService.get_edge_label(displayed_edge)
-        return displayed_edge
-
-    @staticmethod
-    def insert_edge(from_id, to_id):
-        edge_insert_result = mongo.db.edge.insert_one(
-            {
-                "from": from_id,
-                "to": to_id,
-                "scans": [],
-                "exploits": [],
-                "tunnel": False,
-                "exploited": False,
-                "src_label": EdgeService.get_label_for_endpoint(from_id),
-                "dst_label": EdgeService.get_label_for_endpoint(to_id)
-            })
-        return mongo.db.edge.find_one({"_id": edge_insert_result.inserted_id})
-
-    @staticmethod
-    def get_or_create_edge(edge_from, edge_to):
-        tunnel_edge = mongo.db.edge.find_one({"from": edge_from, "to": edge_to})
-        if tunnel_edge is None:
-            tunnel_edge = EdgeService.insert_edge(edge_from, edge_to)
-
-        return tunnel_edge
-
-    @staticmethod
-    def generate_pseudo_edge(edge_id, edge_from, edge_to):
-        edge = \
-            {
-                "id": edge_id,
-                "from": edge_from,
-                "to": edge_to,
-                "group": "island",
-                "src_label": EdgeService.get_label_for_endpoint(edge_from),
-                "dst_label": EdgeService.get_label_for_endpoint(edge_to)
-            }
-        edge["_label"] = EdgeService.get_edge_label(edge)
-        return edge
-
-    @staticmethod
-    def get_monkey_island_pseudo_edges():
-        edges = []
-        monkey_ids = [x["_id"] for x in mongo.db.monkey.find({}) if "tunnel" not in x]
-        # We're using fake ids because the frontend graph module requires unique ids.
-        # Collision with real id is improbable.
-        count = 0
-        for monkey_id in monkey_ids:
-            count += 1
-            edges.append(EdgeService.generate_pseudo_edge(
-                ObjectId(hex(count)[2:].zfill(24)), monkey_id, ObjectId("000000000000000000000000")))
-
-        return edges
-
-    @staticmethod
-    def get_infected_monkey_island_pseudo_edges():
-        monkey = monkey_island.cc.services.node.NodeService.get_monkey_island_monkey()
-        existing_ids = [x["from"] for x in mongo.db.edge.find({"to": monkey["_id"]})]
-        monkey_ids = [x["_id"] for x in mongo.db.monkey.find({})
-                      if ("tunnel" not in x) and (x["_id"] not in existing_ids) and (x["_id"] != monkey["_id"])]
-        edges = []
-
-        # We're using fake ids because the frontend graph module requires unique ids.
-        # Collision with real id is improbable.
-        count = 0
-        for monkey_id in monkey_ids:
-            count += 1
-            edges.append(EdgeService.generate_pseudo_edge(
-                ObjectId(hex(count)[2:].zfill(24)), monkey_id, monkey["_id"]))
-
-        return edges
-
-    @staticmethod
-    def services_to_displayed_services(services, for_report=False):
-        if for_report:
-            return [x for x in services]
-        else:
-            return [x + ": " + (services[x]['name'] if 'name' in services[x] else 'unknown') for x in services]
-
-    @staticmethod
-    def edge_to_net_edge(edge):
-        return \
-            {
-                "id": edge["_id"],
-                "from": edge["from"],
-                "to": edge["to"],
-                "group": EdgeService.get_edge_group(edge),
-                "src_label": edge["src_label"],
-                "dst_label": edge["dst_label"]
-            }
-
-    @staticmethod
-    def get_edge_group(edge):
-        if edge.get("exploited"):
-            return "exploited"
-        if edge.get("tunnel"):
-            return "tunnel"
-        if (len(edge.get("scans", [])) > 0) or (len(edge.get("exploits", [])) > 0):
-            return "scan"
-        return "empty"
-
-    @staticmethod
-    def set_edge_exploited(edge):
-        mongo.db.edge.update(
-            {"_id": edge["_id"]},
-            {"$set": {"exploited": True}}
-        )
-        monkey_island.cc.services.node.NodeService.set_node_exploited(edge["to"])
-
-    @staticmethod
-    def get_edge_label(edge):
-        return "%s %s %s" % (edge['src_label'], RIGHT_ARROW, edge['dst_label'])
-
-    @staticmethod
-    def get_label_for_endpoint(endpoint_id):
-        node_service = monkey_island.cc.services.node.NodeService
-        if endpoint_id == ObjectId("000000000000000000000000"):
-            return 'MonkeyIsland'
-        if Monkey.is_monkey(endpoint_id):
-            return Monkey.get_label_by_id(endpoint_id)
-        else:
-            return node_service.get_node_label(node_service.get_node_by_id(endpoint_id))
-
-    @staticmethod
-    def update_label_by_endpoint(edge, endpoint_id):
-        label = EdgeService.get_label_for_endpoint(endpoint_id)
-        if endpoint_id == edge["to"]:
-            mongo_field = {"dst_label": label}
-        else:
-            mongo_field = {"src_label": label}
-        mongo.db.edge.update({"_id": edge["_id"]},
-                             {"$set": mongo_field})
-
-
-RIGHT_ARROW = "\u2192"

monkey/monkey_island/cc/services/edge/displayed_edge.py

@@ -0,0 +1,100 @@
+from copy import deepcopy
+
+from bson import ObjectId
+
+from monkey_island.cc.database import mongo
+from monkey_island.cc.models import Monkey
+from monkey_island.cc.models.edge import Edge
+from monkey_island.cc.services.edge.edge import EdgeService
+
+__author__ = "itay.mizeretz"
+
+
+class DisplayedEdgeService:
+
+    @staticmethod
+    def get_displayed_edges_by_dst(dst_id, for_report=False):
+        edges = Edge.objects(dst_node_id=ObjectId(dst_id))
+        return [DisplayedEdgeService.edge_to_displayed_edge(edge, for_report) for edge in edges]
+
+    @staticmethod
+    def get_displayed_edge_by_id(edge_id, for_report=False):
+        edge = Edge.objects.get(id=edge_id)
+        displayed_edge = DisplayedEdgeService.edge_to_displayed_edge(edge, for_report)
+        return displayed_edge
+
+    @staticmethod
+    def edge_to_displayed_edge(edge: Edge, for_report=False):
+        services = []
+        os = {}
+
+        if len(edge.scans) > 0:
+            services = DisplayedEdgeService.services_to_displayed_services(edge.scans[-1]["data"]["services"],
+                                                                           for_report)
+            os = edge.scans[-1]["data"]["os"]
+
+        displayed_edge = DisplayedEdgeService.edge_to_net_edge(edge)
+
+        displayed_edge["ip_address"] = edge['ip_address']
+        displayed_edge["services"] = services
+        displayed_edge["os"] = os
+        # we need to deepcopy all mutable edge properties, because weak-reference link is made otherwise,
+        # which is destroyed after method is exited and causes an error later.
+        displayed_edge["exploits"] = deepcopy(edge['exploits'])
+        displayed_edge["_label"] = EdgeService.get_edge_label(displayed_edge)
+        return displayed_edge
+
+    @staticmethod
+    def generate_pseudo_edge(edge_id, src_node_id, dst_node_id, src_label, dst_label):
+        edge = \
+            {
+                "id": edge_id,
+                "from": src_node_id,
+                "to": dst_node_id,
+                "group": "island",
+                "src_label": src_label,
+                "dst_label": dst_label
+            }
+        edge["_label"] = EdgeService.get_edge_label(edge)
+        return edge
+
+    @staticmethod
+    def get_infected_monkey_island_pseudo_edges(monkey_island_monkey):
+        existing_ids = [x.src_node_id for x in Edge.objects(dst_node_id=monkey_island_monkey["_id"])]
+        monkey_ids = [x["_id"] for x in mongo.db.monkey.find({})
+                      if ("tunnel" not in x) and
+                      (x["_id"] not in existing_ids) and
+                      (x["_id"] != monkey_island_monkey["_id"])]
+        edges = []
+
+        # We're using fake ids because the frontend graph module requires unique ids.
+        # Collision with real id is improbable.
+        count = 0
+        for monkey_id in monkey_ids:
+            count += 1
+            edges.append(DisplayedEdgeService.generate_pseudo_edge(
+                ObjectId(hex(count)[2:].zfill(24)), monkey_id, monkey_island_monkey["_id"]))
+
+        return edges
+
+    @staticmethod
+    def services_to_displayed_services(services, for_report=False):
+        if for_report:
+            return [x for x in services]
+        else:
+            return [x + ": " + (services[x]['name'] if 'name' in services[x] else 'unknown') for x in services]
+
+    @staticmethod
+    def edge_to_net_edge(edge: Edge):
+        return \
+            {
+                "id": edge.id,
+                "from": edge.src_node_id,
+                "to": edge.dst_node_id,
+                "group": EdgeService.get_edge_group(edge),
+                "src_label": edge.src_label,
+                "dst_label": edge.dst_label
+            }
+
+
+RIGHT_ARROW = "\u2192"

monkey/monkey_island/cc/services/edge/edge.py

@@ -0,0 +1,78 @@
+import copy
+from typing import Dict
+
+from bson import ObjectId
+from mongoengine import DoesNotExist
+
+from monkey_island.cc.models.edge import Edge
+
+RIGHT_ARROW = "\u2192"
+
+
+class EdgeService:
+
+    @staticmethod
+    def get_or_create_edge(src_node_id, dst_node_id, src_label, dst_label):
+        edge = False
+        try:
+            edge = Edge.objects.get(src_node_id=src_node_id, dst_node_id=dst_node_id)
+        except DoesNotExist:
+            edge = Edge(src_node_id=src_node_id, dst_node_id=dst_node_id)
+        finally:
+            if edge:
+                edge.src_label = src_label
+                edge.dst_label = dst_label
+                edge.save()
+        return edge
+
+    @staticmethod
+    def update_label(edge: Edge, node_id: ObjectId, label: str):
+        if edge.src_node_id == node_id:
+            edge.src_label = label
+        elif edge.dst_node_id == node_id:
+            edge.dst_label = label
+        else:
+            raise DoesNotExist("Node id provided does not match with any endpoint of an edge provided.")
+        edge.save()
+        pass
+
+    @staticmethod
+    def update_based_on_scan_telemetry(edge: Edge, telemetry: Dict):
+        machine_info = copy.deepcopy(telemetry['data']['machine'])
+        new_scan = \
+            {
+                "timestamp": telemetry["timestamp"],
+                "data": machine_info
+            }
+        ip_address = machine_info.pop("ip_addr")
+        domain_name = machine_info.pop("domain_name")
+        edge.scans.append(new_scan)
+        edge.ip_address = ip_address
+        edge.domain_name = domain_name
+        edge.save()
+
+    @staticmethod
+    def update_based_on_exploit(edge: Edge, exploit: Dict):
+        edge.exploits.append(exploit)
+        edge.save()
+        if exploit['result']:
+            EdgeService.set_edge_exploited(edge)
+
+    @staticmethod
+    def get_edge_group(edge: Edge):
+        if edge.exploited:
+            return "exploited"
+        if edge.tunnel:
+            return "tunnel"
+        if edge.scans or edge.exploits:
+            return "scan"
+        return "empty"
+
+    @staticmethod
+    def set_edge_exploited(edge: Edge):
+        edge.exploited = True
+        edge.save()
+
+    @staticmethod
+    def get_edge_label(edge):
+        return "%s %s %s" % (edge['src_label'], RIGHT_ARROW, edge['dst_label'])

monkey/monkey_island/cc/services/edge/test_displayed_edge.py

@@ -0,0 +1,97 @@
+from bson import ObjectId
+
+from monkey_island.cc.models.edge import Edge
+from monkey_island.cc.services.edge.displayed_edge import DisplayedEdgeService
+from monkey_island.cc.services.edge.edge import EdgeService, RIGHT_ARROW
+from monkey_island.cc.testing.IslandTestCase import IslandTestCase
+
+SCAN_DATA_MOCK = [{
+    "timestamp": "2020-05-27T14:59:28.944Z",
+    "data": {
+        "os": {
+            "type": "linux",
+            "version": "Ubuntu-4ubuntu2.8"
+        },
+        "services": {
+            "tcp-8088": {
+                "display_name": "unknown(TCP)",
+                "port": 8088
+            },
+            "tcp-22": {
+                "display_name": "SSH",
+                "port": 22,
+                "banner": "SSH-2.0-OpenSSH_7.2p2 Ubuntu-4ubuntu2.8\r\n",
+                "name": "ssh"
+            }
+        },
+        "monkey_exe": None,
+        "default_tunnel": None,
+        "default_server": None
+    }
+}]
+
+EXPLOIT_DATA_MOCK = [{
+    "result": True,
+    "exploiter": "ElasticGroovyExploiter",
+    "info": {
+        "display_name": "Elastic search",
+        "started": "2020-05-11T08:59:38.105Z",
+        "finished": "2020-05-11T08:59:38.106Z",
+        "vulnerable_urls": [],
+        "vulnerable_ports": [],
+        "executed_cmds": []
+    },
+    "attempts": [],
+    "timestamp": "2020-05-27T14:59:29.048Z"
+}]
+
+
+class TestDisplayedEdgeService(IslandTestCase):
+    def test_get_displayed_edges_by_to(self):
+        self.clean_edge_db()
+
+        dst_id = ObjectId()
+
+        src_id = ObjectId()
+        EdgeService.get_or_create_edge(src_id, dst_id, "Ubuntu-4ubuntu2.8", "Ubuntu-4ubuntu2.8")
+
+        src_id2 = ObjectId()
+        EdgeService.get_or_create_edge(src_id2, dst_id, "Ubuntu-4ubuntu3.2", "Ubuntu-4ubuntu2.8")
+
+        displayed_edges = DisplayedEdgeService.get_displayed_edges_by_dst(dst_id)
+        self.assertEqual(len(displayed_edges), 2)
+
+    def test_edge_to_displayed_edge(self):
+        src_node_id = ObjectId()
+        dst_node_id = ObjectId()
+        edge = Edge(src_node_id=src_node_id,
+                    dst_node_id=dst_node_id,
+                    scans=SCAN_DATA_MOCK,
+                    exploits=EXPLOIT_DATA_MOCK,
+                    exploited=True,
+                    domain_name=None,
+                    ip_address="10.2.2.2",
+                    dst_label="Ubuntu-4ubuntu2.8",
+                    src_label="Ubuntu-4ubuntu3.2")
+
+        displayed_edge = DisplayedEdgeService.edge_to_displayed_edge(edge)
+
+        self.assertEqual(displayed_edge['to'], dst_node_id)
+        self.assertEqual(displayed_edge['from'], src_node_id)
+        self.assertEqual(displayed_edge['ip_address'], "10.2.2.2")
+        self.assertListEqual(displayed_edge['services'], ["tcp-8088: unknown", "tcp-22: ssh"])
+        self.assertEqual(displayed_edge['os'], {"type": "linux",
+                                                "version": "Ubuntu-4ubuntu2.8"})
+        self.assertEqual(displayed_edge['exploits'], EXPLOIT_DATA_MOCK)
+        self.assertEqual(displayed_edge['_label'], "Ubuntu-4ubuntu3.2 " + RIGHT_ARROW + " Ubuntu-4ubuntu2.8")
+        self.assertEqual(displayed_edge['group'], "exploited")
+        return displayed_edge
+
+    def test_services_to_displayed_services(self):
+        services1 = DisplayedEdgeService.services_to_displayed_services(SCAN_DATA_MOCK[-1]["data"]["services"],
+                                                                        True)
+        self.assertEqual(services1, ["tcp-8088", "tcp-22"])
+
+        services2 = DisplayedEdgeService.services_to_displayed_services(SCAN_DATA_MOCK[-1]["data"]["services"],
+                                                                        False)
+        self.assertEqual(services2, ["tcp-8088: unknown", "tcp-22: ssh"])

monkey/monkey_island/cc/services/edge/test_edge.py

@@ -0,0 +1,60 @@
+import logging
+
+from mongomock import ObjectId
+
+from monkey_island.cc.models.edge import Edge
+from monkey_island.cc.services.edge.edge import EdgeService
+from monkey_island.cc.testing.IslandTestCase import IslandTestCase
+
+logger = logging.getLogger(__name__)
+
+
+class TestEdgeService(IslandTestCase):
+    """
+    Make sure to set server environment to `testing` in server_config.json!
+    Otherwise this will mess up your mongo instance and won't work.
+
+    Also, the working directory needs to be the working directory from which you usually run the island so the
+    server_config.json file is found and loaded.
+    """
+
+    def test_get_or_create_edge(self):
+        self.fail_if_not_testing_env()
+        self.clean_edge_db()
+
+        src_id = ObjectId()
+        dst_id = ObjectId()
+
+        test_edge1 = EdgeService.get_or_create_edge(src_id, dst_id, "Mock label 1", "Mock label 2")
+        self.assertEqual(test_edge1.src_node_id, src_id)
+        self.assertEqual(test_edge1.dst_node_id, dst_id)
+        self.assertFalse(test_edge1.exploited)
+        self.assertFalse(test_edge1.tunnel)
+        self.assertListEqual(test_edge1.scans, [])
+        self.assertListEqual(test_edge1.exploits, [])
+        self.assertEqual(test_edge1.src_label, "Mock label 1")
+        self.assertEqual(test_edge1.dst_label, "Mock label 2")
+        self.assertIsNone(test_edge1.group)
+        self.assertIsNone(test_edge1.domain_name)
+        self.assertIsNone(test_edge1.ip_address)
+
+        EdgeService.get_or_create_edge(src_id, dst_id, "Mock label 1", "Mock label 2")
+        self.assertEqual(len(Edge.objects()), 1)
+
+    def test_get_edge_group(self):
+        edge = Edge(src_node_id=ObjectId(),
+                    dst_node_id=ObjectId(),
+                    exploited=True)
+        self.assertEqual("exploited", EdgeService.get_edge_group(edge))
+
+        edge.exploited = False
+        edge.tunnel = True
+        self.assertEqual("tunnel", EdgeService.get_edge_group(edge))
+
+        edge.tunnel = False
+        edge.exploits.append(["mock_exploit_data"])
+        self.assertEqual("scan", EdgeService.get_edge_group(edge))
+
+        edge.exploits = []
+        edge.scans = []
+        self.assertEqual("empty", EdgeService.get_edge_group(edge))

monkey/monkey_island/cc/services/netmap/net_edge.py

@@ -0,0 +1,68 @@
+from bson import ObjectId
+
+from monkey_island.cc.models import Monkey
+from monkey_island.cc.models.edge import Edge
+from monkey_island.cc.services.edge.displayed_edge import DisplayedEdgeService
+from monkey_island.cc.services.node import NodeService
+
+
+class NetEdgeService:
+
+    @staticmethod
+    def get_all_net_edges():
+        edges = NetEdgeService._get_standard_net_edges()
+        if NodeService.get_monkey_island_monkey() is None:
+            edges += NetEdgeService._get_uninfected_island_net_edges()
+        else:
+            monkey_island_monkey = NodeService.get_monkey_island_monkey()
+            edges += NetEdgeService._get_infected_island_net_edges(monkey_island_monkey)
+        return edges
+
+    @staticmethod
+    def _get_standard_net_edges():
+        return [DisplayedEdgeService.edge_to_net_edge(x) for x in Edge.objects()]
+
+    @staticmethod
+    def _get_uninfected_island_net_edges():
+        edges = []
+        monkey_ids = [x.id for x in Monkey.objects() if "tunnel" not in x]
+        count = 0
+        for monkey_id in monkey_ids:
+            count += 1
+            # generating fake ID, because front end requires unique ID's for each edge. Collision improbable
+            fake_id = ObjectId(hex(count)[2:].zfill(24))
+            island_id = ObjectId("000000000000000000000000")
+            monkey_label = NodeService.get_label_for_endpoint(monkey_id)
+            island_label = NodeService.get_label_for_endpoint(island_id)
+            island_pseudo_edge = DisplayedEdgeService.generate_pseudo_edge(edge_id=fake_id,
+                                                                           src_node_id=monkey_id,
+                                                                           dst_node_id=island_id,
+                                                                           src_label=monkey_label,
+                                                                           dst_label=island_label)
+            edges.append(island_pseudo_edge)
+        return edges
+
+    @staticmethod
+    def _get_infected_island_net_edges(monkey_island_monkey):
+        existing_ids = [x.src_node_id for x in Edge.objects(dst_node_id=monkey_island_monkey["_id"])]
+        monkey_ids = [x.id for x in Monkey.objects()
+                      if ("tunnel" not in x) and
+                      (x.id not in existing_ids) and
+                      (x.id != monkey_island_monkey["_id"])]
+        edges = []
+
+        count = 0
+        for monkey_id in monkey_ids:
+            count += 1
+            # generating fake ID, because front end requires unique ID's for each edge. Collision improbable
+            fake_id = ObjectId(hex(count)[2:].zfill(24))
+            src_label = NodeService.get_label_for_endpoint(monkey_id)
+            dst_label = NodeService.get_label_for_endpoint(monkey_island_monkey["_id"])
+            edge = DisplayedEdgeService.generate_pseudo_edge(edge_id=fake_id,
+                                                             src_node_id=monkey_id,
+                                                             dst_node_id=monkey_island_monkey["_id"],
+                                                             src_label=src_label,
+                                                             dst_label=dst_label)
+            edges.append(edge)
+
+        return edges

monkey/monkey_island/cc/services/netmap/net_node.py

@@ -0,0 +1,23 @@
+from monkey_island.cc.database import mongo
+from monkey_island.cc.services.node import NodeService
+
+
+class NetNodeService:
+
+    @staticmethod
+    def get_all_net_nodes():
+        monkeys = NetNodeService._get_monkey_net_nodes()
+        nodes = NetNodeService._get_standard_net_nodes()
+        if NodeService.get_monkey_island_monkey() is None:
+            monkey_island = [NodeService.get_monkey_island_pseudo_net_node()]
+        else:
+            monkey_island = []
+        return monkeys + nodes + monkey_island
+
+    @staticmethod
+    def _get_monkey_net_nodes():
+        return [NodeService.monkey_to_net_node(x) for x in mongo.db.monkey.find({})]
+
+    @staticmethod
+    def _get_standard_net_nodes():
+        return [NodeService.node_to_net_node(x) for x in mongo.db.node.find({})]

monkey/monkey_island/cc/services/node.py

@@ -3,13 +3,16 @@ from typing import Dict
 import socket
 
 from bson import ObjectId
+from mongoengine import DoesNotExist
 
 import monkey_island.cc.services.log
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
-from monkey_island.cc.services.edge import EdgeService
+from monkey_island.cc.models.edge import Edge
+from monkey_island.cc.services.edge.displayed_edge import DisplayedEdgeService
 from monkey_island.cc.network_utils import local_ip_addresses, is_local_ips
 from monkey_island.cc import models
+from monkey_island.cc.services.edge.edge import EdgeService
 from monkey_island.cc.services.utils.node_states import NodeStates
 
 __author__ = "itay.mizeretz"
@@ -48,18 +51,18 @@ class NodeService:
         accessible_from_nodes_hostnames = []
         exploits = []
 
-        edges = EdgeService.get_displayed_edges_by_to(node_id, for_report)
+        edges = DisplayedEdgeService.get_displayed_edges_by_dst(node_id, for_report)
 
         for edge in edges:
-            from_node_id = edge["from"]
+            from_node_id = edge['from']
             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 edge_exploit in edge["exploits"]:
+            for edge_exploit in edge['exploits']:
-                edge_exploit["origin"] = from_node_label
+                edge_exploit['origin'] = from_node_label
                 exploits.append(edge_exploit)
 
         exploits = sorted(exploits, key=lambda exploit: exploit['timestamp'])
@@ -186,23 +189,30 @@ class NodeService:
             {'$unset': {'tunnel': ''}},
             upsert=False)
 
-        mongo.db.edge.update(
+        try:
-            {"from": monkey_id, 'tunnel': True},
+            edge = Edge.objects.get(src_node_id=monkey_id, tunnel=True)
-            {'$set': {'tunnel': False}},
+            edge.tunnel = False
-            upsert=False)
+            edge.save()
+        except DoesNotExist:
+            pass
 
     @staticmethod
     def set_monkey_tunnel(monkey_id, tunnel_host_ip):
         tunnel_host_id = NodeService.get_monkey_by_ip(tunnel_host_ip)["_id"]
         NodeService.unset_all_monkey_tunnels(monkey_id)
         mongo.db.monkey.update(
-            {"_id": monkey_id},
+            {'_id': monkey_id},
             {'$set': {'tunnel': tunnel_host_id}},
             upsert=False)
-        tunnel_edge = EdgeService.get_or_create_edge(monkey_id, tunnel_host_id)
+        monkey_label = NodeService.get_label_for_endpoint(monkey_id)
-        mongo.db.edge.update({"_id": tunnel_edge["_id"]},
+        tunnel_host_label = NodeService.get_label_for_endpoint(tunnel_host_id)
-                             {'$set': {'tunnel': True, 'ip_address': tunnel_host_ip}},
+        tunnel_edge = EdgeService.get_or_create_edge(src_node_id=monkey_id,
-                             upsert=False)
+                                                     dst_node_id=tunnel_host_id,
+                                                     src_label=monkey_label,
+                                                     dst_label=tunnel_host_label)
+        tunnel_edge.tunnel = True
+        tunnel_edge.ip_address = tunnel_host_ip
+        tunnel_edge.save()
 
     @staticmethod
     def insert_node(ip_address, domain_name=''):
@@ -255,12 +265,16 @@ class NodeService:
                 dst_node = NodeService.get_node_or_monkey_by_ip(bootloader_telem['tunnel'])
             else:
                 dst_node = NodeService.get_monkey_island_node()
-            edge = EdgeService.get_or_create_edge(new_node['_id'], dst_node['id'])
+            src_label = NodeService.get_label_for_endpoint(new_node['_id'])
-            mongo.db.edge.update({"_id": edge["_id"]},
+            dst_label = NodeService.get_label_for_endpoint(dst_node['id'])
-                                 {'$set': {'tunnel': bool(bootloader_telem['tunnel']),
+            edge = EdgeService.get_or_create_edge(src_node_id=new_node['_id'],
-                                           'ip_address': bootloader_telem['ips'][0],
+                                                  dst_node_id=dst_node['id'],
-                                           'group': NodeStates.get_by_keywords(['island']).value}},
+                                                  src_label=src_label,
-                                 upsert=False)
+                                                  dst_label=dst_label)
+            edge.tunnel = bool(bootloader_telem['tunnel'])
+            edge.ip_address = bootloader_telem['ips'][0]
+            edge.group = NodeStates.get_by_keywords(['island']).value
+            edge.save()
         return new_node
 
     @staticmethod
@@ -411,6 +425,15 @@ class NodeService:
     def get_hostname_by_id(node_id):
         return NodeService.get_node_hostname(mongo.db.monkey.find_one({'_id': node_id}, {'hostname': 1}))
 
+    @staticmethod
+    def get_label_for_endpoint(endpoint_id):
+        if endpoint_id == ObjectId("000000000000000000000000"):
+            return 'MonkeyIsland'
+        if Monkey.is_monkey(endpoint_id):
+            return Monkey.get_label_by_id(endpoint_id)
+        else:
+            return NodeService.get_node_label(NodeService.get_node_by_id(endpoint_id))
+
 
 class NodeCreationException(Exception):
     pass

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

@@ -2,10 +2,10 @@ import copy
 
 import dateutil
 
-from monkey_island.cc.database import mongo
 from monkey_island.cc.encryptor import encryptor
 from monkey_island.cc.models import Monkey
-from monkey_island.cc.services.edge import EdgeService
+from monkey_island.cc.models.edge import Edge
+from monkey_island.cc.services.edge.displayed_edge import EdgeService
 from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.services.telemetry.processing.utils import get_edge_by_scan_or_exploit_telemetry
 from monkey_island.cc.services.telemetry.zero_trust_tests.machine_exploited import test_machine_exploited
@@ -14,7 +14,7 @@ from monkey_island.cc.services.telemetry.zero_trust_tests.machine_exploited impo
 def process_exploit_telemetry(telemetry_json):
     encrypt_exploit_creds(telemetry_json)
     edge = get_edge_by_scan_or_exploit_telemetry(telemetry_json)
-    update_edge_info_with_new_exploit(edge, telemetry_json)
+    update_network_with_exploit(edge, telemetry_json)
     update_node_credentials_from_successful_attempts(edge, telemetry_json)
 
     test_machine_exploited(
@@ -25,28 +25,25 @@ def process_exploit_telemetry(telemetry_json):
         timestamp=telemetry_json['timestamp'])
 
 
-def update_node_credentials_from_successful_attempts(edge, telemetry_json):
+def update_node_credentials_from_successful_attempts(edge: Edge, telemetry_json):
     for attempt in telemetry_json['data']['attempts']:
         if attempt['result']:
             found_creds = {'user': attempt['user']}
             for field in ['password', 'lm_hash', 'ntlm_hash', 'ssh_key']:
                 if len(attempt[field]) != 0:
                     found_creds[field] = attempt[field]
-            NodeService.add_credentials_to_node(edge['to'], found_creds)
+            NodeService.add_credentials_to_node(edge.dst_node_id, found_creds)
 
 
-def update_edge_info_with_new_exploit(edge, telemetry_json):
+def update_network_with_exploit(edge: Edge, telemetry_json):
     telemetry_json['data']['info']['started'] = dateutil.parser.parse(telemetry_json['data']['info']['started'])
     telemetry_json['data']['info']['finished'] = dateutil.parser.parse(telemetry_json['data']['info']['finished'])
     new_exploit = copy.deepcopy(telemetry_json['data'])
     new_exploit.pop('machine')
     new_exploit['timestamp'] = telemetry_json['timestamp']
-    mongo.db.edge.update(
+    EdgeService.update_based_on_exploit(edge, new_exploit)
-        {'_id': edge['_id']},
-        {'$push': {'exploits': new_exploit}}
-    )
     if new_exploit['result']:
-        EdgeService.set_edge_exploited(edge)
+        NodeService.set_node_exploited(edge.dst_node_id)
 
 
 def encrypt_exploit_creds(telemetry_json):

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

@@ -2,7 +2,8 @@ import copy
 
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
-from monkey_island.cc.services.edge import EdgeService
+from monkey_island.cc.services.edge.edge import EdgeService
+from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.services.telemetry.processing.utils import get_edge_by_scan_or_exploit_telemetry
 from monkey_island.cc.services.telemetry.zero_trust_tests.data_endpoints import test_open_data_endpoints
 from monkey_island.cc.services.telemetry.zero_trust_tests.segmentation import test_segmentation_violation
@@ -19,22 +20,11 @@ def process_scan_telemetry(telemetry_json):
 
 def update_edges_and_nodes_based_on_scan_telemetry(telemetry_json):
     edge = get_edge_by_scan_or_exploit_telemetry(telemetry_json)
-    data = copy.deepcopy(telemetry_json['data']['machine'])
+    EdgeService.update_based_on_scan_telemetry(edge, telemetry_json)
-    ip_address = data.pop("ip_addr")
+
-    domain_name = data.pop("domain_name")
+    node = mongo.db.node.find_one({"_id": edge.dst_node_id})
-    new_scan = \
-        {
-            "timestamp": telemetry_json["timestamp"],
-            "data": data
-        }
-    mongo.db.edge.update(
-        {"_id": edge["_id"]},
-        {"$push": {"scans": new_scan},
-         "$set": {"ip_address": ip_address, 'domain_name': domain_name}}
-    )
-    node = mongo.db.node.find_one({"_id": edge["to"]})
     if node is not None:
-        scan_os = new_scan["data"]["os"]
+        scan_os = telemetry_json['data']['machine']["os"]
         if "type" in scan_os:
             mongo.db.node.update({"_id": node["_id"]},
                                  {"$set": {"os.type": scan_os["type"]}},
@@ -43,4 +33,5 @@ def update_edges_and_nodes_based_on_scan_telemetry(telemetry_json):
             mongo.db.node.update({"_id": node["_id"]},
                                  {"$set": {"os.version": scan_os["version"]}},
                                  upsert=False)
-        EdgeService.update_label_by_endpoint(edge, node["_id"])
+        label = NodeService.get_label_for_endpoint(node["_id"])
+        EdgeService.update_label(edge, node["_id"], label)

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

@@ -1,5 +1,4 @@
 import logging
-from ipaddress import ip_address
 
 from monkey_island.cc.encryptor import encryptor
 from monkey_island.cc.services import mimikatz_utils

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

@@ -1,4 +1,4 @@
-from monkey_island.cc.services.edge import EdgeService
+from monkey_island.cc.services.edge.edge import EdgeService
 from monkey_island.cc.services.node import NodeService
 
 
@@ -10,7 +10,10 @@ def get_edge_by_scan_or_exploit_telemetry(telemetry_json):
     if dst_node is None:
         dst_node = NodeService.get_or_create_node(dst_ip, dst_domain_name)
 
-    return EdgeService.get_or_create_edge(src_monkey["_id"], dst_node["_id"])
+    src_label = NodeService.get_label_for_endpoint(src_monkey["_id"])
+    dst_label = NodeService.get_label_for_endpoint(dst_node["_id"])
+
+    return EdgeService.get_or_create_edge(src_monkey["_id"], dst_node["_id"], src_label, dst_label)
 
 
 def get_tunnel_host_ip_from_proxy_field(telemetry_json):

monkey/monkey_island/cc/testing/IslandTestCase.py

@@ -1,6 +1,7 @@
 import unittest
 from monkey_island.cc.environment.environment import env
 from monkey_island.cc.models import Monkey
+from monkey_island.cc.models.edge import Edge
 from monkey_island.cc.models.zero_trust.finding import Finding
 
 
@@ -12,6 +13,10 @@ class IslandTestCase(unittest.TestCase):
     def clean_monkey_db():
         Monkey.objects().delete()
 
+    @staticmethod
+    def clean_edge_db():
+        Edge.objects().delete()
+
     @staticmethod
     def clean_finding_db():
         Finding.objects().delete()