diff --git a/api/project_aggregation.py b/api/project_aggregation.py index b6420d1..bdf2324 100644 --- a/api/project_aggregation.py +++ b/api/project_aggregation.py @@ -90,13 +90,10 @@ async def get_projects_vul_data(request_data: dict, db=Depends(get_mongo_db), _: @router.post("/project/subdomain/data") -async def get_projects_vul_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)): +async def get_projects_subdomain_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)): filter = request_data.get("filter", {}) project_id = filter["project"][0] project_query = {} - host_filter = "" - if "host" in filter: - host_filter = filter["host"] project_query["_id"] = ObjectId(project_id) doc = await db.project.find_one(project_query, {"_id": 0, "root_domains": 1}) if not doc or "root_domains" not in doc: @@ -125,3 +122,171 @@ async def get_projects_vul_data(request_data: dict, db=Depends(get_mongo_db), _: 'list': results } } + + +@router.post("/project/port/data") +async def get_projects_vul_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)): + query = await get_search_query("asset", request_data) + if query == "": + return {"message": "Search condition parsing error", "code": 500} + pipeline = [ + { + "$match": query # 添加搜索条件 + }, + { + "$group": + { + "_id": "$port", + "count": + { + "$sum": 1 + } + } + }, + { + "$sort": {"count": -1} + } + ] + result = await db['asset'].aggregate(pipeline).to_list(None) + + async def fetch_asset_data(r): + tmp_result = { + "port": r['_id'], + "id": generate_random_string(5) + } + + query_copy = query.copy() + query_copy["port"] = r['_id'] + + cursor: AsyncIOMotorCursor = db['asset'].find(query_copy, { + "_id": 0, + "id": {"$toString": "$_id"}, + "host": 1, + "ip": 1, + "type": 1, + "timestamp": 1, + "domain": 1, + "protocol": 1 + }).sort([("timestamp", DESCENDING)]) + + asset_result = await cursor.to_list(length=None) + children_list = [] + + for asset in asset_result: + children_data = {} + + if asset['type'] == "other": + children_data['protocol'] = asset['protocol'] + children_data['domain'] = asset['host'] + children_data['ip'] = asset['ip'] + else: + children_data['protocol'] = asset['type'] + children_data['domain'] = asset['domain'] + children_data['ip'] = asset['host'] + + children_data['timestamp'] = asset['timestamp'] + children_data['id'] = asset['id'] + children_list.append(children_data) + + if len(children_list) != 0: + tmp_result["children"] = children_list + + return tmp_result + + tasks = [fetch_asset_data(r) for r in result] + result_list = await asyncio.gather(*tasks) + + return { + "code": 200, + "data": { + 'list': result_list + } + } + + +@router.post("/project/service/data") +async def get_projects_service_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)): + query = await get_search_query("asset", request_data) + if query == "": + return {"message": "Search condition parsing error", "code": 500} + + pipeline = [ + { + "$match": query + }, + { + "$group": { + "_id": { + "$cond": { + "if": {"$eq": ["$type", "other"]}, + "then": "$protocol", + "else": "$type" + } + }, + "count": {"$sum": 1} + } + }, + { + "$sort": {"count": -1} + } + ] + + result = await db['asset'].aggregate(pipeline).to_list(None) + + async def fetch_asset_data(r): + tmp_result = { + "service": r['_id'], + "id": generate_random_string(5) + } + + query_copy = query.copy() + if r['_id'] == 'http' or r['_id'] == 'https': + query_copy["type"] = r['_id'] + else: + query_copy["type"] = r['_id'] + if r['_id'] == "": + tmp_result['service'] = 'unknown' + + cursor: AsyncIOMotorCursor = db['asset'].find(query_copy, { + "_id": 0, + "id": {"$toString": "$_id"}, + "host": 1, + "ip": 1, + "type": 1, + "timestamp": 1, + "webServer": 1, + "domain": 1, + "protocol": 1 + }).sort([("timestamp", DESCENDING)]) + + asset_result = await cursor.to_list(length=None) + children_list = [] + + for asset in asset_result: + children_data = {} + + if asset['type'] == "other": + children_data['protocol'] = asset['protocol'] + children_data['domain'] = asset['host'] + children_data['ip'] = asset['ip'] + else: + children_data['service'] = asset['webServer'] + children_data['protocol'] = asset['type'] + children_data['domain'] = asset['domain'] + children_data['ip'] = asset['host'] + + children_data['timestamp'] = asset['timestamp'] + children_data['id'] = asset['id'] + children_list.append(children_data) + + if len(children_list) != 0: + tmp_result["children"] = children_list + + return tmp_result + + tasks = [fetch_asset_data(r) for r in result] + result_list = await asyncio.gather(*tasks) + return { + "code": 200, + "data": result_list + } diff --git a/core/util.py b/core/util.py index f92bd6f..bd09401 100644 --- a/core/util.py +++ b/core/util.py @@ -323,6 +323,7 @@ async def search_to_mongodb(expression_raw, keyword): async def get_search_query(name, request_data): + global tmp_f_q search_query = request_data.get("search", "") search_key_v = { 'sens':{ @@ -425,6 +426,22 @@ async def get_search_query(name, request_data): tmp_or.append({filter_key[f]: v}) if len(tmp_or) != 0: query["$and"].append({"$or": tmp_or}) + fuzzy_query = request_data.get("fq", {}) + fuzzy_query_key = {"sub_host": 'host', "sub_value": "value", "sub_ip": "ip", "port_port": "port", "port_domain":['domain', 'host'], 'port_ip': ['ip', 'host'], 'port_protocol': ['type', 'protocol']} + if fuzzy_query: + if "$and" not in query: + query["$and"] = [] + for q in fuzzy_query: + if fuzzy_query[q] != "": + tmp_f_q = [] + if q in fuzzy_query_key: + if type(fuzzy_query_key[q]) is list: + for key in fuzzy_query_key[q]: + tmp_f_q.append({key: {"$regex": fuzzy_query[q]}}) + else: + tmp_f_q.append({fuzzy_query_key[q]: {"$regex": fuzzy_query[q]}}) + if len(tmp_f_q) != 0: + query["$and"].append({"$or": tmp_f_q}) if "$and" in query: if len(query["$and"]) == 0: query.pop("$and") diff --git a/main.py b/main.py index 3d7373c..f201c7b 100644 --- a/main.py +++ b/main.py @@ -76,11 +76,11 @@ async def update(): @app.on_event("startup") async def startup_db_client(): - await update() file_path = os.path.join(os.getcwd(), 'file') if not os.path.exists(file_path): os.makedirs(file_path) await db.create_database() + await update() scheduler.start() jobs = scheduler.get_jobs() find_page_m = False