481 lines
19 KiB
Python
481 lines
19 KiB
Python
import time
|
|
import traceback
|
|
import asyncio
|
|
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/data")
|
|
async def get_projects_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token),
|
|
background_tasks: BackgroundTasks = BackgroundTasks()):
|
|
background_tasks.add_task(update_project_count)
|
|
search_query = request_data.get("search", "")
|
|
page_index = request_data.get("pageIndex", 1)
|
|
page_size = request_data.get("pageSize", 10)
|
|
|
|
query = {
|
|
"$or": [
|
|
{"name": {"$regex": search_query, "$options": "i"}},
|
|
{"root_domain": {"$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
|
|
|
|
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)
|
|
for result in results:
|
|
result["AssetCount"] = result.get("AssetCount", 0)
|
|
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": {
|
|
"result": result_list,
|
|
"tag": tag_num
|
|
}
|
|
}
|
|
|
|
|
|
@router.get("/project/all")
|
|
async def get_projects_all(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
|
try:
|
|
pipeline = [
|
|
{
|
|
"$group": {
|
|
"_id": "$tag", # 根据 tag 字段分组
|
|
"children": {"$push": {"value": {"$toString": "$_id"}, "label": "$name"}} # 将每个文档的 _id 和 name 放入 children 集合中
|
|
}
|
|
},
|
|
{
|
|
"$project": {
|
|
"_id": 0,
|
|
"label": "$_id",
|
|
"value": {"$literal": ""},
|
|
"children": 1
|
|
}
|
|
}
|
|
]
|
|
result = await db['project'].aggregate(pipeline).to_list(None)
|
|
return {
|
|
"code": 200,
|
|
"data": {
|
|
'list': result
|
|
}
|
|
}
|
|
except Exception as e:
|
|
logger.error(str(e))
|
|
logger.error(traceback.format_exc())
|
|
return {"message": "error","code":500}
|
|
|
|
|
|
async def update_project_count():
|
|
async for db in get_mongo_db():
|
|
cursor = db.project.find({}, {"_id": 0, "id": {"$toString": "$_id"}})
|
|
results = await cursor.to_list(length=None)
|
|
|
|
async def update_count(id):
|
|
query = {"project": {"$eq": id}}
|
|
total_count = await db.asset.count_documents(query)
|
|
if total_count != 0:
|
|
update_document = {
|
|
"$set": {
|
|
"AssetCount": total_count
|
|
}
|
|
}
|
|
await db.project.update_one({"_id": ObjectId(id)}, update_document)
|
|
|
|
fetch_tasks = [update_count(r['id']) for r in results]
|
|
|
|
await asyncio.gather(*fetch_tasks)
|
|
|
|
|
|
@router.post("/project/content")
|
|
async def get_project_content(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
|
project_id = request_data.get("id")
|
|
if not project_id:
|
|
return {"message": "ID is missing in the request data", "code": 400}
|
|
query = {"_id": ObjectId(project_id)}
|
|
doc = await db.project.find_one(query)
|
|
if not doc:
|
|
return {"message": "Content not found for the provided ID", "code": 404}
|
|
project_target_data = await db.ProjectTargetData.find_one({"id": project_id})
|
|
result = {
|
|
"name": doc.get("name", ""),
|
|
"tag": doc.get("tag", ""),
|
|
"target": project_target_data.get("target", ""),
|
|
"node": doc.get("node", []),
|
|
"logo": doc.get("logo", ""),
|
|
"subdomainScan": doc.get("subdomainScan", False),
|
|
"subdomainConfig": doc.get("subdomainConfig", []),
|
|
"urlScan": doc.get("urlScan", False),
|
|
"sensitiveInfoScan": doc.get("sensitiveInfoScan", False),
|
|
"pageMonitoring": doc.get("pageMonitoring", ""),
|
|
"crawlerScan": doc.get("crawlerScan", False),
|
|
"vulScan": doc.get("vulScan", False),
|
|
"vulList": doc.get("vulList", []),
|
|
"portScan": doc.get("portScan"),
|
|
"ports": doc.get("ports"),
|
|
"waybackurl": doc.get("waybackurl"),
|
|
"dirScan": doc.get("dirScan"),
|
|
"scheduledTasks": doc.get("scheduledTasks"),
|
|
"hour": doc.get("hour"),
|
|
"allNode": doc.get("allNode", False),
|
|
"duplicates": doc.get("duplicates")
|
|
}
|
|
return {"code": 200, "data": result}
|
|
|
|
|
|
@router.post("/project/add")
|
|
async def add_project_rule(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token),
|
|
background_tasks: BackgroundTasks = BackgroundTasks()):
|
|
try:
|
|
# Extract values from request data
|
|
name = request_data.get("name")
|
|
target = request_data.get("target")
|
|
runNow = request_data.get("runNow")
|
|
del request_data["runNow"]
|
|
scheduledTasks = request_data.get("scheduledTasks", False)
|
|
hour = request_data.get("hour", 1)
|
|
t_list = []
|
|
tmsg = ''
|
|
root_domains = []
|
|
for t in target.split('\n'):
|
|
if t not in t_list:
|
|
targetTmp = t.replace('http://', "").replace('https://', "").replace("\n", "").replace("\r", "").strip()
|
|
if targetTmp != "":
|
|
root_domain = get_root_domain(targetTmp)
|
|
if root_domain not in root_domains:
|
|
root_domains.append(root_domain)
|
|
tmsg += targetTmp + '\n'
|
|
# Create a new SensitiveRule document
|
|
tmsg = tmsg.strip().strip("\n")
|
|
request_data["root_domains"] = root_domains
|
|
del request_data['target']
|
|
if "All Poc" in request_data['vulList']:
|
|
request_data['vulList'] = ["All Poc"]
|
|
cursor = db.project.find({"name": {"$eq": name}}, {"_id": 1})
|
|
results = await cursor.to_list(length=None)
|
|
if len(results) != 0:
|
|
return {"code": 400, "message": "name already exists"}
|
|
# Insert the new document into the SensitiveRule collection
|
|
result = await db.project.insert_one(request_data)
|
|
# Check if the insertion was successful
|
|
if result.inserted_id:
|
|
await db.ProjectTargetData.insert_one({"id": str(result.inserted_id), "target": tmsg})
|
|
if scheduledTasks:
|
|
scheduler.add_job(scheduler_project, 'interval', hours=hour, args=[str(result.inserted_id)],
|
|
id=str(result.inserted_id), jobstore='mongo')
|
|
next_time = scheduler.get_job(str(result.inserted_id)).next_run_time
|
|
formatted_time = next_time.strftime("%Y-%m-%d %H:%M:%S")
|
|
db.ScheduledTasks.insert_one(
|
|
{"id": str(result.inserted_id), "name": name, 'hour': hour, 'type': 'Project', 'state': True,
|
|
'lastTime': get_now_time(), 'nextTime': formatted_time, 'runner_id': str(result.inserted_id)})
|
|
if runNow:
|
|
await scheduler_project(str(result.inserted_id))
|
|
background_tasks.add_task(update_project, root_domains, str(result.inserted_id))
|
|
await refresh_config('all', 'project')
|
|
Project_List[name] = str(result.inserted_id)
|
|
return {"code": 200, "message": "Project added successfully"}
|
|
else:
|
|
return {"code": 400, "message": "Failed to add Project"}
|
|
|
|
except Exception as e:
|
|
logger.error(str(e))
|
|
# Handle exceptions as needed
|
|
return {"message": "error", "code": 500}
|
|
|
|
|
|
@router.post("/project/delete")
|
|
async def delete_project_rules(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token),
|
|
background_tasks: BackgroundTasks = BackgroundTasks()):
|
|
try:
|
|
pro_ids = request_data.get("ids", [])
|
|
delA = request_data.get("delA", False)
|
|
if delA:
|
|
background_tasks.add_task(delete_asset, pro_ids, True)
|
|
obj_ids = [ObjectId(poc_id) for poc_id in pro_ids]
|
|
result = await db.project.delete_many({"_id": {"$in": obj_ids}})
|
|
await db.ProjectTargetData.delete_many({"id": {"$in": pro_ids}})
|
|
# Check if the deletion was successful
|
|
if result.deleted_count > 0:
|
|
for pro_id in pro_ids:
|
|
job = scheduler.get_job(pro_id)
|
|
if job:
|
|
scheduler.remove_job(pro_id)
|
|
background_tasks.add_task(delete_asset_project_handler, pro_id)
|
|
for project_id in Project_List:
|
|
if pro_id == Project_List[project_id]:
|
|
del Project_List[project_id]
|
|
break
|
|
await db.ScheduledTasks.delete_many({"id": {"$in": pro_ids}})
|
|
return {"code": 200, "message": "Project deleted successfully"}
|
|
else:
|
|
return {"code": 404, "message": "Project not found"}
|
|
|
|
except Exception as e:
|
|
logger.error(str(e))
|
|
# Handle exceptions as needed
|
|
return {"message": "error", "code": 500}
|
|
|
|
|
|
@router.post("/project/update")
|
|
async def update_project_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token),
|
|
background_tasks: BackgroundTasks = BackgroundTasks()):
|
|
try:
|
|
# Get the ID from the request data
|
|
pro_id = request_data.get("id")
|
|
hour = request_data.get("hour")
|
|
runNow = request_data.get("runNow")
|
|
del request_data["runNow"]
|
|
# Check if ID is provided
|
|
if not pro_id:
|
|
return {"message": "ID is missing in the request data", "code": 400}
|
|
query = {"id": pro_id}
|
|
doc = await db.ScheduledTasks.find_one(query)
|
|
|
|
newScheduledTasks = request_data.get("scheduledTasks")
|
|
if doc:
|
|
oldScheduledTasks = doc["state"]
|
|
old_hour = doc["hour"]
|
|
if oldScheduledTasks != newScheduledTasks:
|
|
if newScheduledTasks:
|
|
scheduler.add_job(scheduler_project, 'interval', hours=hour, args=[pro_id],
|
|
id=str(pro_id), jobstore='mongo')
|
|
await db.ScheduledTasks.update_one({"id": pro_id}, {"$set": {'state': True}})
|
|
else:
|
|
job = scheduler.get_job(pro_id)
|
|
if job is not None:
|
|
scheduler.remove_job(pro_id)
|
|
await db.ScheduledTasks.update_one({"id": pro_id}, {"$set": {'state': False}})
|
|
else:
|
|
if newScheduledTasks:
|
|
if hour != old_hour:
|
|
job = scheduler.get_job(pro_id)
|
|
if job is not None:
|
|
scheduler.remove_job(pro_id)
|
|
scheduler.add_job(scheduler_project, 'interval', hours=hour, args=[pro_id],
|
|
id=str(pro_id), jobstore='mongo')
|
|
else:
|
|
if newScheduledTasks:
|
|
scheduler.add_job(scheduler_project, 'interval', hours=hour, args=[str(pro_id)],
|
|
id=str(pro_id), jobstore='mongo')
|
|
next_time = scheduler.get_job(str(pro_id)).next_run_time
|
|
formatted_time = next_time.strftime("%Y-%m-%d %H:%M:%S")
|
|
db.ScheduledTasks.insert_one(
|
|
{"id": str(pro_id), "name": request_data['name'], 'hour': hour, 'type': 'Project', 'state': True,
|
|
'lastTime': get_now_time(), 'nextTime': formatted_time})
|
|
|
|
result = await db.project.find_one({"_id": ObjectId(pro_id)})
|
|
project_target_data = await db.ProjectTargetData.find_one({"id": pro_id})
|
|
old_targets = project_target_data['target']
|
|
old_name = result['name']
|
|
new_name = request_data['name']
|
|
new_targets = request_data['target']
|
|
if old_targets != new_targets.strip().strip('\n'):
|
|
new_root_domain = []
|
|
update_document = {
|
|
"$set": {
|
|
"target": new_targets.strip().strip('\n')
|
|
}
|
|
}
|
|
for n_t in new_targets.strip().strip('\n').split('\n'):
|
|
t_root_domain = get_root_domain(n_t)
|
|
if t_root_domain not in new_root_domain:
|
|
new_root_domain.append(t_root_domain)
|
|
request_data["root_domains"] = new_root_domain
|
|
await db.ProjectTargetData.update_one({"id": pro_id}, update_document)
|
|
background_tasks.add_task(update_project, new_root_domain, pro_id, True)
|
|
if old_name != new_name:
|
|
del Project_List[old_name]
|
|
Project_List[new_name] = pro_id
|
|
await db.ScheduledTasks.update_one({"id": pro_id}, {"$set": {'name': new_name}})
|
|
request_data.pop("id")
|
|
del request_data['target']
|
|
update_document = {
|
|
"$set": request_data
|
|
}
|
|
result = await db.project.update_one({"_id": ObjectId(pro_id)}, update_document)
|
|
# Check if the update was successful
|
|
if result:
|
|
if runNow:
|
|
await scheduler_project(str(pro_id))
|
|
return {"message": "Task updated successfully", "code": 200}
|
|
else:
|
|
return {"message": "Failed to update data", "code": 404}
|
|
|
|
except Exception as e:
|
|
logger.error(str(e))
|
|
logger.error(traceback.format_exc())
|
|
# Handle exceptions as needed
|
|
return {"message": "error", "code": 500}
|
|
|
|
async def update_project(root_domain, project_id, change=False):
|
|
asset_collection_list = {
|
|
'asset': ["url", "host", "ip"],
|
|
'subdomain': ["host", "ip"],
|
|
'DirScanResult': ["url"],
|
|
'vulnerability': ["url"],
|
|
'SubdoaminTakerResult': ["input"],
|
|
'PageMonitoring': ["url"],
|
|
'SensitiveResult': ["url"],
|
|
'UrlScan': ["input"],
|
|
'crawler': ["url"]}
|
|
async for db in get_mongo_db():
|
|
for a in asset_collection_list:
|
|
if change:
|
|
await asset_update_project(root_domain, asset_collection_list[a], a, db, project_id)
|
|
else:
|
|
await asset_add_project(root_domain, asset_collection_list[a], a, db, project_id)
|
|
|
|
|
|
async def asset_add_project(root_domain, db_key, doc_name, db, project_id):
|
|
regex_patterns = [f".*{domain}.*" for domain in root_domain]
|
|
pattern = "|".join(regex_patterns)
|
|
# 构建查询条件
|
|
query = {
|
|
"$and": [
|
|
{
|
|
"$or": [
|
|
{key: {"$regex": pattern, "$options": "i"}} for key in db_key
|
|
]
|
|
},
|
|
{"project": {"$exists": True, "$eq": ""}}
|
|
]
|
|
}
|
|
update_query = {
|
|
"$set": {
|
|
"project": project_id
|
|
}
|
|
}
|
|
result = await db[doc_name].update_many(query, update_query)
|
|
# 打印更新的文档数量
|
|
logger.info(f"Updated {doc_name} {result.modified_count} documents")
|
|
|
|
|
|
async def asset_update_project(root_domain, db_key, doc_name, db, project_id):
|
|
regex_patterns = [f".*{domain}.*" for domain in root_domain]
|
|
pattern = "|".join(regex_patterns)
|
|
# 构建查询条件
|
|
query = {
|
|
"$and": [
|
|
{"project": project_id},
|
|
{
|
|
"$nor": [
|
|
{key: {"$regex": pattern, "$options": "i"}} for key in db_key
|
|
]
|
|
}
|
|
]
|
|
}
|
|
update_query = {
|
|
"$set": {
|
|
"project": ""
|
|
}
|
|
}
|
|
result = await db[doc_name].update_many(query, update_query)
|
|
# 打印更新的文档数量
|
|
logger.info(f"Updated {doc_name} {result.modified_count} documents to null ")
|
|
await asset_add_project(root_domain, db_key, doc_name, db, project_id)
|
|
|
|
|
|
async def delete_asset_project(db, collection, project_id):
|
|
try:
|
|
query = {"project": project_id}
|
|
|
|
cursor = db[collection].find(query)
|
|
|
|
async for document in cursor:
|
|
await db[collection].update_one({"_id": document["_id"]}, {"$set": {"project": ""}})
|
|
|
|
except Exception as e:
|
|
logger.error(f"delete_asset_project error:{e}")
|
|
|
|
|
|
async def delete_asset_project_handler(project_id):
|
|
async for db in get_mongo_db():
|
|
asset_collection_list = ['asset', 'subdomain', 'DirScanResult', 'vulnerability', 'SubdoaminTakerResult',
|
|
'PageMonitoring', 'SensitiveResult', 'UrlScan', 'crawler']
|
|
for c in asset_collection_list:
|
|
await delete_asset_project(db, c, project_id)
|
|
|
|
|
|
async def scheduler_project(id):
|
|
logger.info(f"Scheduler project {id}")
|
|
async for db in get_mongo_db():
|
|
async for redis in get_redis_pool():
|
|
job = scheduler.get_job(id)
|
|
task_id = generate_random_string(15)
|
|
if job:
|
|
next_time = job.next_run_time
|
|
formatted_time = next_time.strftime("%Y-%m-%d %H:%M:%S")
|
|
doc = await db.ScheduledTasks.find_one({"id": id})
|
|
run_id_last = doc.get("runner_id", "")
|
|
if run_id_last != "":
|
|
progresskeys = await redis.keys(f"TaskInfo:progress:{run_id_last}:*")
|
|
for pgk in progresskeys:
|
|
await redis.delete(pgk)
|
|
update_document = {
|
|
"$set": {
|
|
"lastTime": get_now_time(),
|
|
"nextTime": formatted_time,
|
|
"runner_id": task_id
|
|
}
|
|
}
|
|
await db.ScheduledTasks.update_one({"id": id}, update_document)
|
|
query = {"_id": ObjectId(id)}
|
|
doc = await db.project.find_one(query)
|
|
targetList = []
|
|
target_data = await db.ProjectTargetData.find_one({"id": id})
|
|
for t in target_data.get('target', '').split("\n"):
|
|
t.replace("http://", "").replace("https://", "")
|
|
t = t.strip("\n").strip("\r").strip()
|
|
if t != "" and t not in targetList:
|
|
targetList.append(t)
|
|
await create_scan_task(doc, task_id, targetList, redis)
|