ScopeSentry/api/vulnerability.py

79 lines
2.7 KiB
Python
Raw Normal View History

2024-06-05 13:39:34 +00:00
# -------------------------------------
# @file : vulnerability.py
# @author : Autumn
# @contact : rainy-autumn@outlook.com
# @time : 2024/4/27 13:25
# -------------------------------------------
from fastapi import APIRouter, Depends
from motor.motor_asyncio import AsyncIOMotorCursor
from pymongo import DESCENDING
from api.users import verify_token
from core.config import POC_LIST
from core.db import get_mongo_db
from core.util import search_to_mongodb
from loguru import logger
router = APIRouter()
@router.post("/vul/data")
async def get_vul_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
try:
search_query = request_data.get("search", "")
page_index = request_data.get("pageIndex", 1)
page_size = request_data.get("pageSize", 10)
# MongoDB collection for SensitiveRule
# Fuzzy search based on the name field
keyword = {
'url': 'url',
'vulname': 'vulname',
'project': 'project',
'matched': 'matched',
'request': 'request',
'response': 'response',
'level': 'level'
}
query = await search_to_mongodb(search_query, keyword)
if query == "" or query is None:
return {"message": "Search condition parsing error", "code": 500}
query = query[0]
# Get the total count of documents matching the search criteria
total_count = await db.vulnerability.count_documents(query)
if total_count == 0:
return {
"code": 200,
"data": {
'list': [],
'total': 0
}
}
# Perform pagination query
cursor: AsyncIOMotorCursor = db.vulnerability.find(query).skip((page_index - 1) * page_size).limit(page_size).sort([("timestamp", DESCENDING)])
result = await cursor.to_list(length=None)
# Process the result as needed
response_data = []
for doc in result:
data = {
"id": str(doc["_id"]),
"url": doc["url"],
"vulnerability": doc["vulname"],
"matched": doc["matched"],
"time": doc["time"],
"request": doc["request"],
"response": doc["response"],
}
if doc["vulnid"] in POC_LIST:
data["level"] = POC_LIST[doc["vulnid"]]
response_data.append(data)
return {
"code": 200,
"data": {
'list': response_data,
'total': total_count
}
}
except Exception as e:
logger.error(str(e))
# Handle exceptions as needed
return {"message": "error","code":500}