From 092482ad8755d1720ba97773c72949545eac66d2 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Thu, 28 May 2020 20:12:10 +0300
Subject: [PATCH 1/4] Refactored edge from mongo queries to DAL

---
 monkey/monkey_island/cc/models/edge.py        |  16 ++
 monkey/monkey_island/cc/resources/edge.py     |   7 +-
 monkey/monkey_island/cc/resources/monkey.py   |   7 +-
 monkey/monkey_island/cc/resources/netmap.py   |  21 +--
 .../cc/services/attack/attack_schema.py       |   2 +-
 monkey/monkey_island/cc/services/edge.py      | 173 ------------------
 .../cc/services/edge/__init__.py              |   0
 .../cc/services/edge/displayed_edge.py        | 100 ++++++++++
 monkey/monkey_island/cc/services/edge/edge.py |  78 ++++++++
 .../cc/services/edge/test_displayed_edge.py   |  97 ++++++++++
 .../cc/services/edge/test_edge.py             |  60 ++++++
 .../cc/services/netmap/net_edge.py            |  68 +++++++
 .../cc/services/netmap/net_node.py            |  23 +++
 monkey/monkey_island/cc/services/node.py      |  63 +++++--
 .../services/telemetry/processing/exploit.py  |  19 +-
 .../cc/services/telemetry/processing/scan.py  |  25 +--
 .../telemetry/processing/system_info.py       |   1 -
 .../cc/services/telemetry/processing/utils.py |   7 +-
 .../cc/testing/IslandTestCase.py              |   5 +
 19 files changed, 527 insertions(+), 245 deletions(-)
 create mode 100644 monkey/monkey_island/cc/models/edge.py
 delete mode 100644 monkey/monkey_island/cc/services/edge.py
 create mode 100644 monkey/monkey_island/cc/services/edge/__init__.py
 create mode 100644 monkey/monkey_island/cc/services/edge/displayed_edge.py
 create mode 100644 monkey/monkey_island/cc/services/edge/edge.py
 create mode 100644 monkey/monkey_island/cc/services/edge/test_displayed_edge.py
 create mode 100644 monkey/monkey_island/cc/services/edge/test_edge.py
 create mode 100644 monkey/monkey_island/cc/services/netmap/net_edge.py
 create mode 100644 monkey/monkey_island/cc/services/netmap/net_node.py

diff --git a/monkey/monkey_island/cc/models/edge.py b/monkey/monkey_island/cc/models/edge.py
new file mode 100644
index 000000000..31a51598e
--- /dev/null
+++ b/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()
diff --git a/monkey/monkey_island/cc/resources/edge.py b/monkey/monkey_island/cc/resources/edge.py
index f3ce94ee3..3d284e82c 100644
--- a/monkey/monkey_island/cc/resources/edge.py
+++ b/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 {}
diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py
index 181ce94b2..5b5e1af4a 100644
--- a/monkey/monkey_island/cc/resources/monkey.py
+++ b/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}):
-                mongo.db.edge.update({"_id": edge["_id"]}, {"$set": {"to": new_monkey_id}})
+            for edge in Edge.objects(dst_node_id=node_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})
diff --git a/monkey/monkey_island/cc/resources/netmap.py b/monkey/monkey_island/cc/resources/netmap.py
index 3b7e471d8..3b2c4c12e 100644
--- a/monkey/monkey_island/cc/resources/netmap.py
+++ b/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.node import NodeService
-from monkey_island.cc.database import mongo
+from monkey_island.cc.services.netmap.net_edge import NetEdgeService
+from monkey_island.cc.services.netmap.net_node import NetNodeService
 
 __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({})]
-        nodes = [NodeService.node_to_net_node(x) for x in mongo.db.node.find({})]
-        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()
+        net_nodes = NetNodeService.get_all_net_nodes()
+        net_edges = NetEdgeService.get_all_net_edges()
 
         return \
             {
-                "nodes": monkeys + nodes + monkey_island,
-                "edges": edges
+                "nodes": net_nodes,
+                "edges": net_edges
             }
diff --git a/monkey/monkey_island/cc/services/attack/attack_schema.py b/monkey/monkey_island/cc/services/attack/attack_schema.py
index 3c3a451f2..177dc9eaa 100644
--- a/monkey/monkey_island/cc/services/attack/attack_schema.py
+++ b/monkey/monkey_island/cc/services/attack/attack_schema.py
@@ -78,7 +78,7 @@ SCHEMA = {
                     "necessary": False,
                     "link": "https://attack.mitre.org/techniques/T1136",
                     "description": "Adversaries with a sufficient level of access "
-                                    "may create a local system, domain, or cloud tenant account."
+                                   "may create a local system, domain, or cloud tenant account."
                 }
             }
         },
