monkey/monkey_island/cc/resources/monkey.py

import json
from datetime import datetime

import dateutil.parser
from flask import request
import flask_restful

from cc.database import mongo
from cc.services.config import ConfigService
from cc.services.node import NodeService

__author__ = 'Barak'

# TODO: separate logic from interface


class Monkey(flask_restful.Resource):
    def get(self, guid=None, **kw):
        NodeService.update_dead_monkeys()  # refresh monkeys status
        if not guid:
            guid = request.args.get('guid')
        timestamp = request.args.get('timestamp')

        if guid:
            monkey_json = mongo.db.monkey.find_one_or_404({"guid": guid})
            return monkey_json
        else:
            result = {'timestamp': datetime.now().isoformat()}
            find_filter = {}
            if timestamp is not None:
                find_filter['modifytime'] = {'$gt': dateutil.parser.parse(timestamp)}
            result['objects'] = [x for x in mongo.db.monkey.find(find_filter)]
            return result

    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 'config_error' in monkey_json:
            update['$set']['config_error'] = monkey_json['config_error']

        if 'tunnel' in monkey_json:
            tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "")
            NodeService.set_monkey_tunnel(monkey["_id"], tunnel_host_ip)

        return mongo.db.monkey.update({"_id": monkey["_id"]}, update, upsert=False)

    def post(self, **kw):
        monkey_json = json.loads(request.data)
        monkey_json['creds'] = {}
        if 'keepalive' in monkey_json:
            monkey_json['keepalive'] = dateutil.parser.parse(monkey_json['keepalive'])
        else:
            monkey_json['keepalive'] = datetime.now()

        monkey_json['modifytime'] = datetime.now()

        # if new monkey telem, change config according to "new monkeys" config.
        db_monkey = mongo.db.monkey.find_one({"guid": monkey_json["guid"]})
        if not db_monkey:
            new_config = ConfigService.get_flat_config()
            monkey_json['config'] = monkey_json.get('config', {})
            monkey_json['config'].update(new_config)
        else:
            db_config = db_monkey.get('config', {})
            if 'current_server' in db_config:
                del db_config['current_server']
            monkey_json.get('config', {}).update(db_config)

        # try to find new monkey parent
        parent = monkey_json.get('parent')
        parent_to_add = (monkey_json.get('guid'), None)  # default values in case of manual run
        if parent and parent != monkey_json.get('guid'):  # current parent is known
            exploit_telem = [x for x in
                             mongo.db.telemetry.find({'telem_type': {'$eq': 'exploit'}, 'data.result': {'$eq': True},
                                                      'data.machine.ip_addr': {'$in': monkey_json['ip_addresses']},
                                                      'monkey_guid': {'$eq': parent}})]
            if 1 == len(exploit_telem):
                parent_to_add = (exploit_telem[0].get('monkey_guid'), exploit_telem[0].get('data').get('exploiter'))
            else:
                parent_to_add = (parent, None)
        elif (not parent or parent == monkey_json.get('guid')) and 'ip_addresses' in  monkey_json:
            exploit_telem = [x for x in
                             mongo.db.telemetry.find({'telem_type': {'$eq': 'exploit'}, 'data.result': {'$eq': True},
                                                      'data.machine.ip_addr': {'$in': monkey_json['ip_addresses']}})]

            if 1 == len(exploit_telem):
                parent_to_add = (exploit_telem[0].get('monkey_guid'), exploit_telem[0].get('data').get('exploiter'))

        if not db_monkey:
            monkey_json['parent'] = [parent_to_add]
        else:
            monkey_json['parent'] = db_monkey.get('parent') + [parent_to_add]

        tunnel_host_ip = None
        if 'tunnel' in monkey_json:
            tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "")
            monkey_json.pop('tunnel')

        mongo.db.monkey.update({"guid": monkey_json["guid"]},
                               {"$set": monkey_json},
                               upsert=True)

        # Merge existing scanned node with new monkey

        new_monkey_id = mongo.db.monkey.find_one({"guid": monkey_json["guid"]})["_id"]

        if tunnel_host_ip is not None:
            NodeService.set_monkey_tunnel(new_monkey_id, tunnel_host_ip)

        existing_node = mongo.db.node.find_one({"ip_addresses": {"$in": monkey_json["ip_addresses"]}})

        if existing_node:
            node_id = existing_node["_id"]
            for edge in mongo.db.edge.find({"to": node_id}):
                mongo.db.edge.update({"_id": edge["_id"]}, {"$set": {"to": new_monkey_id}})
            for user in existing_node['creds']:
                NodeService.add_credentials_to_monkey(new_monkey_id, user, existing_node['creds'][user])
            mongo.db.node.remove({"_id": node_id})

        return {"id": new_monkey_id}