forked from p15670423/monkey
modification on tunnel handling
This commit is contained in:
parent
78e989e5fc
commit
9fef5f2da2
|
@ -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:
|
||||
|
|
|
@ -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'])
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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}},
|
||||
|
|
Loading…
Reference in New Issue