diff --git a/monkey/monkey_island/cc/services/edge.py b/monkey/monkey_island/cc/services/edge.py
deleted file mode 100644
index cd4ef737b..000000000
--- a/monkey/monkey_island/cc/services/edge.py
+++ /dev/null
@@ -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"
diff --git a/monkey/monkey_island/cc/services/edge/__init__.py b/monkey/monkey_island/cc/services/edge/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/monkey/monkey_island/cc/services/edge/displayed_edge.py b/monkey/monkey_island/cc/services/edge/displayed_edge.py
new file mode 100644
index 000000000..761dfa555
--- /dev/null
+++ b/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"
diff --git a/monkey/monkey_island/cc/services/edge/edge.py b/monkey/monkey_island/cc/services/edge/edge.py
new file mode 100644
index 000000000..e34ce56d1
--- /dev/null
+++ b/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'])
diff --git a/monkey/monkey_island/cc/services/edge/test_displayed_edge.py b/monkey/monkey_island/cc/services/edge/test_displayed_edge.py
new file mode 100644
index 000000000..340efe771
--- /dev/null
+++ b/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"])
diff --git a/monkey/monkey_island/cc/services/edge/test_edge.py b/monkey/monkey_island/cc/services/edge/test_edge.py
new file mode 100644
index 000000000..8053f45e3
--- /dev/null
+++ b/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))
diff --git a/monkey/monkey_island/cc/services/netmap/net_edge.py b/monkey/monkey_island/cc/services/netmap/net_edge.py
new file mode 100644
index 000000000..f9d5e1932
--- /dev/null
+++ b/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
diff --git a/monkey/monkey_island/cc/services/netmap/net_node.py b/monkey/monkey_island/cc/services/netmap/net_node.py
new file mode 100644
index 000000000..796167cf5
--- /dev/null
+++ b/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({})]
diff --git a/monkey/monkey_island/cc/services/node.py b/monkey/monkey_island/cc/services/node.py
index 9c0921580..ecce06c85 100644
--- a/monkey/monkey_island/cc/services/node.py
+++ b/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"]:
-                edge_exploit["origin"] = from_node_label
+            for edge_exploit in edge['exploits']:
+                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(
-            {"from": monkey_id, 'tunnel': True},
-            {'$set': {'tunnel': False}},
-            upsert=False)
+        try:
+            edge = Edge.objects.get(src_node_id=monkey_id, tunnel=True)
+            edge.tunnel = 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)
-        mongo.db.edge.update({"_id": tunnel_edge["_id"]},
-                             {'$set': {'tunnel': True, 'ip_address': tunnel_host_ip}},
-                             upsert=False)
+        monkey_label = NodeService.get_label_for_endpoint(monkey_id)
+        tunnel_host_label = NodeService.get_label_for_endpoint(tunnel_host_id)
+        tunnel_edge = EdgeService.get_or_create_edge(src_node_id=monkey_id,
+                                                     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'])
-            mongo.db.edge.update({"_id": edge["_id"]},
-                                 {'$set': {'tunnel': bool(bootloader_telem['tunnel']),
-                                           'ip_address': bootloader_telem['ips'][0],
-                                           'group': NodeStates.get_by_keywords(['island']).value}},
-                                 upsert=False)
+            src_label = NodeService.get_label_for_endpoint(new_node['_id'])
+            dst_label = NodeService.get_label_for_endpoint(dst_node['id'])
+            edge = EdgeService.get_or_create_edge(src_node_id=new_node['_id'],
+                                                  dst_node_id=dst_node['id'],
+                                                  src_label=src_label,
+                                                  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
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/exploit.py b/monkey/monkey_island/cc/services/telemetry/processing/exploit.py
index 9b25c97ef..8cb5db0d6 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/exploit.py
+++ b/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(
-        {'_id': edge['_id']},
-        {'$push': {'exploits': new_exploit}}
-    )
+    EdgeService.update_based_on_exploit(edge, 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):
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/scan.py b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
index d57aef7c0..6566b91f2 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/scan.py
+++ b/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'])
-    ip_address = data.pop("ip_addr")
-    domain_name = data.pop("domain_name")
-    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"]})
+    EdgeService.update_based_on_scan_telemetry(edge, telemetry_json)
+
+    node = mongo.db.node.find_one({"_id": edge.dst_node_id})
     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)
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 844724163..5b842df0b 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/system_info.py
+++ b/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
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/utils.py b/monkey/monkey_island/cc/services/telemetry/processing/utils.py
index 466b81bf1..df898945e 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/utils.py
+++ b/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):
diff --git a/monkey/monkey_island/cc/testing/IslandTestCase.py b/monkey/monkey_island/cc/testing/IslandTestCase.py
index 35e568399..e85d2f9ec 100644
--- a/monkey/monkey_island/cc/testing/IslandTestCase.py
+++ b/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()

