# -*- coding:utf-8 -*-　　 # @name: sensitive # @auth: rainy-autumn@outlook.com # @version: import asyncio import json import traceback from loguru import logger from bson import ObjectId from fastapi import APIRouter, Depends, BackgroundTasks from pymongo import DESCENDING from api.users import verify_token from motor.motor_asyncio import AsyncIOMotorCursor from core.apscheduler_handler import scheduler from core.db import get_mongo_db from core.redis_handler import get_redis_pool, check_redis_task_target_is_null from core.util import * from api.node import get_redis_online_data, get_node_all from api.page_monitoring import get_page_monitoring_data router = APIRouter() @router.post("/task/data") async def get_task_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token), background_tasks: BackgroundTasks = BackgroundTasks(), redis_con=Depends(get_redis_pool)): try: background_tasks.add_task(task_progress) search_query = request_data.get("search", "") page_index = request_data.get("pageIndex", 1) page_size = request_data.get("pageSize", 10) # Fuzzy search based on the name field query = {"name": {"$regex": search_query, "$options": "i"}} # Get the total count of documents matching the search criteria total_count = await db.task.count_documents(query) # Perform pagination query cursor: AsyncIOMotorCursor = db.task.find(query).skip((page_index - 1) * page_size).limit(page_size).sort([("creatTime", DESCENDING)]) result = await cursor.to_list(length=None) # Process the result as needed response_data = [{"id": str(doc["_id"]), "name": doc["name"], "taskNum": doc["taskNum"], "progress": doc["progress"], "creatTime": doc["creatTime"], "endTime": doc["endTime"]} for doc in result] # response_data = [] # for doc in result: # transformed_doc = { # "id": str(doc["_id"]), # "name": doc["name"], # "taskNum": doc["taskNum"], # "progress": doc["progress"], # "creatTime": doc["creatTime"], # "endTime": doc["endTime"] # } # # if doc["progress"] != 100: # # background_tasks.add_task(check_redis_task_target_is_null, str(doc["_id"]), "", redis_con) # response_data.append(transformed_doc) 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} @router.post("/task/add") async def add_task(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token), redis_con=Depends(get_redis_pool)): try: name = request_data.get("name") cursor = db.task.find({"name": {"$eq": name}}, {"_id": 1}) results = await cursor.to_list(length=None) if len(results) != 0: return {"code": 400, "message": "name already exists"} target = request_data.get("target", "") node = request_data.get("node") if name == "" or target == "" or node == []: return {"message": "Null", "code": 500} scheduledTasks = request_data.get("scheduledTasks", False) hour = request_data.get("hour", 1) targetList = [] targetTmp = "" for t in 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) targetTmp += t + "\n" taskNum = len(targetList) request_data['taskNum'] = taskNum request_data['target'] = targetTmp.strip("\n") request_data['progress'] = 0 request_data["creatTime"] = get_now_time() request_data["endTime"] = "" if "All Poc" in request_data['vulList']: request_data['vulList'] = ["All Poc"] result = await db.task.insert_one(request_data) # Check if the insertion was successful if result.inserted_id: if scheduledTasks: scheduler.add_job(scheduler_scan_task, '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': 'Scan', 'state': True, 'lastTime': get_now_time(), 'nextTime': formatted_time, 'runner_id': str(result.inserted_id)}) f = await create_scan_task(request_data, result.inserted_id, targetList, redis_con) if f: return {"code": 200, "message": "Task added successfully"} else: return {"code": 400, "message": "Failed to add Task"} else: return {"code": 400, "message": "Failed to add Task"} except Exception as e: logger.error(str(e)) # Handle exceptions as needed return {"message": "error", "code": 500} @router.post("/task/content") async def task_content(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)): try: # Get the ID from the request data task_id = request_data.get("id") # Check if ID is provided if not task_id: return {"message": "ID is missing in the request data", "code": 400} # Query the database for content based on ID query = {"_id": ObjectId(task_id)} doc = await db.task.find_one(query) if not doc: return {"message": "Content not found for the provided ID", "code": 404} result = { "name": doc.get("name", ""), "target": doc.get("target", ""), "node": doc.get("node", []), "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"), "duplicates": doc.get("duplicates") } return {"code": 200, "data": result} except Exception as e: logger.error(str(e)) # Handle exceptions as needed return {"message": "error", "code": 500} @router.post("/task/delete") async def delete_task(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token), redis_con=Depends(get_redis_pool), background_tasks: BackgroundTasks = BackgroundTasks()): try: # Extract the list of IDs from the request_data dictionary task_ids = request_data.get("ids", []) delA = request_data.get("delA", False) # Convert the provided rule_ids to ObjectId obj_ids = [] redis_key = [] for task_id in task_ids: obj_ids.append(ObjectId(task_id)) redis_key.append("TaskInfo:" + task_id) job = scheduler.get_job(task_id) if job: scheduler.remove_job(task_id) await db.ScheduledTasks.delete_many({"id": {"$in": task_ids}}) if delA: background_tasks.add_task(delete_asset, task_ids) await redis_con.delete(*redis_key) # Delete the SensitiveRule documents based on the provided IDs result = await db.task.delete_many({"_id": {"$in": obj_ids}}) # Check if the deletion was successful if result.deleted_count > 0: return {"code": 200, "message": "Task deleted successfully"} else: return {"code": 404, "message": "Task not found"} except Exception as e: logger.error(str(e)) # Handle exceptions as needed return {"message": "error", "code": 500} @router.post("/task/retest") async def retest_task(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token), redis_con=Depends(get_redis_pool)): try: # Get the ID from the request data task_id = request_data.get("id") # Check if ID is provided if not task_id: return {"message": "ID is missing in the request data", "code": 400} # Query the database for content based on ID query = {"_id": ObjectId(task_id)} doc = await db.task.find_one(query) if not doc: return {"message": "Content not found for the provided ID", "code": 404} target = doc['target'] targetList = target.split("\n") keys_to_delete = [ f"TaskInfo:tmp:{task_id}", f"TaskInfo:{task_id}", f"TaskInfo:time:{task_id}", f"duplicates:url:{task_id}", f"duplicates:domain:{task_id}", f"duplicates:sensresp:{task_id}", f"duplicates:craw:{task_id}" ] progresskeys = await redis_con.keys(f"TaskInfo:progress:{task_id}:*") keys_to_delete.extend(progresskeys) if keys_to_delete: await redis_con.delete(*keys_to_delete) f = await create_scan_task(doc, task_id, targetList, redis_con) if f: update_document = { "$set": { "progress": 0, "creatTime": get_now_time(), "endTime": "" } } await db.task.update_one({"_id": ObjectId(task_id)}, update_document) return {"code": 200, "message": "Task added successfully"} else: return {"code": 400, "message": "Failed to add Task"} except Exception as e: logger.error(str(e)) # Handle exceptions as needed return {"message": "error", "code": 500} async def create_scan_task(request_data, id, targetList, redis_con): try: request_data["id"] = str(id) if request_data['allNode']: request_data["node"] = await get_node_all(redis_con) keys_to_delete = [ f"TaskInfo:tmp:{id}", f"TaskInfo:{id}", f"TaskInfo:time:{id}", f"duplicates:url:{id}", f"duplicates:domain:{id}", f"duplicates:sensresp:{id}", f"duplicates:craw:{id}" ] progresskeys = await redis_con.keys(f"TaskInfo:progress:{id}:*") keys_to_delete.extend(progresskeys) if keys_to_delete: await redis_con.delete(*keys_to_delete) add_redis_task_data = transform_db_redis(request_data) async with redis_con as redis: await redis.lpush(f"TaskInfo:{id}", *targetList) for name in request_data["node"]: await redis.rpush(f"NodeTask:{name}", json.dumps(add_redis_task_data)) return True except Exception as e: logger.error(str(e)) # Handle exceptions as needed return False @router.post("/task/update") async def update_task_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)): try: # Get the ID from the request data task_id = request_data.get("id") hour = request_data.get("hour") # Check if ID is provided if not task_id: return {"message": "ID is missing in the request data", "code": 400} query = {"id": task_id} doc = await db.ScheduledTasks.find_one(query) oldScheduledTasks = doc["state"] old_hour = doc["hour"] newScheduledTasks = request_data.get("scheduledTasks") if oldScheduledTasks != newScheduledTasks: if newScheduledTasks: scheduler.add_job(scheduler_scan_task, 'interval', hours=hour, args=[task_id], id=str(task_id), jobstore='mongo') await db.ScheduledTasks.update_one({"id": task_id}, {"$set": {'state': True}}) else: scheduler.remove_job(task_id) await db.ScheduledTasks.update_one({"id": task_id}, {"$set": {'state': False}}) if newScheduledTasks: if hour != old_hour: await db.ScheduledTasks.update_one({"id": task_id}, {"$set": {'hour': hour}}) job = scheduler.get_job(task_id) if job is not None: scheduler.remove_job(task_id) scheduler.add_job(scheduler_scan_task, 'interval', hours=hour, args=[task_id], id=str(task_id), jobstore='mongo') else: scheduler.add_job(scheduler_scan_task, 'interval', hours=hour, args=[task_id], id=str(task_id), jobstore='mongo') request_data.pop("id") update_document = { "$set": request_data } result = await db.task.update_one({"_id": ObjectId(task_id)}, update_document) # Check if the update was successful if result: 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()) return {"message": "error", "code": 500} async def task_progress(): async for db in get_mongo_db(): async for redis in get_redis_pool(): query = {"progress": {"$ne": 100}} cursor: AsyncIOMotorCursor = db.task.find(query) result = await cursor.to_list(length=None) if len(result) == 0: return True for r in result: id = str(r["_id"]) key = f"TaskInfo:tmp:{id}" exists = await redis.exists(key) if exists: count = await redis.llen(key) progress_tmp = round(count / r['taskNum'], 2) progress_tmp = round(progress_tmp * 100, 1) if progress_tmp > 100: progress_tmp = 100 if progress_tmp == 100: time_key = f"TaskInfo:time:{id}" time_value = await redis.get(time_key) await db.task.update_one({"_id": r["_id"]}, {"$set": {"endTime": time_value}}) await db.task.update_one({"_id": r["_id"]}, {"$set": {"progress": progress_tmp}}) else: await db.task.update_one({"_id": r["_id"]}, {"$set": {"progress": 0}}) return # @router.post("/task/progress/info") # async def progress_info(request_data: dict, _: dict = Depends(verify_token), redis_con=Depends(get_redis_pool), db=Depends(get_mongo_db)): # task_id = request_data.get("id") # type = request_data.get("type") # runner = request_data.get("runner") # # Check if ID is provided # if not task_id: # return {"message": "ID is missing in the request data", "code": 400} # query = {"_id": ObjectId(task_id)} # if type == "scan": # doc = await db.task.find_one(query) # else: # doc = await db.project.find_one(query) # target_data = await db.ProjectTargetData.find_one({"id": task_id}) # doc["target"] = target_data["target"] # if not doc: # return {"message": "Content not found for the provided ID", "code": 404} # target = doc['target'] # result_list = [] # for t in target.split("\n"): # progress_result = { # "subdomain": ["", ""], # "subdomainTakeover": ["", ""], # "portScan": ["", ""], # "assetMapping": ["", ""], # "urlScan": ["", ""], # "sensitive": ["", ""], # "crawler": ["", ""], # "dirScan": ["", ""], # "vulnerability": ["", ""], # "all": ["", ""], # "target": "" # } # if not doc['subdomainScan']: # progress_result['subdomain'] = ['', '', ''] # if not doc['portScan']: # progress_result['portScan'] = ['', '', ''] # if not doc['urlScan']: # progress_result['urlScan'] = ['', '', ''] # if not doc['sensitiveInfoScan']: # progress_result['sensitive'] = ['', '', ''] # if not doc['crawlerScan']: # progress_result['crawler'] = ['', '', ''] # if not doc['dirScan']: # progress_result['dirScan'] = ['', '', ''] # if not doc['vulScan']: # progress_result['vulnerability'] = ['', '', ''] # if runner != "": # key = "TaskInfo:progress:" + runner + ":" + t # else: # key = "TaskInfo:progress:" + task_id + ":" + t # data = await redis_con.hgetall(key) # progress_result["target"] = t # if not data: # result_list.append(progress_result) # else: # progress_result['subdomain'][0] = data.get("subdomain_start","") # progress_result['subdomain'][1] = data.get("subdomain_end", "") # progress_result['subdomainTakeover'][0] = data.get("subdomainTakeover_start", "") # progress_result['subdomainTakeover'][1] = data.get("subdomainTakeover_end", "") # progress_result['portScan'][0] = data.get("portScan_start", "") # progress_result['portScan'][1] = data.get("portScan_end", "") # progress_result['assetMapping'][0] = data.get("assetMapping_start", "") # progress_result['assetMapping'][1] = data.get("assetMapping_end", "") # progress_result['urlScan'][0] = data.get("urlScan_start", "") # progress_result['urlScan'][1] = data.get("urlScan_end", "") # progress_result['sensitive'][0] = data.get("sensitive_start", "") # progress_result['sensitive'][1] = data.get("sensitive_end", "") # progress_result['crawler'][0] = data.get("crawler_start", "") # progress_result['crawler'][1] = data.get("crawler_end", "") # progress_result['dirScan'][0] = data.get("dirScan_start", "") # progress_result['dirScan'][1] = data.get("dirScan_end", "") # progress_result['vulnerability'][0] = data.get("vulnerability_start", "") # progress_result['vulnerability'][1] = data.get("vulnerability_end", "") # progress_result['all'][0] = data.get("scan_start", "") # progress_result['all'][1] = data.get("scan_end", "") # result_list.append(progress_result) # return { # "code": 200, # "data": { # 'list': result_list, # "total": len(result_list) # } # } @router.post("/task/progress/info") async def progress_info(request_data: dict, _: dict = Depends(verify_token), redis_con=Depends(get_redis_pool), db=Depends(get_mongo_db)): task_id = request_data.get("id") type = request_data.get("type") runner = request_data.get("runner") if not task_id: return {"message": "ID is missing in the request data", "code": 400} query = {"_id": ObjectId(task_id)} if type == "scan": doc = await db.task.find_one(query) else: doc, target_data = await asyncio.gather( db.project.find_one(query), db.ProjectTargetData.find_one({"id": task_id}) ) if target_data: doc["target"] = target_data["target"] if not doc: return {"message": "Content not found for the provided ID", "code": 404} target = doc['target'] result_list = [] tasks = [] for t in target.split("\n"): key = f"TaskInfo:progress:{runner or task_id}:{t}" tasks.append(redis_con.hgetall(key)) redis_results = await asyncio.gather(*tasks) for t, data in zip(target.split("\n"), redis_results): progress_result = { "subdomain": ["", ""], "subdomainTakeover": ["", ""], "portScan": ["", ""], "assetMapping": ["", ""], "urlScan": ["", ""], "sensitive": ["", ""], "crawler": ["", ""], "dirScan": ["", ""], "vulnerability": ["", ""], "all": ["", ""], "target": t } if not data: result_list.append(progress_result) continue progress_result['subdomain'][0] = data.get("subdomain_start", "") progress_result['subdomain'][1] = data.get("subdomain_end", "") progress_result['subdomainTakeover'][0] = data.get("subdomainTakeover_start", "") progress_result['subdomainTakeover'][1] = data.get("subdomainTakeover_end", "") progress_result['portScan'][0] = data.get("portScan_start", "") progress_result['portScan'][1] = data.get("portScan_end", "") progress_result['assetMapping'][0] = data.get("assetMapping_start", "") progress_result['assetMapping'][1] = data.get("assetMapping_end", "") progress_result['urlScan'][0] = data.get("urlScan_start", "") progress_result['urlScan'][1] = data.get("urlScan_end", "") progress_result['sensitive'][0] = data.get("sensitive_start", "") progress_result['sensitive'][1] = data.get("sensitive_end", "") progress_result['crawler'][0] = data.get("crawler_start", "") progress_result['crawler'][1] = data.get("crawler_end", "") progress_result['dirScan'][0] = data.get("dirScan_start", "") progress_result['dirScan'][1] = data.get("dirScan_end", "") progress_result['vulnerability'][0] = data.get("vulnerability_start", "") progress_result['vulnerability'][1] = data.get("vulnerability_end", "") progress_result['all'][0] = data.get("scan_start", "") progress_result['all'][1] = data.get("scan_end", "") result_list.append(progress_result) return { "code": 200, "data": { 'list': result_list, "total": len(result_list) } } async def scheduler_scan_task(id): logger.info(f"Scheduler scan {id}") async for db in get_mongo_db(): async for redis in get_redis_pool(): next_time = scheduler.get_job(id).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 != "" and id != run_id_last: progresskeys = await redis.keys(f"TaskInfo:progress:{run_id_last}:*") for pgk in progresskeys: await redis.delete(pgk) task_id = generate_random_string(15) 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.task.find_one(query) targetList = [] for t in doc['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) async def delete_asset(task_ids, is_project = False): async for db in get_mongo_db(): key = ["asset", "subdomain", "SubdoaminTakerResult", "UrlScan", "crawler", "SensitiveResult", "DirScanResult", "vulnerability", "PageMonitoring"] del_query = {"taskId": {"$in": task_ids}} if is_project: del_query = { "$or": [ {"taskId": {"$in": task_ids}}, {"project": {"$in": task_ids}} ] } for k in key: result = await db[k].delete_many(del_query) if result.deleted_count > 0: logger.info("Deleted {} {} documents".format(k, result.deleted_count)) else: logger.info("Deleted {} None documents".format(k))