from myutils.ConfigManager import myCongif from mycode.TaskObject import TaskObject from mycode.DBManager import app_DBM from myutils.PickleManager import g_PKM from myutils.MyLogger_logger import LogHandler from mycode.TargetManager import TargetManager # 从模块导入类 from mycode.PythonTManager import PythonTManager import time import threading class TaskManager: def __init__(self): self.logger = LogHandler().get_logger("TaskManager") self.TargetM = TargetManager() #获取系统信息 -- 用户可修改的都放在DB中，不修改的放config self.num_threads = myCongif.get_data("Task_max_threads") data = app_DBM.get_system_info() self.local_ip = data[0] self.version = data[1] self.brun = True self.tm_th = None #任务控制线程 #----临时任务相关------- self.web_cur_task = 0 #web端当前显示的 -- web端可以操作的都为临时任务 self.tasks_lock = threading.Lock() self.tasks = {} # 执行中的任务，test_id为key self.max_run_task = myCongif.get_data("max_run_task") self.cur_task_run_num = 0 #-----巡检任务相关------- self.polling_llmtype = myCongif.get_data("polling_llm_type") self.polling_tasks_lock = threading.Lock() self.polling_tasks = {} #独立一个队列 self.max_polling_task = myCongif.get_data("max_polling_task") self.cur_polling_task_num = 0 self.work_type = 1 #巡检工作模式就是自动了 #python 进程池 提到程序主线程层面 --不随任务进行启停 self.PythonM = PythonTManager(myCongif.get_data("Python_max_procs")) # 启动python子进程池 --进程池的启动，目前是没全面确认原子进程池的状态，直接None self.PythonM.start_pool() #开启进程池 def __del__(self): # 停止子进程池 self.PythonM.shutdown_pool() #判断目标是不是在当前执行任务中，---没加锁，最多跟删除有冲突，问题应该不大 def is_target_in_tasks(self,task_target,itype = 1): '''itype：1--临时任务，2--巡检任务''' if itype ==1: for task in self.tasks.values(): if task_target == task.target: return True else: for p_task in self.polling_tasks.values(): if task_target == p_task.target: return True return False #程序启动后，加载未完成的测试任务 def load_tasks(self): '''程序启动时，加载未执行完成的,未点击结束的任务 -- task_status<>4 #若不是异常停止，正常情况下，任务都应该是没有待办MQ的 ''' 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] safe_rank = data[6] fake_target = data[7] target_id = data[8] # 创建任务对象 task = TaskObject(task_target, cookie_info, work_type, llm_type, self.num_threads, self.local_ip,fake_target,self,target_id,safe_rank) #读取attact_tree---load的任务应该都要有attact_tree attack_tree = g_PKM.ReadData(str(task_id)) if attack_tree: #恢复数据--遍历节点恢复状态 --临时和polling适用 task.init_task(task_id,attack_tree) #开始任务 ---根据task_status来判断是否需要启动工作线程 if task_status == 1: #执行中的任务 bcan_start = False if target_id == 0: #临时任务 if self.cur_task_run_num < self.max_run_task: #load 是程序刚起，只有主线程，不加锁 bcan_start = True else: #巡检任务 if self.cur_polling_task_num < self.max_polling_task: bcan_start = True if bcan_start: isuc, strout = task.start_task() if isuc == 1: #启动工作线程成功 if target_id == 0: self.cur_task_run_num += 1 else: self.cur_polling_task_num += 1 elif isuc == -1: # 存在未结束的线程--load这边应该不会发生 self.logger.debug(f"load-task遇到未结束线程的任务--{task_id}") task.update_task_status(2) else: # 0 --没有待办事项 self.logger.debug(f"load-task遇到无待办实现的任务--{task_id}") task.update_task_status(3) else: task.update_task_status(0) #等其他task结束后会启动 # 内存保留task对象 if target_id == 0: self.tasks[task_id] = task else: self.polling_tasks[task_id] = task else: self.logger.error(f"{task_id}任务的节点树数据缺失，需要检查!") #新建测试任务--2025-4-28调整为可以批量添加--cookie_info信息取消 def create_task(self, test_targets,llm_type,work_type): """创建新任务--create和load复用？-- 1.创建和初始化task_object; 2.创建task_id return 失败的list """ fail_list = [] target_list = test_targets.split(",") for target in target_list: #这里判断目标的合法性 bok,target,type,fake_target = self.TargetM.validate_and_extract(target) #若是url，target是域名部分 if not bok:#目标不合法 fail_list.append(target) continue #判断是否已在执行列表 if self.is_target_in_tasks(target,1): fail_list.append(target) continue #raise ValueError(f"Task {test_target} already exists") #创建任务对象--cookie参数取消 task = TaskObject(target,"",work_type,llm_type,self.num_threads,self.local_ip,fake_target,self) #获取task_id -- test_target,cookie_info,work_type,llm_type 入数据库 task_id = app_DBM.start_task(target,"",work_type,llm_type,fake_target) if task_id >0: #2025-4-28调整批量添加任务，默认不启动线程start_task task.init_task(task_id) #保留task对象 self.tasks[task_id] = task #尝试启动task _,_ = self.start_task_TM(task_id) else: fail_list.append(target) result = ",".join(fail_list) return result def create_polling_task(self,target): #创建巡检任务 -- 使用check_target task_target = target[2] if self.is_target_in_tasks(task_target,2): #巡检任务一个周期还没有执行完。 #更新检查时间,本次巡检轮次不执行工作 bok = app_DBM.update_polling_last_time(target[0]) else:#创建任务 if target[9] == "IPv4" or target[9] == "IPv6": fake_target = self.TargetM.get_fake_target(1) else: fake_terget = self.TargetM.get_fake_target(2) #创建任务实例 polling_task = TaskObject(target, "", self.work_type, self.polling_llmtype, self.num_threads, self.local_ip, fake_target, self,target[0]) # 获取task_id -- test_target,cookie_info,work_type,llm_type 入数据库 task_id = app_DBM.start_task(target, "", self.work_type, self.polling_llmtype, fake_target,target[0]) #增加一个target_id if task_id >0: polling_task.init_task(task_id)#初始化任务信息，并提交到待处理节点队列 #保留task对象--巡检任务 self.polling_tasks[task_id] = polling_task #尝试启动task _,_ = self.start_polling_task_TM(polling_task) def is_can_start_task(self,itype): if itype ==1: if self.cur_task_run_num < self.max_run_task: return True else: if self.cur_polling_task_num < self.max_polling_task: return True return False def start_next_task(self,task_id,target_id): if target_id == 0: #临时任务 with self.tasks_lock: self.cur_task_run_num -= 1 for task in self.tasks.values(): if task.task_status == 0: #只有是0状态会自动启动，2暂停状态需要人工启动 isuc, error = task.start_task() if isuc == 1: self.cur_task_run_num +=1 break else: #巡检任务--是直接结束置4 with self.polling_tasks_lock: del self.polling_tasks[task_id] # 删除缓存 self.cur_polling_task_num -=1 for pl_task in self.polling_tasks: if pl_task.task_status == 0: isuc ,error = pl_task.start_task() if isuc == 1: self.cur_polling_task_num += 1 break #开启task任务 def start_task_TM(self,task_id): task = self.tasks[task_id] if task: with self.tasks_lock: if self.cur_task_run_num < self.max_run_task: #启动工作线程和进程 isuc,error = task.start_task() if isuc ==1: self.cur_task_run_num += 1 return True,"启动成功" else: return False,error else: return False,f"已到达最大的启动数--{self.max_run_task}" return False,"该任务不存在，程序逻辑存在问题！" def start_polling_task_TM(self,pl_task): if pl_task: with self.polling_tasks_lock: if self.cur_polling_task_num < self.max_polling_task: isuc,error = pl_task.start_task() if isuc == 1: self.cur_polling_task_num += 1 return True,"启动成功" else: return False, error else: return False, f"已到达最大的启动数--{self.max_run_task}" #停止task任务--暂停---要启动下一个task---只针对临时任务 def stop_task_TM(self,task_id): task = self.tasks[task_id] if task: if task.brun and task.task_status ==1: #是运行状态 with self.tasks_lock: task.stop_task() #停止线程应该就没什么失败需要处理的 self.cur_task_run_num -= 1 #启动下一个任务-只针对临时任务 for tmptask in self.tasks.values(): if tmptask.task_status == 0: # 只有是0状态会自动启动，2暂停状态需要人工启动 isuc, error = tmptask.start_task() if isuc == 1: self.cur_task_run_num += 1 break return True,"停止任务成功" else: return True,"该任务已处于停止状态" return False,"该任务不存在，程序逻辑存在问题！" #结束任务-只针对临时任务 def over_task(self,task_id): #先尝试停止 bsuccess,_ = self.stop_task_TM(task_id) time.sleep(1) if bsuccess: #再结束 bsuccess = app_DBM.over_task(task_id) # 不管是不是修改（置4）成功，都执行结束 del self.tasks[task_id] # 删除缓存 return True,"结束任务成功" else: return False,"该任务不存在，程序逻辑存在问题！" #删除任务--历史任务处理 def del_task(self,task_id): #删除数据记录 app_DBM.del_task(task_id) #删除节点树 g_PKM.DelData(str(task_id)) return True,"删除任务成功！" #控制task启停----线程不停 --2025-4-28 配合批量任务，需要停止工作线程了 def control_taks(self,task_id): task = self.tasks[task_id] if task: if task.task_status in (0,2): # 0-未启动，1-执行中，2-暂停中，3-已完成，4-已结束 bsuc,error = self.start_task_TM(task_id) #任务是否存在的状态有点重复 elif task.task_status == 1: bsuc,error = self.stop_task_TM(task_id) else: return False,"当前任务状态不允许修改，请联系管理员！",task.task_status else: return False,"没有找到对应的任务",None return bsuc,error,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 get_his_node_tree(self,task_id): attack_tree = g_PKM.ReadData(str(task_id)) if attack_tree: tree_dict = 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 in (0,2): task.update_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,step_num): task = self.tasks[task_id] if task: bsuccess,error = await task.put_one_task(step_num) 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.cur_messages,tmpMsg[0] #待提交消息正常应该只有一条 else: return node.cur_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: bsuccess,error = node.updatemsg(newtype,newcontent,node.parent.parent_messages,node.parent.cur_messages,0) return bsuccess,error 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 get_his_tasks(self,target_name,safe_rank,llm_type,start_time,end_time): tasks = app_DBM.get_his_tasks(target_name,safe_rank,llm_type,start_time,end_time) return tasks g_TaskM = TaskManager() #单一实例