liangsheng8708
/
ScopeSentry
mirror of https://mirror.ghproxy.com/https://github.com/Autumn-27/ScopeSentry.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add

This commit is contained in:
Autumn.home 2024-07-08 23:30:40 +08:00
parent 30ed71b838
commit 017b229240
3 changed files with 105 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

89
api/project.py
View File

@ -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
}

57
api/project_aggregation.py Normal file
View File

@ -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
View File

@ -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")