Island: add semaphores to avoid race condition

If user kills all monkeys during the waking up of a monkey, all monkeys will get flagged as dead except the one that just woke up

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

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,6 +67,7 @@ class Monkey(flask_restful.Resource):
     # Called on monkey wakeup to initialize local configuration
     @TestTelemStore.store_exported_telem
     def post(self, **kw):
+        with AGENT_KILLING_SEMAPHORE:
             monkey_json = json.loads(request.data)
             monkey_json["creds"] = []
             monkey_json["dead"] = False
@@ -108,7 +110,9 @@ class Monkey(flask_restful.Resource):
                     )
                 else:
                     parent_to_add = (parent, None)
-        elif (not parent or parent == monkey_json.get("guid")) and "ip_addresses" in monkey_json:
+            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(
@@ -139,7 +143,9 @@ class Monkey(flask_restful.Resource):
             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)
+            mongo.db.monkey.update(
+                {"guid": monkey_json["guid"]}, {"$set": monkey_json}, upsert=True
+            )
 
             # Merge existing scanned node with new monkey
 
@@ -154,7 +160,9 @@ class Monkey(flask_restful.Resource):
 
             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)
+                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})

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()