diff --git a/monkey_island/cc/main.py b/monkey_island/cc/main.py index 9641e2ae3..b3234b4d7 100644 --- a/monkey_island/cc/main.py +++ b/monkey_island/cc/main.py @@ -189,7 +189,8 @@ class Telemetry(restful.Resource): upsert=False) else: mongo.db.monkey.update({"guid": telemetry_json['monkey_guid']}, - {'$unset': {'tunnel_guid': ''}, 'modifytime': datetime.now()}, + {'$unset': {'tunnel_guid': ''}, + '$set': {'modifytime': datetime.now()}}, upsert=False) elif telemetry_json.get('telem_type') == 'state': if telemetry_json['data']['done']: @@ -200,7 +201,30 @@ class Telemetry(restful.Resource): mongo.db.monkey.update({"guid": telemetry_json['monkey_guid']}, {'$set': {'dead': False, 'modifytime': datetime.now()}}, 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 # Update credentials DB @@ -217,6 +241,29 @@ class Telemetry(restful.Resource): 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): def get(self): @@ -281,6 +328,8 @@ class Root(restful.Resource): mongo.db.telemetry.drop() mongo.db.usernames.drop() mongo.db.passwords.drop() + mongo.db.nodes.drop() + mongo.db.edges.drop() init_db() return { 'status': 'OK',