def test_get_or_create_edge(self):
self.fail_if_not_testing_env()
self.clean_edge_db()
src_id = ObjectId()
dst_id = ObjectId()
test_edge1 = EdgeService.get_or_create_edge(src_id, dst_id,
"Mock label 1",
"Mock label 2")
self.assertEqual(test_edge1.src_node_id, src_id)
self.assertEqual(test_edge1.dst_node_id, dst_id)
self.assertFalse(test_edge1.exploited)
self.assertFalse(test_edge1.tunnel)
self.assertListEqual(test_edge1.scans, [])
self.assertListEqual(test_edge1.exploits, [])
self.assertEqual(test_edge1.src_label, "Mock label 1")
self.assertEqual(test_edge1.dst_label, "Mock label 2")
self.assertIsNone(test_edge1.group)
self.assertIsNone(test_edge1.domain_name)
self.assertIsNone(test_edge1.ip_address)
EdgeService.get_or_create_edge(src_id, dst_id, "Mock label 1",
"Mock label 2")
self.assertEqual(len(Edge.objects()), 1)
def test_get_displayed_edges_by_to(self):
dst_id = ObjectId()
src_id = ObjectId()
EdgeService.get_or_create_edge(src_id, dst_id, "Ubuntu-4ubuntu2.8", "Ubuntu-4ubuntu2.8")
src_id2 = ObjectId()
EdgeService.get_or_create_edge(src_id2, dst_id, "Ubuntu-4ubuntu3.2", "Ubuntu-4ubuntu2.8")
displayed_edges = DisplayedEdgeService.get_displayed_edges_by_dst(str(dst_id))
assert len(displayed_edges) == 2
def test_get_displayed_edges_by_to(self):
self.clean_edge_db()
dst_id = ObjectId()
src_id = ObjectId()
EdgeService.get_or_create_edge(src_id, dst_id, "Ubuntu-4ubuntu2.8",
"Ubuntu-4ubuntu2.8")
src_id2 = ObjectId()
EdgeService.get_or_create_edge(src_id2, dst_id, "Ubuntu-4ubuntu3.2",
"Ubuntu-4ubuntu2.8")
displayed_edges = DisplayedEdgeService.get_displayed_edges_by_dst(
str(dst_id))
self.assertEqual(len(displayed_edges), 2)
def test_edge_to_displayed_edge(self):
src_node_id = ObjectId()
dst_node_id = ObjectId()
edge = EdgeService(src_node_id=src_node_id,
dst_node_id=dst_node_id,
scans=SCAN_DATA_MOCK,
exploits=EXPLOIT_DATA_MOCK,
exploited=True,
domain_name=None,
ip_address="10.2.2.2",
dst_label="Ubuntu-4ubuntu2.8",
src_label="Ubuntu-4ubuntu3.2")
displayed_edge = DisplayedEdgeService.edge_to_displayed_edge(edge)
self.assertEqual(displayed_edge['to'], dst_node_id)
self.assertEqual(displayed_edge['from'], src_node_id)
self.assertEqual(displayed_edge['ip_address'], "10.2.2.2")
self.assertListEqual(displayed_edge['services'],
["tcp-8088: unknown", "tcp-22: ssh"])
self.assertEqual(displayed_edge['os'], {
"type": "linux",
"version": "Ubuntu-4ubuntu2.8"
})
self.assertEqual(displayed_edge['exploits'], EXPLOIT_DATA_MOCK)
self.assertEqual(
displayed_edge['_label'],
"Ubuntu-4ubuntu3.2 " + RIGHT_ARROW + " Ubuntu-4ubuntu2.8")
self.assertEqual(displayed_edge['group'], "exploited")
return displayed_edge
def get_or_create_node_from_bootloader_telem(bootloader_telem: Dict, will_monkey_run: bool) -> Dict:
if is_local_ips(bootloader_telem['ips']):
raise NodeCreationException("Bootloader ran on island, no need to create new node.")
new_node = mongo.db.node.find_one({"ip_addresses": {"$in": bootloader_telem['ips']}})
# Temporary workaround to not create a node after monkey finishes
monkey_node = mongo.db.monkey.find_one({"ip_addresses": {"$in": bootloader_telem['ips']}})
if monkey_node:
# Don't create new node, monkey node is already present
return monkey_node
if new_node is None:
new_node = NodeService.create_node_from_bootloader_telem(bootloader_telem, will_monkey_run)
if bootloader_telem['tunnel']:
dst_node = NodeService.get_node_or_monkey_by_ip(bootloader_telem['tunnel'])
else:
dst_node = NodeService.get_monkey_island_node()
src_label = NodeService.get_label_for_endpoint(new_node['_id'])
dst_label = NodeService.get_label_for_endpoint(dst_node['id'])
edge = EdgeService.get_or_create_edge(src_node_id=new_node['_id'],
dst_node_id=dst_node['id'],
src_label=src_label,
dst_label=dst_label)
edge.tunnel = bool(bootloader_telem['tunnel'])
edge.ip_address = bootloader_telem['ips'][0]
edge.group = NodeStates.get_by_keywords(['island']).value
edge.save()
return new_node
def test_edge_to_displayed_edge(self):
src_node_id = ObjectId()
dst_node_id = ObjectId()
edge = EdgeService(
src_node_id=src_node_id,
dst_node_id=dst_node_id,
scans=SCAN_DATA_MOCK,
exploits=EXPLOIT_DATA_MOCK,
exploited=True,
domain_name=None,
ip_address="10.2.2.2",
dst_label="Ubuntu-4ubuntu2.8",
src_label="Ubuntu-4ubuntu3.2",
)
displayed_edge = DisplayedEdgeService.edge_to_displayed_edge(edge)
assert displayed_edge["to"] == dst_node_id
assert displayed_edge["from"] == src_node_id
assert displayed_edge["ip_address"] == "10.2.2.2"
assert displayed_edge["services"] == [
"tcp-8088: unknown", "tcp-22: ssh"
]
assert displayed_edge["os"] == {
"type": "linux",
"version": "Ubuntu-4ubuntu2.8"
}
assert displayed_edge["exploits"] == EXPLOIT_DATA_MOCK
assert displayed_edge[
"_label"] == "Ubuntu-4ubuntu3.2 " + RIGHT_ARROW + " Ubuntu-4ubuntu2.8"
assert displayed_edge["group"] == "exploited"
return displayed_edge
def _get_infected_island_net_edges(monkey_island_monkey):
existing_ids = [
x.src_node_id for x in EdgeService.get_by_dst_node(
dst_node_id=monkey_island_monkey["_id"])
]
monkey_ids = [
x.id for x in Monkey.objects()
if ("tunnel" not in x) and (x.id not in existing_ids) and (
x.id != monkey_island_monkey["_id"])
]
edges = []
count = 0
for monkey_id in monkey_ids:
count += 1
# generating fake ID, because front end requires unique ID's for each edge. Collision
# improbable
fake_id = ObjectId(hex(count)[2:].zfill(24))
src_label = NodeService.get_label_for_endpoint(monkey_id)
dst_label = NodeService.get_label_for_endpoint(
monkey_island_monkey["_id"])
edge = DisplayedEdgeService.generate_pseudo_edge(
edge_id=fake_id,
src_node_id=monkey_id,
dst_node_id=monkey_island_monkey["_id"],
src_label=src_label,
dst_label=dst_label,
)
edges.append(edge)
return edges
def test_get_edge_group(self):
edge = Edge(src_node_id=ObjectId(),
dst_node_id=ObjectId(),
exploited=True)
assert "exploited" == EdgeService.get_group(edge)
edge.exploited = False
edge.tunnel = True
assert "tunnel" == EdgeService.get_group(edge)
edge.tunnel = False
edge.exploits.append(["mock_exploit_data"])
assert "scan" == EdgeService.get_group(edge)
edge.exploits = []
edge.scans = []
assert "empty" == EdgeService.get_group(edge)
def test_get_edge_group(self):
edge = Edge(src_node_id=ObjectId(),
dst_node_id=ObjectId(),
exploited=True)
self.assertEqual("exploited", EdgeService.get_group(edge))
edge.exploited = False
edge.tunnel = True
self.assertEqual("tunnel", EdgeService.get_group(edge))
edge.tunnel = False
edge.exploits.append(["mock_exploit_data"])
self.assertEqual("scan", EdgeService.get_group(edge))
edge.exploits = []
edge.scans = []
self.assertEqual("empty", EdgeService.get_group(edge))
def edge_to_net_edge(edge: EdgeService):
return {
"id": edge.id,
"from": edge.src_node_id,
"to": edge.dst_node_id,
"group": edge.get_group(),
"src_label": edge.src_label,
"dst_label": edge.dst_label,
}
def unset_all_monkey_tunnels(monkey_id):
mongo.db.monkey.update({"_id": monkey_id}, {'$unset': {
'tunnel': ''
}},
upsert=False)
edges = EdgeService.get_tunnel_edges_by_src(monkey_id)
for edge in edges:
edge.disable_tunnel()
def test_get_or_create_edge(self):
src_id = ObjectId()
dst_id = ObjectId()
test_edge1 = EdgeService.get_or_create_edge(src_id, dst_id,
"Mock label 1",
"Mock label 2")
assert test_edge1.src_node_id == src_id
assert test_edge1.dst_node_id == dst_id
assert not test_edge1.exploited
assert not test_edge1.tunnel
assert test_edge1.scans == []
assert test_edge1.exploits == []
assert test_edge1.src_label == "Mock label 1"
assert test_edge1.dst_label == "Mock label 2"
assert test_edge1.group is None
assert test_edge1.domain_name is None
assert test_edge1.ip_address is None
EdgeService.get_or_create_edge(src_id, dst_id, "Mock label 1",
"Mock label 2")
assert len(Edge.objects()) == 1
def get_edge_by_scan_or_exploit_telemetry(telemetry_json):
dst_ip = telemetry_json['data']['machine']['ip_addr']
dst_domain_name = telemetry_json['data']['machine']['domain_name']
src_monkey = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid'])
dst_node = NodeService.get_monkey_by_ip(dst_ip)
if dst_node is None:
dst_node = NodeService.get_or_create_node(dst_ip, dst_domain_name)
src_label = NodeService.get_label_for_endpoint(src_monkey["_id"])
dst_label = NodeService.get_label_for_endpoint(dst_node["_id"])
return EdgeService.get_or_create_edge(src_monkey["_id"], dst_node["_id"],
src_label, dst_label)
def set_monkey_tunnel(monkey_id, tunnel_host_ip):
tunnel_host_id = NodeService.get_monkey_by_ip(tunnel_host_ip)["_id"]
NodeService.unset_all_monkey_tunnels(monkey_id)
mongo.db.monkey.update(
{'_id': monkey_id},
{'$set': {'tunnel': tunnel_host_id}},
upsert=False)
monkey_label = NodeService.get_label_for_endpoint(monkey_id)
tunnel_host_label = NodeService.get_label_for_endpoint(tunnel_host_id)
tunnel_edge = EdgeService.get_or_create_edge(src_node_id=monkey_id,
dst_node_id=tunnel_host_id,
src_label=monkey_label,
dst_label=tunnel_host_label)
tunnel_edge.tunnel = True
tunnel_edge.ip_address = tunnel_host_ip
tunnel_edge.save()
def edge_to_displayed_edge(edge: EdgeService, for_report=False):
services = []
os = {}
if len(edge.scans) > 0:
services = DisplayedEdgeService.services_to_displayed_services(
edge.scans[-1]["data"]["services"], for_report)
os = edge.scans[-1]["data"]["os"]
displayed_edge = DisplayedEdgeService.edge_to_net_edge(edge)
displayed_edge["ip_address"] = edge.ip_address
displayed_edge["services"] = services
displayed_edge["os"] = os
# we need to deepcopy all mutable edge properties, because weak-reference link is made otherwise,
# which is destroyed after method is exited and causes an error later.
displayed_edge["exploits"] = deepcopy(edge.exploits)
displayed_edge["_label"] = edge.get_label()
return displayed_edge
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"),
)
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})
return {"id": new_monkey_id}
def get_displayed_edges_by_dst(dst_id: str, for_report=False):
edges = EdgeService.get_by_dst_node(dst_node_id=ObjectId(dst_id))
return [
DisplayedEdgeService.edge_to_displayed_edge(edge, for_report)
for edge in edges
]
def get_displayed_edge_by_id(edge_id: str, for_report=False):
edge = EdgeService.get_edge_by_id(ObjectId(edge_id))
displayed_edge = DisplayedEdgeService.edge_to_displayed_edge(
edge, for_report)
return displayed_edge
def _get_standard_net_edges():
return [DisplayedEdgeService.edge_to_net_edge(x) for x in EdgeService.get_all_edges()]