forked from p15670423/monkey
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
This commit is contained in:
parent
31cdd29edb
commit
e4280660df
|
@ -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):
|
||||
|
|
|
@ -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}
|
||||
|
|
|
@ -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()
|
Loading…
Reference in New Issue