diff --git a/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py b/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py index 8c91ac739..d84ef2f8b 100644 --- a/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py +++ b/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py @@ -4,7 +4,8 @@ import flask_restful from flask import make_response, request from monkey_island.cc.resources.auth.auth import jwt_required -from monkey_island.cc.services.infection_lifecycle import set_stop_all +from monkey_island.cc.resources.utils.semaphores import AGENT_KILLING_SEMAPHORE +from monkey_island.cc.services.infection_lifecycle import set_stop_all, was_monkey_killed class StopAllAgents(flask_restful.Resource): diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py index 3853b58ed..ffae35fba 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/monkey/monkey_island/cc/resources/monkey.py @@ -8,6 +8,7 @@ from flask import request from monkey_island.cc.database import mongo from monkey_island.cc.models.monkey_ttl import create_monkey_ttl_document from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore +from monkey_island.cc.resources.utils.semaphores import AGENT_KILLING_SEMAPHORE from monkey_island.cc.server_utils.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS from monkey_island.cc.services.config import ConfigService from monkey_island.cc.services.edge.edge import EdgeService @@ -66,97 +67,104 @@ class Monkey(flask_restful.Resource): # Called on monkey wakeup to initialize local configuration @TestTelemStore.store_exported_telem def post(self, **kw): - monkey_json = json.loads(request.data) - monkey_json["creds"] = [] - monkey_json["dead"] = False - 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() - - ConfigService.save_initial_config_if_needed() - - # if new monkey telem, change config according to "new monkeys" config. - db_monkey = mongo.db.monkey.find_one({"guid": monkey_json["guid"]}) - - # Update monkey configuration - new_config = ConfigService.get_flat_config(False, False) - monkey_json["config"] = monkey_json.get("config", {}) - monkey_json["config"].update(new_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_category": {"$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"), - ) + with AGENT_KILLING_SEMAPHORE: + monkey_json = json.loads(request.data) + monkey_json["creds"] = [] + monkey_json["dead"] = False + if "keepalive" in monkey_json: + monkey_json["keepalive"] = dateutil.parser.parse(monkey_json["keepalive"]) 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_category": {"$eq": "exploit"}, - "data.result": {"$eq": True}, - "data.machine.ip_addr": {"$in": monkey_json["ip_addresses"]}, - } + monkey_json["keepalive"] = datetime.now() + + monkey_json["modifytime"] = datetime.now() + + ConfigService.save_initial_config_if_needed() + + # if new monkey telem, change config according to "new monkeys" config. + db_monkey = mongo.db.monkey.find_one({"guid": monkey_json["guid"]}) + + # Update monkey configuration + new_config = ConfigService.get_flat_config(False, False) + monkey_json["config"] = monkey_json.get("config", {}) + monkey_json["config"].update(new_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_category": {"$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_category": {"$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") + + ttl = create_monkey_ttl_document(DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS) + monkey_json["ttl_ref"] = ttl.id + + 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"] + EdgeService.update_all_dst_nodes( + old_dst_node_id=node_id, new_dst_node_id=new_monkey_id ) - ] + for creds in existing_node["creds"]: + NodeService.add_credentials_to_monkey(new_monkey_id, creds) + mongo.db.node.remove({"_id": node_id}) - 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") - - ttl = create_monkey_ttl_document(DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS) - monkey_json["ttl_ref"] = ttl.id - - 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"] - EdgeService.update_all_dst_nodes(old_dst_node_id=node_id, new_dst_node_id=new_monkey_id) - for creds in existing_node["creds"]: - NodeService.add_credentials_to_monkey(new_monkey_id, creds) - mongo.db.node.remove({"_id": node_id}) - - return {"id": new_monkey_id} + return {"id": new_monkey_id} diff --git a/monkey/monkey_island/cc/resources/utils/semaphores.py b/monkey/monkey_island/cc/resources/utils/semaphores.py new file mode 100644 index 000000000..97f36e441 --- /dev/null +++ b/monkey/monkey_island/cc/resources/utils/semaphores.py @@ -0,0 +1,5 @@ +from gevent.lock import BoundedSemaphore + +# Semaphore avoids race condition between monkeys +# being marked dead and monkey waking up as alive +AGENT_KILLING_SEMAPHORE = BoundedSemaphore()