305 lines
9.6 KiB
Python
305 lines
9.6 KiB
Python
# -------------------------------------
|
|
# @file : project_aggregation.py
|
|
# @author : Autumn
|
|
# @contact : rainy-autumn@outlook.com
|
|
# @time : 2024/7/8 21:02
|
|
# -------------------------------------------
|
|
import asyncio
|
|
import time
|
|
import traceback
|
|
|
|
from bson import ObjectId
|
|
from fastapi import APIRouter, Depends, BackgroundTasks
|
|
from pymongo import DESCENDING
|
|
|
|
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
|
|
}}
|
|
|
|
|
|
@router.post("/project/vul/statistics")
|
|
async def get_projects_vul_statistics(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
|
id = request_data.get("id", "")
|
|
pipeline = [
|
|
{"$match": {"project": id}},
|
|
{
|
|
"$group": {
|
|
"_id": "$level",
|
|
"count": {"$sum": 1}
|
|
}
|
|
}
|
|
]
|
|
result = await db['vulnerability'].aggregate(pipeline).to_list(None)
|
|
return {"code": 200, "data": result}
|
|
|
|
|
|
@router.post("/project/vul/data")
|
|
async def get_projects_vul_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
|
id = request_data.get("id", "")
|
|
cursor: AsyncIOMotorCursor = db.vulnerability.find({"project": id},
|
|
{"_id": 0, "url": 1, "vulname": 1, "level": 1, "time": 1,
|
|
"matched": 1}).sort([("time", DESCENDING)])
|
|
result = await cursor.to_list(length=None)
|
|
return {
|
|
"code": 200,
|
|
"data": {
|
|
'list': result
|
|
}
|
|
}
|
|
|
|
|
|
@router.post("/project/subdomain/data")
|
|
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 = {}
|
|
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:
|
|
return {"code": 404, "message": "domain is null"}
|
|
query = await get_search_query("subdomain", request_data)
|
|
if query == "":
|
|
return {"message": "Search condition parsing error", "code": 500}
|
|
results = []
|
|
for root_domain in doc["root_domains"]:
|
|
query["$and"].append({"host": {"$regex": f"{root_domain}$"}})
|
|
cursor: AsyncIOMotorCursor = db['subdomain'].find(query, {
|
|
"_id": 0, "id": {"$toString": "$_id"}, "host": 1, "type": 1, "value": 1, "ip": 1, "time": 1
|
|
}).sort([("time", -1)])
|
|
result = await cursor.to_list(length=None)
|
|
result_list = []
|
|
for r in result:
|
|
if r['value'] is None:
|
|
r['value'] = []
|
|
if r['ip'] is None:
|
|
r['ip'] = []
|
|
result_list.append(r)
|
|
results.append({"host": root_domain, "type": "", "value": [], "ip": [], "id": generate_random_string(5), "children": result_list, "count": len(result_list)})
|
|
return {
|
|
"code": 200,
|
|
"data": {
|
|
'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),
|
|
"count": r['count']
|
|
}
|
|
|
|
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'],
|
|
"count": r['count'],
|
|
"id": generate_random_string(5),
|
|
"domain": "",
|
|
"ip": "",
|
|
"timestamp": "",
|
|
"port": ""
|
|
}
|
|
|
|
query_copy = query.copy()
|
|
if r['_id'] == 'http' or r['_id'] == 'https':
|
|
query_copy["type"] = r['_id']
|
|
else:
|
|
query_copy["protocol"] = 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,
|
|
"port": 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['service'] = ''
|
|
children_data['domain'] = asset['host']
|
|
children_data['ip'] = asset['ip']
|
|
else:
|
|
if 'webServer' in asset:
|
|
children_data['service'] = asset['webServer']
|
|
else:
|
|
children_data['service'] = ''
|
|
children_data['domain'] = asset['domain']
|
|
children_data['ip'] = asset['host']
|
|
|
|
children_data['timestamp'] = asset['timestamp']
|
|
children_data['port'] = asset['port']
|
|
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
|
|
}
|
|
}
|