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