Add nodes and edges
This commit is contained in:
parent
54f054a4e7
commit
2e5fe25117
|
@ -189,7 +189,8 @@ class Telemetry(restful.Resource):
|
||||||
upsert=False)
|
upsert=False)
|
||||||
else:
|
else:
|
||||||
mongo.db.monkey.update({"guid": telemetry_json['monkey_guid']},
|
mongo.db.monkey.update({"guid": telemetry_json['monkey_guid']},
|
||||||
{'$unset': {'tunnel_guid': ''}, 'modifytime': datetime.now()},
|
{'$unset': {'tunnel_guid': ''},
|
||||||
|
'$set': {'modifytime': datetime.now()}},
|
||||||
upsert=False)
|
upsert=False)
|
||||||
elif telemetry_json.get('telem_type') == 'state':
|
elif telemetry_json.get('telem_type') == 'state':
|
||||||
if telemetry_json['data']['done']:
|
if telemetry_json['data']['done']:
|
||||||
|
@ -200,7 +201,30 @@ class Telemetry(restful.Resource):
|
||||||
mongo.db.monkey.update({"guid": telemetry_json['monkey_guid']},
|
mongo.db.monkey.update({"guid": telemetry_json['monkey_guid']},
|
||||||
{'$set': {'dead': False, 'modifytime': datetime.now()}},
|
{'$set': {'dead': False, 'modifytime': datetime.now()}},
|
||||||
upsert=False)
|
upsert=False)
|
||||||
except:
|
elif telemetry_json.get('telem_type') == 'scan':
|
||||||
|
dst_ip = telemetry_json['data']['machine']['ip_addr']
|
||||||
|
src_monkey = mongo.db.monkey.find_one({"guid": telemetry_json['monkey_guid']})
|
||||||
|
dst_monkey = mongo.db.monkey.find_one({"ip_addresses": dst_ip})
|
||||||
|
if dst_monkey:
|
||||||
|
edge = mongo.db.edges.find_one({"from": src_monkey["_id"], "to": dst_monkey["_id"]})
|
||||||
|
|
||||||
|
if edge is None:
|
||||||
|
edge = self.insert_edge(src_monkey["_id"], dst_monkey["_id"])
|
||||||
|
|
||||||
|
else:
|
||||||
|
dst_node = mongo.db.nodes.find_one({"ip_addresses": dst_ip})
|
||||||
|
if dst_node is None:
|
||||||
|
dst_node_insert_result = mongo.db.nodes.insert_one({"ip_addresses": [dst_ip]})
|
||||||
|
dst_node = mongo.db.nodes.find_one({"_id": dst_node_insert_result.inserted_id})
|
||||||
|
|
||||||
|
edge = mongo.db.edges.find_one({"from": src_monkey["_id"], "to": dst_node["_id"]})
|
||||||
|
|
||||||
|
if edge is None:
|
||||||
|
edge = self.insert_edge(src_monkey["_id"], dst_node["_id"])
|
||||||
|
|
||||||
|
self.add_scan_to_edge(edge, telemetry_json)
|
||||||
|
|
||||||
|
except StandardError as e:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
# Update credentials DB
|
# Update credentials DB
|
||||||
|
@ -217,6 +241,29 @@ class Telemetry(restful.Resource):
|
||||||
|
|
||||||
return mongo.db.telemetry.find_one_or_404({"_id": telem_id})
|
return mongo.db.telemetry.find_one_or_404({"_id": telem_id})
|
||||||
|
|
||||||
|
def add_scan_to_edge(self, edge, telemetry_json):
|
||||||
|
data = telemetry_json['data']['machine']
|
||||||
|
data.pop("ip_addr")
|
||||||
|
new_scan = \
|
||||||
|
{
|
||||||
|
"timestamp": telemetry_json["timestamp"],
|
||||||
|
"data": data,
|
||||||
|
"scanner": telemetry_json['data']['scanner']
|
||||||
|
}
|
||||||
|
mongo.db.edges.update(
|
||||||
|
{"_id": edge["_id"]},
|
||||||
|
{"$push": {"scans": new_scan}}
|
||||||
|
)
|
||||||
|
|
||||||
|
def insert_edge(self, from_id, to_id):
|
||||||
|
edge_insert_result = mongo.db.edges.insert_one(
|
||||||
|
{
|
||||||
|
"from": from_id,
|
||||||
|
"to": to_id,
|
||||||
|
"scans": []
|
||||||
|
})
|
||||||
|
return mongo.db.edges.find_one({"_id": edge_insert_result.inserted_id})
|
||||||
|
|
||||||
|
|
||||||
class LocalRun(restful.Resource):
|
class LocalRun(restful.Resource):
|
||||||
def get(self):
|
def get(self):
|
||||||
|
@ -281,6 +328,8 @@ class Root(restful.Resource):
|
||||||
mongo.db.telemetry.drop()
|
mongo.db.telemetry.drop()
|
||||||
mongo.db.usernames.drop()
|
mongo.db.usernames.drop()
|
||||||
mongo.db.passwords.drop()
|
mongo.db.passwords.drop()
|
||||||
|
mongo.db.nodes.drop()
|
||||||
|
mongo.db.edges.drop()
|
||||||
init_db()
|
init_db()
|
||||||
return {
|
return {
|
||||||
'status': 'OK',
|
'status': 'OK',
|
||||||
|
|
Loading…
Reference in New Issue