bug fixes on monkey wakeup

This commit is contained in:
Barak Hoffer 2015-10-08 13:37:38 +03:00
parent a19f820ec8
commit ef6474f7b6
1 changed files with 30 additions and 14 deletions

View File

@ -66,11 +66,19 @@ class Monkey(restful.Resource):
def patch(self, guid): def patch(self, guid):
monkey_json = json.loads(request.data); monkey_json = json.loads(request.data);
update = {"$set" : {'modifytime':datetime.now(), 'keepalive':datetime.now()}} update = {"$set" : {'modifytime':datetime.now()}}
if monkey_json.has_key('keepalive'):
update['$set']['keepalive'] = dateutil.parser.parse(monkey_json['keepalive'])
else:
update['$set']['keepalive'] = datetime.now()
if monkey_json.has_key('config'): if monkey_json.has_key('config'):
update['$set']['config'] = monkey_json['config'] update['$set']['config'] = monkey_json['config']
return mongo.db.monkey.update({"guid": guid},\ if monkey_json.has_key('tunnel'):
update,\ update['$set']['tunnel'] = monkey_json['tunnel']
return mongo.db.monkey.update({"guid": guid},
update,
upsert=False) upsert=False)
def post(self, **kw): def post(self, **kw):
@ -82,6 +90,21 @@ class Monkey(restful.Resource):
monkey_json['modifytime'] = datetime.now() monkey_json['modifytime'] = datetime.now()
# if new monkey, change config according to general config.
db_monkey = mongo.db.monkey.find_one({"guid": monkey_json["guid"]})
if not db_monkey:
general_config = mongo.db.config.find_one({'name' : 'generalconfig'}) or {}
monkey_json['config'] = monkey_json.get('config', {})
monkey_json['config'].update(general_config)
else:
db_config = db_monkey.get('config', {})
if db_config.has_key('current_server'):
del db_config['current_server']
monkey_json.get('config', {}).update(db_config)
if not monkey_json.has_key('parent') and db_monkey.get('parent'):
monkey_json['parent'] = db_monkey.get('parent')
# try to find new monkey parent # try to find new monkey parent
parent = monkey_json.get('parent') parent = monkey_json.get('parent')
if (not parent or parent == monkey_json.get('guid')) and monkey_json.has_key('ip_addresses'): if (not parent or parent == monkey_json.get('guid')) and monkey_json.has_key('ip_addresses'):
@ -90,16 +113,9 @@ class Monkey(restful.Resource):
if 1 == len(exploit_telem): if 1 == len(exploit_telem):
monkey_json['parent'] = exploit_telem[0].get('monkey_guid') monkey_json['parent'] = exploit_telem[0].get('monkey_guid')
# if new monkey, change config according to general config. return mongo.db.monkey.update({"guid": monkey_json["guid"]},
db_monkey = mongo.db.monkey.find_one({"guid": monkey_json["guid"]}) {"$set" : monkey_json},
if not db_monkey: upsert=True)
general_config = mongo.db.config.find_one({'name' : 'generalconfig'}) or {}
monkey_json['config']= monkey_json.get('config', {})
monkey_json['config'].update(general_config)
else:
monkey_json.update(db_monkey)
return mongo.db.monkey.update({"guid": monkey_json["guid"]}, {"$set" : monkey_json}, upsert=True)
class Telemetry(restful.Resource): class Telemetry(restful.Resource):
def get(self, **kw): def get(self, **kw):