From a160e3396b542369860ae6d9e1995d2e4151446b Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Fri, 29 May 2020 11:42:27 +0300
Subject: [PATCH 2/4] Small PEP improvements all around, unused method deleted
 on displayed_edge.py

---
 monkey/monkey_island/cc/bootloader_server.py  |  3 ++-
 monkey/monkey_island/cc/resources/monkey.py   |  6 ++++--
 .../cc/resources/test/clear_caches.py         |  4 ++--
 .../attack/technique_reports/T1082.py         |  3 ++-
 .../cc/services/edge/displayed_edge.py        | 21 -------------------
 monkey/monkey_island/cc/services/edge/edge.py | 10 ++++-----
 .../cc/services/reporting/pth_report.py       |  3 ++-
 .../cc/services/telemetry/processing/scan.py  |  2 --
 .../system_info_telemetry_dispatcher.py       |  5 +++--
 .../test_system_info_telemetry_dispatcher.py  |  8 +++++--
 10 files changed, 26 insertions(+), 39 deletions(-)

diff --git a/monkey/monkey_island/cc/bootloader_server.py b/monkey/monkey_island/cc/bootloader_server.py
index b1f7ec484..30196ff34 100644
--- a/monkey/monkey_island/cc/bootloader_server.py
+++ b/monkey/monkey_island/cc/bootloader_server.py
@@ -29,7 +29,8 @@ class BootloaderHTTPRequestHandler(BaseHTTPRequestHandler):
         post_data = self.rfile.read(content_length).decode()
         island_server_path = BootloaderHTTPRequestHandler.get_bootloader_resource_url(self.request.getsockname()[0])
         island_server_path = parse.urljoin(island_server_path, self.path[1:])
-        # The island server doesn't always have a correct SSL cert installed (By default it comes with a self signed one),
+        # The island server doesn't always have a correct SSL cert installed
+        # (By default it comes with a self signed one),
         # that's why we're not verifying the cert in this request.
         r = requests.post(url=island_server_path, data=post_data, verify=False)  # noqa: DUO123
 
diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py
index 5b5e1af4a..ae34c624d 100644
--- a/monkey/monkey_island/cc/resources/monkey.py
+++ b/monkey/monkey_island/cc/resources/monkey.py
@@ -88,7 +88,8 @@ class Monkey(flask_restful.Resource):
         parent_to_add = (monkey_json.get('guid'), None)  # default values in case of manual run
         if parent and parent != monkey_json.get('guid'):  # current parent is known
             exploit_telem = [x for x in
-                             mongo.db.telemetry.find({'telem_category': {'$eq': 'exploit'}, 'data.result': {'$eq': True},
+                             mongo.db.telemetry.find({'telem_category': {'$eq': 'exploit'},
+                                                      'data.result': {'$eq': True},
                                                       'data.machine.ip_addr': {'$in': monkey_json['ip_addresses']},
                                                       'monkey_guid': {'$eq': parent}})]
             if 1 == len(exploit_telem):
@@ -97,7 +98,8 @@ class Monkey(flask_restful.Resource):
                 parent_to_add = (parent, None)
         elif (not parent or parent == monkey_json.get('guid')) and 'ip_addresses' in monkey_json:
             exploit_telem = [x for x in
-                             mongo.db.telemetry.find({'telem_category': {'$eq': 'exploit'}, 'data.result': {'$eq': True},
+                             mongo.db.telemetry.find({'telem_category': {'$eq': 'exploit'},
+                                                      'data.result': {'$eq': True},
                                                       'data.machine.ip_addr': {'$in': monkey_json['ip_addresses']}})]
 
             if 1 == len(exploit_telem):
diff --git a/monkey/monkey_island/cc/resources/test/clear_caches.py b/monkey/monkey_island/cc/resources/test/clear_caches.py
index f17193821..8d510a8bb 100644
--- a/monkey/monkey_island/cc/resources/test/clear_caches.py
+++ b/monkey/monkey_island/cc/resources/test/clear_caches.py
@@ -13,8 +13,8 @@ logger = logging.getLogger(__name__)
 
 class ClearCaches(flask_restful.Resource):
     """
-    Used for timing tests - we want to get actual execution time of functions in BlackBox without caching - so we use this
-    to clear the caches.
+    Used for timing tests - we want to get actual execution time of functions in BlackBox without caching -
+    so we use this to clear the caches.
     :note: DO NOT CALL THIS IN PRODUCTION CODE as this will slow down the user experience.
     """
     @jwt_required()
diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1082.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1082.py
index cdbdb42ec..c80a3d476 100644
--- a/monkey/monkey_island/cc/services/attack/technique_reports/T1082.py
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1082.py
@@ -23,7 +23,8 @@ class T1082(AttackTechnique):
                            'collections': [
                                {'used': {'$and': [{'$ifNull': ['$netstat', False]}, {'$gt': ['$aws', {}]}]},
                                 'name': {'$literal': 'Amazon Web Services info'}},
-                               {'used': {'$and': [{'$ifNull': ['$process_list', False]}, {'$gt': ['$process_list', {}]}]},
+                               {'used': {'$and': [{'$ifNull': ['$process_list', False]},
+                                                  {'$gt': ['$process_list', {}]}]},
                                 'name': {'$literal': 'Running process list'}},
                                {'used': {'$and': [{'$ifNull': ['$netstat', False]}, {'$ne': ['$netstat', []]}]},
                                 'name': {'$literal': 'Network connections'}},
diff --git a/monkey/monkey_island/cc/services/edge/displayed_edge.py b/monkey/monkey_island/cc/services/edge/displayed_edge.py
index 761dfa555..8f94a6ffa 100644
--- a/monkey/monkey_island/cc/services/edge/displayed_edge.py
+++ b/monkey/monkey_island/cc/services/edge/displayed_edge.py
@@ -2,8 +2,6 @@ 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
 
@@ -58,25 +56,6 @@ class DisplayedEdgeService:
         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:
diff --git a/monkey/monkey_island/cc/services/edge/edge.py b/monkey/monkey_island/cc/services/edge/edge.py
index e34ce56d1..ab4c3c114 100644
--- a/monkey/monkey_island/cc/services/edge/edge.py
+++ b/monkey/monkey_island/cc/services/edge/edge.py
@@ -58,6 +58,11 @@ class EdgeService:
         if exploit['result']:
             EdgeService.set_edge_exploited(edge)
 
+    @staticmethod
+    def set_edge_exploited(edge: Edge):
+        edge.exploited = True
+        edge.save()
+
     @staticmethod
     def get_edge_group(edge: Edge):
         if edge.exploited:
@@ -68,11 +73,6 @@ class EdgeService:
             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'])
diff --git a/monkey/monkey_island/cc/services/reporting/pth_report.py b/monkey/monkey_island/cc/services/reporting/pth_report.py
index ecb209c69..f6d7b615a 100644
--- a/monkey/monkey_island/cc/services/reporting/pth_report.py
+++ b/monkey/monkey_island/cc/services/reporting/pth_report.py
@@ -106,7 +106,8 @@ class PTHReportService(object):
                 {
                     'username': user['name'],
                     'domain_name': user['domain_name'],
-                    'hostname': NodeService.get_hostname_by_id(ObjectId(user['machine_id'])) if user['machine_id'] else None
+                    'hostname': NodeService.get_hostname_by_id(ObjectId(user['machine_id'])) if
+                    user['machine_id'] else None
                 } for user in doc['Docs']
             ]
             users_cred_groups.append({'cred_groups': users_list})
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/scan.py b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
index 6566b91f2..48c1f11c3 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/scan.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
@@ -1,5 +1,3 @@
-import copy
-
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.services.edge.edge import EdgeService
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py
index b5f2d24ea..af477ebb4 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py
@@ -21,8 +21,9 @@ SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSORS = {
 class SystemInfoTelemetryDispatcher(object):
     def __init__(self, collector_to_parsing_functions: typing.Mapping[str, typing.List[typing.Callable]] = None):
         """
-        :param collector_to_parsing_functions: Map between collector names and a list of functions that process the output of
-        that collector. If `None` is supplied, uses the default one; This should be the normal flow, overriding the
+        :param collector_to_parsing_functions: Map between collector names and a list of functions
+        that process the output of that collector.
+        If `None` is supplied, uses the default one; This should be the normal flow, overriding the
         collector->functions mapping is useful mostly for testing.
         """
         if collector_to_parsing_functions is None:
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py
index c5cc7aca2..f5a72405d 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py
@@ -22,9 +22,13 @@ class SystemInfoTelemetryDispatcherTest(IslandTestCase):
         bad_empty_telem_json = {}
         self.assertRaises(KeyError, dispatcher.dispatch_collector_results_to_relevant_processors, bad_empty_telem_json)
         bad_no_data_telem_json = {"monkey_guid": "bla"}
-        self.assertRaises(KeyError, dispatcher.dispatch_collector_results_to_relevant_processors, bad_no_data_telem_json)
+        self.assertRaises(KeyError,
+                          dispatcher.dispatch_collector_results_to_relevant_processors,
+                          bad_no_data_telem_json)
         bad_no_monkey_telem_json = {"data": {"collectors": {"AwsCollector": "Bla"}}}
-        self.assertRaises(KeyError, dispatcher.dispatch_collector_results_to_relevant_processors, bad_no_monkey_telem_json)
+        self.assertRaises(KeyError,
+                          dispatcher.dispatch_collector_results_to_relevant_processors,
+                          bad_no_monkey_telem_json)
 
         # Telem JSON with no collectors - nothing gets dispatched
         good_telem_no_collectors = {"monkey_guid": "bla", "data": {"bla": "bla"}}

From fb595319699d16fc70fa089302ee8f9afacc4592 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Mon, 8 Jun 2020 10:29:04 +0300
Subject: [PATCH 3/4] Refactored EdgeService into a boundary object.

---
 monkey/monkey_island/cc/models/edge.py        |   3 +
 monkey/monkey_island/cc/resources/monkey.py   |   7 +-
 .../cc/services/edge/displayed_edge.py        |  28 +++--
 monkey/monkey_island/cc/services/edge/edge.py | 103 +++++++++++-------
 .../cc/services/edge/test_displayed_edge.py   |  20 ++--
 .../cc/services/edge/test_edge.py             |   8 +-
 .../cc/services/netmap/net_edge.py            |   6 +-
 monkey/monkey_island/cc/services/node.py      |   9 +-
 .../services/telemetry/processing/exploit.py  |   6 +-
 .../cc/services/telemetry/processing/scan.py  |   4 +-
 10 files changed, 111 insertions(+), 83 deletions(-)

diff --git a/monkey/monkey_island/cc/models/edge.py b/monkey/monkey_island/cc/models/edge.py
index 31a51598e..09af04680 100644
--- a/monkey/monkey_island/cc/models/edge.py
+++ b/monkey/monkey_island/cc/models/edge.py
@@ -2,6 +2,9 @@ from mongoengine import Document, ObjectIdField, ListField, DynamicField, Boolea
 
 
 class Edge(Document):
+
+    meta = {'allow_inheritance': True}
+
     # SCHEMA
     src_node_id = ObjectIdField(required=True)
     dst_node_id = ObjectIdField(required=True)
diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py
index ae34c624d..a39aaf199 100644
--- a/monkey/monkey_island/cc/resources/monkey.py
+++ b/monkey/monkey_island/cc/resources/monkey.py
@@ -4,7 +4,6 @@ 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
 
@@ -12,6 +11,7 @@ from monkey_island.cc.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECOND
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models.monkey_ttl import create_monkey_ttl_document
 from monkey_island.cc.services.config import ConfigService
+from monkey_island.cc.services.edge.edge import EdgeService
 from monkey_island.cc.services.node import NodeService
 
 __author__ = 'Barak'
@@ -133,9 +133,8 @@ class Monkey(flask_restful.Resource):
 
         if existing_node:
             node_id = existing_node["_id"]
-            for edge in Edge.objects(dst_node_id=node_id):
-                edge.dst_node_id = new_monkey_id
-                edge.save()
+            EdgeService.update_all_dst_nodes(old_dst_node_id=node_id,
+                                             new_dst_node_id=new_monkey_id)
             for creds in existing_node['creds']:
                 NodeService.add_credentials_to_monkey(new_monkey_id, creds)
             mongo.db.node.remove({"_id": node_id})
diff --git a/monkey/monkey_island/cc/services/edge/displayed_edge.py b/monkey/monkey_island/cc/services/edge/displayed_edge.py
index 8f94a6ffa..f7a0664bf 100644
--- a/monkey/monkey_island/cc/services/edge/displayed_edge.py
+++ b/monkey/monkey_island/cc/services/edge/displayed_edge.py
@@ -1,8 +1,8 @@
 from copy import deepcopy
+from typing import Dict
 
 from bson import ObjectId
 
-from monkey_island.cc.models.edge import Edge
 from monkey_island.cc.services.edge.edge import EdgeService
 
 __author__ = "itay.mizeretz"
@@ -11,18 +11,18 @@ __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))
+    def get_displayed_edges_by_dst(dst_id: str, for_report=False):
+        edges = EdgeService.get_by_dst_node(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)
+    def get_displayed_edge_by_id(edge_id: str, for_report=False):
+        edge = EdgeService.get_edge_by_id(ObjectId(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):
+    def edge_to_displayed_edge(edge: EdgeService, for_report=False):
         services = []
         os = {}
 
@@ -33,13 +33,13 @@ class DisplayedEdgeService:
 
         displayed_edge = DisplayedEdgeService.edge_to_net_edge(edge)
 
-        displayed_edge["ip_address"] = edge['ip_address']
+        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)
+        displayed_edge["exploits"] = deepcopy(edge.exploits)
+        displayed_edge["_label"] = edge.get_label()
         return displayed_edge
 
     @staticmethod
@@ -53,9 +53,13 @@ class DisplayedEdgeService:
                 "src_label": src_label,
                 "dst_label": dst_label
             }
-        edge["_label"] = EdgeService.get_edge_label(edge)
+        edge["_label"] = DisplayedEdgeService.get_pseudo_label(edge)
         return edge
 
+    @staticmethod
+    def get_pseudo_label(edge: Dict):
+        return f"{edge['src_label']} {RIGHT_ARROW} {edge['dst_label']}"
+
     @staticmethod
     def services_to_displayed_services(services, for_report=False):
         if for_report:
@@ -64,13 +68,13 @@ class DisplayedEdgeService:
             return [x + ": " + (services[x]['name'] if 'name' in services[x] else 'unknown') for x in services]
 
     @staticmethod
-    def edge_to_net_edge(edge: Edge):
+    def edge_to_net_edge(edge: EdgeService):
         return \
             {
                 "id": edge.id,
                 "from": edge.src_node_id,
                 "to": edge.dst_node_id,
-                "group": EdgeService.get_edge_group(edge),
+                "group": edge.get_group(),
                 "src_label": edge.src_label,
                 "dst_label": edge.dst_label
             }
diff --git a/monkey/monkey_island/cc/services/edge/edge.py b/monkey/monkey_island/cc/services/edge/edge.py
index ab4c3c114..4c9ef57d7 100644
--- a/monkey/monkey_island/cc/services/edge/edge.py
+++ b/monkey/monkey_island/cc/services/edge/edge.py
@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import copy
-from typing import Dict
+from typing import Dict, List
 
 from bson import ObjectId
 from mongoengine import DoesNotExist
@@ -9,35 +11,60 @@ from monkey_island.cc.models.edge import Edge
 RIGHT_ARROW = "\u2192"
 
 
-class EdgeService:
+class EdgeService(Edge):
 
     @staticmethod
-    def get_or_create_edge(src_node_id, dst_node_id, src_label, dst_label):
-        edge = False
+    def get_all_edges() -> List[EdgeService]:
+        return EdgeService.objects()
+
+    @staticmethod
+    def get_or_create_edge(src_node_id, dst_node_id, src_label, dst_label) -> EdgeService:
+        edge = None
         try:
-            edge = Edge.objects.get(src_node_id=src_node_id, dst_node_id=dst_node_id)
+            edge = EdgeService.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)
+            edge = EdgeService(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()
+                edge.update_label(node_id=src_node_id, label=src_label)
+                edge.update_label(node_id=dst_node_id, label=dst_label)
         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
+    def get_by_dst_node(dst_node_id: ObjectId) -> List[EdgeService]:
+        return EdgeService.objects(dst_node_id=dst_node_id)
 
     @staticmethod
-    def update_based_on_scan_telemetry(edge: Edge, telemetry: Dict):
+    def get_edge_by_id(edge_id: ObjectId) -> EdgeService:
+        return EdgeService.objects.get(id=edge_id)
+
+    def update_label(self, node_id: ObjectId, label: str):
+        if self.src_node_id == node_id:
+            self.src_label = label
+        elif self.dst_node_id == node_id:
+            self.dst_label = label
+        else:
+            raise DoesNotExist("Node id provided does not match with any endpoint of an self provided.")
+        self.save()
+
+    @staticmethod
+    def update_all_dst_nodes(old_dst_node_id: ObjectId, new_dst_node_id: ObjectId):
+        for edge in EdgeService.objects(dst_node_id=old_dst_node_id):
+            edge.dst_node_id = new_dst_node_id
+            edge.save()
+
+    @staticmethod
+    def get_tunnel_edges_by_src(src_node_id) -> List[EdgeService]:
+        try:
+            return EdgeService.objects(src_node_id=src_node_id, tunnel=True)
+        except DoesNotExist:
+            return []
+
+    def disable_tunnel(self):
+        self.tunnel = False
+        self.save()
+
+    def update_based_on_scan_telemetry(self, telemetry: Dict):
         machine_info = copy.deepcopy(telemetry['data']['machine'])
         new_scan = \
             {
@@ -46,33 +73,29 @@ class EdgeService:
             }
         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()
+        self.scans.append(new_scan)
+        self.ip_address = ip_address
+        self.domain_name = domain_name
+        self.save()
 
-    @staticmethod
-    def update_based_on_exploit(edge: Edge, exploit: Dict):
-        edge.exploits.append(exploit)
-        edge.save()
+    def update_based_on_exploit(self, exploit: Dict):
+        self.exploits.append(exploit)
+        self.save()
         if exploit['result']:
-            EdgeService.set_edge_exploited(edge)
+            self.set_exploited()
 
-    @staticmethod
-    def set_edge_exploited(edge: Edge):
-        edge.exploited = True
-        edge.save()
+    def set_exploited(self):
+        self.exploited = True
+        self.save()
 
-    @staticmethod
-    def get_edge_group(edge: Edge):
-        if edge.exploited:
+    def get_group(self) -> str:
+        if self.exploited:
             return "exploited"
-        if edge.tunnel:
+        if self.tunnel:
             return "tunnel"
-        if edge.scans or edge.exploits:
+        if self.scans or self.exploits:
             return "scan"
         return "empty"
 
-    @staticmethod
-    def get_edge_label(edge):
-        return "%s %s %s" % (edge['src_label'], RIGHT_ARROW, edge['dst_label'])
+    def get_label(self) -> str:
+        return f"{self.src_label} {RIGHT_ARROW} {self.dst_label}"
diff --git a/monkey/monkey_island/cc/services/edge/test_displayed_edge.py b/monkey/monkey_island/cc/services/edge/test_displayed_edge.py
index 340efe771..c134bce33 100644
--- a/monkey/monkey_island/cc/services/edge/test_displayed_edge.py
+++ b/monkey/monkey_island/cc/services/edge/test_displayed_edge.py
@@ -58,21 +58,21 @@ class TestDisplayedEdgeService(IslandTestCase):
         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)
+        displayed_edges = DisplayedEdgeService.get_displayed_edges_by_dst(str(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")
+        edge = EdgeService(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)
 
diff --git a/monkey/monkey_island/cc/services/edge/test_edge.py b/monkey/monkey_island/cc/services/edge/test_edge.py
index 8053f45e3..8ebf45343 100644
--- a/monkey/monkey_island/cc/services/edge/test_edge.py
+++ b/monkey/monkey_island/cc/services/edge/test_edge.py
@@ -45,16 +45,16 @@ class TestEdgeService(IslandTestCase):
         edge = Edge(src_node_id=ObjectId(),
                     dst_node_id=ObjectId(),
                     exploited=True)
-        self.assertEqual("exploited", EdgeService.get_edge_group(edge))
+        self.assertEqual("exploited", EdgeService.get_group(edge))
 
         edge.exploited = False
         edge.tunnel = True
-        self.assertEqual("tunnel", EdgeService.get_edge_group(edge))
+        self.assertEqual("tunnel", EdgeService.get_group(edge))
 
         edge.tunnel = False
         edge.exploits.append(["mock_exploit_data"])
-        self.assertEqual("scan", EdgeService.get_edge_group(edge))
+        self.assertEqual("scan", EdgeService.get_group(edge))
 
         edge.exploits = []
         edge.scans = []
-        self.assertEqual("empty", EdgeService.get_edge_group(edge))
+        self.assertEqual("empty", EdgeService.get_group(edge))
diff --git a/monkey/monkey_island/cc/services/netmap/net_edge.py b/monkey/monkey_island/cc/services/netmap/net_edge.py
index f9d5e1932..44e097630 100644
--- a/monkey/monkey_island/cc/services/netmap/net_edge.py
+++ b/monkey/monkey_island/cc/services/netmap/net_edge.py
@@ -3,6 +3,7 @@ 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.edge.edge import EdgeService
 from monkey_island.cc.services.node import NodeService
 
 
@@ -20,7 +21,7 @@ class NetEdgeService:
 
     @staticmethod
     def _get_standard_net_edges():
-        return [DisplayedEdgeService.edge_to_net_edge(x) for x in Edge.objects()]
+        return [DisplayedEdgeService.edge_to_net_edge(x) for x in EdgeService.get_all_edges()]
 
     @staticmethod
     def _get_uninfected_island_net_edges():
@@ -44,7 +45,8 @@ class NetEdgeService:
 
     @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"])]
+        existing_ids = [x.src_node_id for x
+                        in EdgeService.get_by_dst_node(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
diff --git a/monkey/monkey_island/cc/services/node.py b/monkey/monkey_island/cc/services/node.py
index ecce06c85..f6f5362c3 100644
--- a/monkey/monkey_island/cc/services/node.py
+++ b/monkey/monkey_island/cc/services/node.py
@@ -189,12 +189,9 @@ class NodeService:
             {'$unset': {'tunnel': ''}},
             upsert=False)
 
-        try:
-            edge = Edge.objects.get(src_node_id=monkey_id, tunnel=True)
-            edge.tunnel = False
-            edge.save()
-        except DoesNotExist:
-            pass
+        edges = EdgeService.get_tunnel_edges_by_src(monkey_id)
+        for edge in edges:
+            edge.disable_tunnel()
 
     @staticmethod
     def set_monkey_tunnel(monkey_id, tunnel_host_ip):
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/exploit.py b/monkey/monkey_island/cc/services/telemetry/processing/exploit.py
index 8cb5db0d6..f8ea52b6e 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/exploit.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/exploit.py
@@ -25,7 +25,7 @@ def process_exploit_telemetry(telemetry_json):
         timestamp=telemetry_json['timestamp'])
 
 
-def update_node_credentials_from_successful_attempts(edge: Edge, telemetry_json):
+def update_node_credentials_from_successful_attempts(edge: EdgeService, telemetry_json):
     for attempt in telemetry_json['data']['attempts']:
         if attempt['result']:
             found_creds = {'user': attempt['user']}
@@ -35,13 +35,13 @@ def update_node_credentials_from_successful_attempts(edge: Edge, telemetry_json)
             NodeService.add_credentials_to_node(edge.dst_node_id, found_creds)
 
 
-def update_network_with_exploit(edge: Edge, telemetry_json):
+def update_network_with_exploit(edge: EdgeService, 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']
-    EdgeService.update_based_on_exploit(edge, new_exploit)
+    edge.update_based_on_exploit(new_exploit)
     if new_exploit['result']:
         NodeService.set_node_exploited(edge.dst_node_id)
 
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/scan.py b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
index 48c1f11c3..7a7f0b19c 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/scan.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
@@ -18,7 +18,7 @@ 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)
-    EdgeService.update_based_on_scan_telemetry(edge, telemetry_json)
+    edge.update_based_on_scan_telemetry(telemetry_json)
 
     node = mongo.db.node.find_one({"_id": edge.dst_node_id})
     if node is not None:
@@ -32,4 +32,4 @@ def update_edges_and_nodes_based_on_scan_telemetry(telemetry_json):
                                  {"$set": {"os.version": scan_os["version"]}},
                                  upsert=False)
         label = NodeService.get_label_for_endpoint(node["_id"])
-        EdgeService.update_label(edge, node["_id"], label)
+        edge.update_label(node["_id"], label)

From 4c8319669f29cecb3adceea8f66be642c524c68f Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Mon, 8 Jun 2020 10:31:31 +0300
Subject: [PATCH 4/4] Map jiggle small-fix: wait less for map to stop moving

---
 monkey/monkey_island/cc/ui/src/components/map/MapOptions.js | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/monkey/monkey_island/cc/ui/src/components/map/MapOptions.js b/monkey/monkey_island/cc/ui/src/components/map/MapOptions.js
index e09e93124..e09040b01 100644
--- a/monkey/monkey_island/cc/ui/src/components/map/MapOptions.js
+++ b/monkey/monkey_island/cc/ui/src/components/map/MapOptions.js
@@ -45,7 +45,7 @@ export const basic_options = {
       springLength: 100,
       springConstant: 0.025
     },
-    minVelocity: 0.3,
+    minVelocity: 0.7,
     maxVelocity: 25
   }
 };