From 9fef5f2da2bcb59081d139ba6839d752a6340ba6 Mon Sep 17 00:00:00 2001 From: Itay Mizeretz Date: Tue, 12 Sep 2017 13:48:36 +0300 Subject: [PATCH] modification on tunnel handling --- monkey_island/cc/resources/monkey.py | 23 +++++++++++++++++++---- monkey_island/cc/resources/telemetry.py | 3 ++- monkey_island/cc/services/edge.py | 3 ++- monkey_island/cc/services/node.py | 12 +++++++++++- 4 files changed, 34 insertions(+), 7 deletions(-) diff --git a/monkey_island/cc/resources/monkey.py b/monkey_island/cc/resources/monkey.py index 90eddd015..e59a82b43 100644 --- a/monkey_island/cc/resources/monkey.py +++ b/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: diff --git a/monkey_island/cc/resources/telemetry.py b/monkey_island/cc/resources/telemetry.py index bfc38c582..7f0510793 100644 --- a/monkey_island/cc/resources/telemetry.py +++ b/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']) diff --git a/monkey_island/cc/services/edge.py b/monkey_island/cc/services/edge.py index e1811bef9..c4921bc8d 100644 --- a/monkey_island/cc/services/edge.py +++ b/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 diff --git a/monkey_island/cc/services/node.py b/monkey_island/cc/services/node.py index 6046726fc..db5a88abe 100644 --- a/monkey_island/cc/services/node.py +++ b/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}},