diff --git a/api/asset_info.py b/api/asset_info.py index dcd5a08..8a4678e 100644 --- a/api/asset_info.py +++ b/api/asset_info.py @@ -5,6 +5,7 @@ # @time : 2024/4/14 17:14 # ------------------------------------------- import json +import traceback from bson import ObjectId from fastapi import APIRouter, Depends @@ -74,7 +75,7 @@ async def asset_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Dep .skip((page_index - 1) * page_size) .limit(page_size)) .sort([("timestamp", DESCENDING)])) - result = await cursor.to_list(length=None) + result = cursor.to_list(length=None) result_list = [] for r in result: tmp = {} @@ -127,6 +128,7 @@ async def asset_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Dep } except Exception as e: logger.error(str(e)) + logger.error(traceback.format_exc()) # Handle exceptions as needed return {"message": "error", "code": 500} diff --git a/api/export.py b/api/export.py index 33f5a4a..101bb52 100644 --- a/api/export.py +++ b/api/export.py @@ -5,6 +5,7 @@ # @time : 2024/6/16 16:11 # ------------------------------------------- import os +import traceback from bson import ObjectId from fastapi import APIRouter, Depends, BackgroundTasks @@ -150,7 +151,7 @@ async def fetch_data(db, collection, query, quantity, project_list): {"$project": {"_id": 0, "vulnid": 0}} ] - cursor = db[collection].aggregate(pipeline) + cursor = await db[collection].aggregate(pipeline) return cursor @@ -300,6 +301,7 @@ async def export_data_from_mongodb(quantity, query, file_name, index): await db.export.update_one({"file_name": file_name}, update_document) except Exception as e: logger.error(str(e)) + logger.error(traceback.format_exc()) update_document = { "$set": { "state": 2, diff --git a/api/task.py b/api/task.py index 6a6682f..7d7da5b 100644 --- a/api/task.py +++ b/api/task.py @@ -4,6 +4,8 @@ # @version: import asyncio import json +import traceback + from loguru import logger from bson import ObjectId from fastapi import APIRouter, Depends, BackgroundTasks @@ -305,9 +307,13 @@ async def update_task_data(request_data: dict, db=Depends(get_mongo_db), _: dict 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(job) + 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') @@ -324,7 +330,7 @@ async def update_task_data(request_data: dict, db=Depends(get_mongo_db), _: dict except Exception as e: logger.error(str(e)) - # Handle exceptions as needed + logger.error(traceback.format_exc()) return {"message": "error", "code": 500}