diff --git a/api/project.py b/api/project.py index a07a58a..df1ef7c 100644 --- a/api/project.py +++ b/api/project.py @@ -1,6 +1,6 @@ import time import traceback - +import asyncio from bson import ObjectId from fastapi import APIRouter, Depends, BackgroundTasks @@ -24,56 +24,59 @@ async def get_projects_data(request_data: dict, db=Depends(get_mongo_db), _: dic search_query = request_data.get("search", "") page_index = request_data.get("pageIndex", 1) page_size = request_data.get("pageSize", 10) - if search_query == "": - query = {} - else: - query = { - "$or": [ - {"name": {"$regex": search_query, "$options": "i"}}, - {"target": {"$regex": search_query, "$options": "i"}} - ] - } - tag_num = {} - tag_result = await db.project.aggregate([{ - "$group": { - "_id": "$tag", - "count": {"$sum": 1} - } - }]).to_list(None) - all_num = 0 - for tag in tag_result: - tag_num[tag["_id"]] = tag["count"] - all_num += tag["count"] - result_list = {} + + query = { + "$or": [ + {"name": {"$regex": search_query, "$options": "i"}}, + {"target": {"$regex": search_query, "$options": "i"}} + ] + } if search_query else {} + + # 获取标签统计信息 + tag_result = await db.project.aggregate([ + {"$group": {"_id": "$tag", "count": {"$sum": 1}}}, + {"$sort": {"count": -1}} + ]).to_list(None) + + tag_num = {tag["_id"]: tag["count"] for tag in tag_result} + all_num = sum(tag_num.values()) tag_num["All"] = all_num - for tag in tag_num: - if tag != "All": - tag_query = { - "$and": [ - query, - {"tag": tag} - ] - } - else: - tag_query = query - cursor = db.project.find(tag_query, {"_id": 0, - "id": {"$toString": "$_id"}, - "name": 1, - "logo": 1, - "AssetCount": 1, - "tag": 1 - }).sort("AssetCount", -1).skip((page_index - 1) * page_size).limit( - page_size) + + result_list = {} + + async def fetch_projects(tag, tag_query): + cursor = db.project.find(tag_query, { + "_id": 0, + "id": {"$toString": "$_id"}, + "name": 1, + "logo": 1, + "AssetCount": 1, + "tag": 1 + }).sort("AssetCount", -1).skip((page_index - 1) * page_size).limit(page_size) + results = await cursor.to_list(length=None) - result_list[tag] = [] for result in results: result["AssetCount"] = result.get("AssetCount", 0) - result_list[tag].append(result) background_tasks.add_task(update_project_count, id=result["id"]) + return results + + fetch_tasks = [] + for tag in tag_num: + if tag != "All": + tag_query = {"$and": [query, {"tag": tag}]} + else: + tag_query = query + + fetch_tasks.append(fetch_projects(tag, tag_query)) + + fetch_results = await asyncio.gather(*fetch_tasks) + + for tag, results in zip(tag_num, fetch_results): + result_list[tag] = results return { "code": 200, - 'data': { + "data": { "result": result_list, "tag": tag_num } diff --git a/api/project_aggregation.py b/api/project_aggregation.py new file mode 100644 index 0000000..154fd7d --- /dev/null +++ b/api/project_aggregation.py @@ -0,0 +1,57 @@ +# ------------------------------------- +# @file : project_aggregation.py +# @author : Autumn +# @contact : rainy-autumn@outlook.com +# @time : 2024/7/8 21:02 +# ------------------------------------------- + +import time +import traceback + +from bson import ObjectId +from fastapi import APIRouter, Depends, BackgroundTasks + +from api.task import create_scan_task, delete_asset +from api.users import verify_token +from motor.motor_asyncio import AsyncIOMotorCursor + +from core.config import Project_List +from core.db import get_mongo_db +from core.redis_handler import refresh_config, get_redis_pool +from loguru import logger +from core.util import * +from core.apscheduler_handler import scheduler + +router = APIRouter() + + +@router.post("/project/info") +async def get_projects_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)): + id = request_data.get("id", "") + result = await db.project.find_one({"_id": ObjectId(id)}, { + "_id": 0, + "tag": 1, + "hour": 1, + "scheduledTasks": 1, + "AssetCount": 1, + "root_domains": 1, + "name":1 + } + ) + if result['scheduledTasks']: + job = scheduler.get_job(id) + if job is not None: + next_time = job.next_run_time.strftime("%Y-%m-%d %H:%M:%S") + result['next_time'] = next_time + return {"code": 200, "data": result} + + +@router.post("/project/asset/count") +async def get_projects_asset_count(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)): + id = request_data.get("id", "") + subdomain_count = await db['subdomain'].count_documents({"project": id}) + vulnerability_count = await db['vulnerability'].count_documents({"project": id}) + return {"code": 200, "data": { + "subdomainCount": subdomain_count, + "vulCount": vulnerability_count + }} \ No newline at end of file diff --git a/main.py b/main.py index 2eacdb4..3d7373c 100644 --- a/main.py +++ b/main.py @@ -104,7 +104,7 @@ async def http_exception_handler(request, exc): os.chdir(os.path.dirname(os.path.abspath(__file__))) from api import users, sensitive, dictionary, poc, configuration, fingerprint, node, project, task, asset_info, \ - page_monitoring, vulnerability, SubdoaminTaker, scheduled_tasks, notification, system, export + page_monitoring, vulnerability, SubdoaminTaker, scheduled_tasks, notification, system, export, project_aggregation app.include_router(users.router, prefix='/api') app.include_router(sensitive.router, prefix='/api') @@ -124,8 +124,8 @@ app.include_router(dirscan.router, prefix='/api') app.include_router(notification.router, prefix='/api') app.include_router(system.router, prefix='/api') app.include_router(export.router, prefix='/api') +app.include_router(project_aggregation.router, prefix='/api/project_aggregation') app.mount("/assets", StaticFiles(directory="static/assets"), name="assets") - @app.get("/logo.png", response_class=FileResponse) async def get_logo(request: Request): return FileResponse("static/logo.png")