modification on tunnel handling

This commit is contained in:
Itay Mizeretz 2017-09-12 13:48:36 +03:00
parent 78e989e5fc
commit 9fef5f2da2
4 changed files with 34 additions and 7 deletions

View File

@ -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:

View File

@ -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'])

View File

@ -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

View File

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