from myutils.ConfigManager import myCongif from mycode.TaskObject import TaskObject from mycode.DBManager import app_DBM from myutils.PickleManager import g_PKM import threading class TaskManager: def __init__(self): self.tasks = {} # 执行中的任务，test_id为key self.num_threads = myCongif.get_data("Task_max_threads") #获取系统信息 -- 用户可修改的都放在DB中，不修改的放config data = app_DBM.get_system_info() self.local_ip = data[0] self.version = data[1] self.tasks_lock = threading.Lock() #加个线程锁?不使用1，quart要使用的异步锁，2.限制只允许一个用户登录，3.遍历到删除的问题不大 self.web_cur_task = 0 #web端当前显示的 #判断目标是不是在当前执行任务中，---没加锁，最多跟删除有冲突，问题应该不大 def is_target_in_tasks(self,task_target): for task in self.tasks.values(): if task_target == task.target: return True return False #程序启动后，加载未完成的测试任务 def load_tasks(self): '''程序启动时，加载未执行完成的任务''' datas = app_DBM.get_run_tasks() for data in datas: task_id = data[0] task_target = data[1] task_status = data[2] work_type = data[3] cookie_info = data[4] llm_type = data[5] # 创建任务对象 task = TaskObject(task_target, cookie_info, work_type, llm_type, self.num_threads, self.local_ip,self) #读取attact_tree attack_tree = g_PKM.ReadData(str(task_id)) #开始任务 ---会根据task_status来判断是否需要启动工作线程 task.start_task(task_id,task_status,attack_tree) # 保留task对象 self.tasks[task_id] = task #新建测试任务 def create_task(self, test_target,cookie_info,llm_type,work_type): """创建新任务--create和load复用？-- 1.创建和初始化task_object; 2.创建task_id 3.启动工作线程 return T/F """ if self.is_target_in_tasks(test_target): raise ValueError(f"Task {test_target} already exists") #创建任务对象 task = TaskObject(test_target,cookie_info,work_type,llm_type,self.num_threads,self.local_ip,self) #获取task_id -- test_target,cookie_info,work_type,llm_type 入数据库 task_id = app_DBM.start_task(test_target,cookie_info,work_type,llm_type) if task_id >0: #创建后启动工作--同时赋值task_id task.start_task(task_id) #保留task对象 self.tasks[task_id] = task return True else: return False #控制task启停----线程不停 def control_taks(self,task_id): task = self.tasks[task_id] if task: if task.task_status == 0: # 0-暂停，1-执行中，2-已完成 task.task_status = 1 elif task.task_status == 1: task.task_status = 0 else: return False,"当前任务状态不允许修改，请联系管理员！",task.task_status else: return False,"没有找到对应的任务",None return True,"",task.task_status # 获取任务list def get_task_list(self): tasks = [] for task in self.tasks.values(): one_json = {"taskID": task.task_id, "testTarget": task.target, "taskStatus": task.task_status, "safeRank": task.safe_rank, "workType": task.work_type} tasks.append(one_json) rejson = {"tasks": tasks} return rejson #获取节点树 def get_node_tree(self,task_id): task = self.tasks[task_id] if task: self.web_cur_task = task_id tree_dict = task.attack_tree.get_node_dict() return tree_dict return None #修改任务的工作模式，只有在暂停状态才能修改 def update_task_work_type(self,task_id,new_work_type): task = self.tasks[task_id] if task: if task.task_status == 0: task.work_type = new_work_type return True return False #控制节点的工作状态 def node_bwork_control(self,task_id,node_path): task = self.tasks[task_id] if task: bsuccess,new_bwork = task.attack_tree.update_node_bwork(node_path) if bsuccess: pass #是否要更新IO数据？----待验证是否有只修改部分数据的方案 return bsuccess,new_bwork return False,False #节点单步--只允許web端调用 async def node_one_step(self,task_id,node_path): task = self.tasks[task_id] node = task.attack_tree.find_node_by_nodepath(node_path) #web端触发的任务，需要判断当前的执行状态 bsuccess,error = await task.put_one_node(node) return bsuccess,error #任务单点--只允许web端调用 async def task_one_step(self,task_id): task = self.tasks[task_id] if task: bsuccess,error = await task.put_one_task() return bsuccess,error else: return False,"task_id值存在问题！" #获取节点待执行任务 def get_task_node_todo_instr(self,task_id,nodepath): todoinstr = [] task = self.tasks[task_id] if task: node = task.attack_tree.find_node_by_nodepath(nodepath) if node: todoinstr = node.get_instr_user() return todoinstr #获取节点的MSG信息 def get_task_node_MSG(self,task_id,nodepath): task = self.tasks[task_id] if task: node = task.attack_tree.find_node_by_nodepath(nodepath) if node: tmpMsg = node.get_res_user() if tmpMsg: return node.messages,tmpMsg[0] #待提交消息正常应该只有一条 else: return node.messages,{} return [],{} def update_node_MSG(self,task_id,nodepath,newtype,newcontent): task = self.tasks[task_id] if task: node = task.attack_tree.find_node_by_nodepath(nodepath) if node: work_status = node.get_work_status() if work_status == 0 or work_status == 3: bsuccess,error = node.updatemsg(newtype,newcontent,0) #取的第一条，也就修改第一条 return bsuccess,error else: return False,"当前节点的工作状态不允许修改MSG！" return False,"找不到对应节点！" def del_node_instr(self,task_id,nodepath,instr): task = self.tasks[task_id] if task: node = task.attack_tree.find_node_by_nodepath(nodepath) if node: return node.del_instr(instr) return False,"找不到对应节点！" #------------以下函数还未验证处理----------- def start_task(self, task_id): """启动指定任务""" task = self.tasks.get(task_id) if task: task.start(self.num_threads) else: print(f"Task {task_id} not found") def stop_task(self, task_id): """停止指定任务""" task = self.tasks.get(task_id) if task: task.stop() else: print(f"Task {task_id} not found") g_TaskM = TaskManager() #单一实例