Edge refactoring to DAL and minor PIP improvements

monkey/monkey_island/cc/bootloader_server.py

@@ -29,7 +29,8 @@ def do_POST(self):
         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
 

monkey/monkey_island/cc/models/edge.py

@@ -0,0 +1,19 @@
+from mongoengine import Document, ObjectIdField, ListField, DynamicField, BooleanField, StringField
+
+
+class Edge(Document):
+
+    meta = {'allow_inheritance': True}
+
+    # 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,15 +1,16 @@
-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'
 
 
 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,13 +3,15 @@
 
 import dateutil.parser
 import flask_restful
+
 from monkey_island.cc.resources.test.utils.telem_store import TestTelemStore
 from flask import request
 
 from monkey_island.cc.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
 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'
@@ -86,7 +88,8 @@ def post(self, **kw):
         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):
@@ -95,7 +98,8 @@ def post(self, **kw):
                 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):
@@ -129,8 +133,8 @@ def post(self, **kw):
 
         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}})
+            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})

monkey/monkey_island/cc/resources/netmap.py

@@ -1,29 +1,20 @@
 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'
 
 
 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
             }

monkey/monkey_island/cc/resources/test/clear_caches.py

@@ -13,8 +13,8 @@
 
 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()

monkey/monkey_island/cc/services/attack/attack_schema.py

@@ -78,7 +78,7 @@
                     "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."
                 }
             }
         },

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'}},