modification on tunnel handling

monkey_island/cc/resources/monkey.py

@@ -6,9 +6,12 @@ from flask import request
 import flask_restful
 
 from cc.database import mongo
+from cc.services.node import NodeService
 
 __author__ = 'Barak'
 
+# TODO: separate logic from interface
+
 
 def update_dead_monkeys():
     # Update dead monkeys only if no living monkey transmitted keepalive in the last 10 minutes
@@ -45,19 +48,22 @@ class Monkey(flask_restful.Resource):
     def patch(self, guid):
         monkey_json = json.loads(request.data)
         update = {"$set": {'modifytime': datetime.now()}}
-
+        monkey = NodeService.get_monkey_by_guid(guid)
         if 'keepalive' in monkey_json:
             update['$set']['keepalive'] = dateutil.parser.parse(monkey_json['keepalive'])
         else:
             update['$set']['keepalive'] = datetime.now()
         if 'config' in monkey_json:
             update['$set']['config'] = monkey_json['config']
-        if 'tunnel' in monkey_json:
-            update['$set']['tunnel'] = monkey_json['tunnel']
         if 'config_error' in monkey_json:
             update['$set']['config_error'] = monkey_json['config_error']
 
-        return mongo.db.monkey.update({"guid": guid}, update, upsert=False)
+        if 'tunnel' in monkey_json:
+            host = monkey_json['tunnel'].split(":")[-2].replace("//", "")
+            tunnel_host_id = NodeService.get_monkey_by_ip(host)["_id"]
+            NodeService.set_monkey_tunnel(monkey["_id"], tunnel_host_id)
+
+        return mongo.db.monkey.update({"_id": monkey["_id"]}, update, upsert=False)
 
     def post(self, **kw):
         monkey_json = json.loads(request.data)
@@ -105,6 +111,12 @@ class Monkey(flask_restful.Resource):
         else:
             monkey_json['parent'] = db_monkey.get('parent') + [parent_to_add]
 
+        tunnel_host_id = None
+        if 'tunnel' in monkey_json:
+            host = monkey_json['tunnel'].split(":")[-2].replace("//", "")
+            tunnel_host_id = NodeService.get_monkey_by_ip(host)["_id"]
+            monkey_json.pop('tunnel')
+
         mongo.db.monkey.update({"guid": monkey_json["guid"]},
                                {"$set": monkey_json},
                                upsert=True)
@@ -113,6 +125,9 @@ class Monkey(flask_restful.Resource):
 
         new_monkey_id = mongo.db.monkey.find_one({"guid": monkey_json["guid"]})["_id"]
 
+        if tunnel_host_id is not None:
+            NodeService.set_monkey_tunnel(new_monkey_id, tunnel_host_id)
+
         existing_node = mongo.db.node.find_one({"ip_addresses": {"$in": monkey_json["ip_addresses"]}})
 
         if existing_node:

monkey_island/cc/resources/telemetry.py

@@ -61,11 +61,12 @@ class Telemetry(flask_restful.Resource):
 
     def process_tunnel_telemetry(self, telemetry_json):
         monkey_id = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid'])["_id"]
-        NodeService.unset_all_monkey_tunnels(monkey_id)
         if telemetry_json['data']:
             host = telemetry_json['data'].split(":")[-2].replace("//", "")
             tunnel_host_id = NodeService.get_monkey_by_ip(host)["_id"]
             NodeService.set_monkey_tunnel(monkey_id, tunnel_host_id)
+        else:
+            NodeService.unset_all_monkey_tunnels(monkey_id)
 
     def process_state_telemetry(self, telemetry_json):
         monkey = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid'])

monkey_island/cc/services/edge.py

@@ -115,7 +115,8 @@ class EdgeService:
     def get_monkey_island_pseudo_edges():
         edges = []
         monkey_ids = [x["_id"] for x in mongo.db.monkey.find({}) if "tunnel" not in x]
-        # TODO: find solution for these ids.
+        # 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

monkey_island/cc/services/node.py

@@ -1,4 +1,4 @@
-import datetime
+from datetime import datetime
 from bson import ObjectId
 
 from cc.database import mongo
@@ -121,6 +121,11 @@ class NodeService:
 
     @staticmethod
     def unset_all_monkey_tunnels(monkey_id):
+        mongo.db.monkey.update(
+            {"_id": monkey_id},
+            {'$unset': {'tunnel': ''}},
+            upsert=False)
+
         mongo.db.edge.update(
             {"from": monkey_id, 'tunnel': True},
             {'$set': {'tunnel': False}},
@@ -128,6 +133,11 @@ class NodeService:
 
     @staticmethod
     def set_monkey_tunnel(monkey_id, tunnel_host_id):
+        NodeService.unset_all_monkey_tunnels(monkey_id)
+        mongo.db.monkey.update(
+            {"_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}},