From 72d266a1b824e103e2a770d2490c09600cbdee8f Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Wed, 12 Feb 2020 11:09:11 +0200
Subject: [PATCH] Added node and edge based on bootloader telemetry.

---
 monkey/monkey_island/cc/bootloader_server.py  |  1 +
 .../monkey_island/cc/resources/bootloader.py  | 15 +++++---
 monkey/monkey_island/cc/services/edge.py      |  7 ++--
 monkey/monkey_island/cc/services/node.py      | 36 +++++++++++++++++++
 4 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/monkey/monkey_island/cc/bootloader_server.py b/monkey/monkey_island/cc/bootloader_server.py
index 1638c6316..7706817f8 100644
--- a/monkey/monkey_island/cc/bootloader_server.py
+++ b/monkey/monkey_island/cc/bootloader_server.py
@@ -27,6 +27,7 @@ class BootloaderHTTPRequestHandler(BaseHTTPRequestHandler):
             conf = self.server.mongo_client['monkeyisland']['config'].find_one({'name': 'initial'})
         island_server_path = BootloaderHTTPRequestHandler.get_bootloader_resource_path_from_config(conf)
         r = requests.post(url=island_server_path, data=post_data, verify=False)
+
         if r.status_code != 200:
             self.send_response(404)
         else:
diff --git a/monkey/monkey_island/cc/resources/bootloader.py b/monkey/monkey_island/cc/resources/bootloader.py
index 80a7b56d7..eab4cd5c9 100644
--- a/monkey/monkey_island/cc/resources/bootloader.py
+++ b/monkey/monkey_island/cc/resources/bootloader.py
@@ -1,6 +1,11 @@
+import json
+
 import flask_restful
 from flask import request, make_response
 
+from monkey_island.cc.database import mongo
+from monkey_island.cc.services.node import NodeService
+
 WINDOWS_VERSIONS = {
     "5.0": "Windows 2000",
     "5.1": "Windows XP",
@@ -17,9 +22,11 @@ class Bootloader(flask_restful.Resource):
 
     # Used by monkey. can't secure.
     def post(self, **kw):
-        os_version = request.data.decode().split(" ")
-        if (os_version[0][0] == "W"):
-            os_type = "windows"
-            os_version = os_version[1:]
+        data = json.loads(request.data.decode().replace("\n", ""))
+        local_addr = [i for i in data["ips"] if i.startswith("127")]
+        if local_addr:
+            data["ips"].remove(local_addr[0])
+        mongo.db.bootloader_telems.insert(data)
+        node_id = NodeService.get_or_create_node_from_bootloader_telem(data)
 
         return make_response({"status": "OK"}, 200)
diff --git a/monkey/monkey_island/cc/services/edge.py b/monkey/monkey_island/cc/services/edge.py
index d5e38768d..a8016f5a8 100644
--- a/monkey/monkey_island/cc/services/edge.py
+++ b/monkey/monkey_island/cc/services/edge.py
@@ -2,7 +2,7 @@ from bson import ObjectId
 
 from monkey_island.cc.database import mongo
 import monkey_island.cc.services.node
-from monkey_island.cc.models import Monkey
+from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError
 
 __author__ = "itay.mizeretz"
 
@@ -145,7 +145,10 @@ class EdgeService:
         from_id = edge["from"]
         to_id = edge["to"]
 
-        from_label = Monkey.get_label_by_id(from_id)
+        try:
+            from_label = Monkey.get_label_by_id(from_id)
+        except MonkeyNotFoundError:
+            from_label = node_service.get_node_by_id(from_id)['domain_name']
 
         if to_id == ObjectId("000000000000000000000000"):
             to_label = 'MonkeyIsland'
diff --git a/monkey/monkey_island/cc/services/node.py b/monkey/monkey_island/cc/services/node.py
index 3109f7a78..feb42cc24 100644
--- a/monkey/monkey_island/cc/services/node.py
+++ b/monkey/monkey_island/cc/services/node.py
@@ -207,6 +207,42 @@ class NodeService:
             })
         return mongo.db.node.find_one({"_id": new_node_insert_result.inserted_id})
 
+    @staticmethod
+    def create_node_from_bootloader_telem(bootloader_telem):
+        new_node_insert_result = mongo.db.node.insert_one(
+            {
+                "ip_addresses": bootloader_telem['ips'],
+                "domain_name": bootloader_telem['hostname'],
+                "exploited": False,
+                "creds": [],
+                "os":
+                    {
+                        "type": bootloader_telem['system'],
+                        "version": bootloader_telem['os_version']
+                    }
+            })
+        return mongo.db.node.find_one({"_id": new_node_insert_result.inserted_id})
+
+    @staticmethod
+    def get_or_create_node_from_bootloader_telem(bootloader_telem):
+        new_node = mongo.db.node.find_one({"domain_name": bootloader_telem['hostname'],
+                                           "ip_addresses": bootloader_telem['ips']})
+        if new_node is None:
+            new_node = NodeService.create_node_from_bootloader_telem(bootloader_telem)
+            if bootloader_telem['tunnel']:
+                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']),
+                                           'exploited': (not bool(bootloader_telem['tunnel'])),
+                                           'ip_address': bootloader_telem['ips'][0],
+                                           'group': 'island'}},
+                                 upsert=False)
+
+        return new_node
+
     @staticmethod
     def get_or_create_node(ip_address, domain_name=''):
         new_node = mongo.db.node.find_one({"ip_addresses": ip_address})