add
This commit is contained in:
parent
30ed71b838
commit
017b229240
|
@ -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 = {}
|
||||
} 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,
|
||||
|
||||
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)
|
||||
}).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
|
||||
}
|
||||
|
|
|
@ -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
|
||||
}}
|
4
main.py
4
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")
|
||||
|
|
Loading…
Reference in New Issue