diff --git a/mycode/AttackMap.py b/mycode/AttackMap.py index c54a945..d4f145d 100644 --- a/mycode/AttackMap.py +++ b/mycode/AttackMap.py @@ -197,6 +197,8 @@ class TreeNode: self._instr_queue = [] #针对当前节点的待执行指令----重要约束：一个节点只能有一个线程在执行指令 self.his_instr = [] #保留执行指令的记录｛“instr”:***,"result":***｝ + #单步相关 + self.step_num = 0 #单步执行次数 #用户补充信息 self.cookie = "" self.ext_info = "" diff --git a/mycode/TaskManager.py b/mycode/TaskManager.py index 0ec44d8..57e9f17 100644 --- a/mycode/TaskManager.py +++ b/mycode/TaskManager.py @@ -220,10 +220,10 @@ class TaskManager: return bsuccess,error #任务单点--只允许web端调用 - async def task_one_step(self,task_id): + async def task_one_step(self,task_id,step_num): task = self.tasks[task_id] if task: - bsuccess,error = await task.put_one_task() + bsuccess,error = await task.put_one_task(step_num) return bsuccess,error else: return False,"task_id值存在问题！" diff --git a/mycode/TaskObject.py b/mycode/TaskObject.py index 50fc6f4..0f3dfa9 100644 --- a/mycode/TaskObject.py +++ b/mycode/TaskObject.py @@ -61,9 +61,11 @@ class TaskObject: self.llmth_list = [None] * self.llm_max_nums # llm线程list self.doing_llm_list = [""] * self.llm_max_nums self.llm_node_queue = queue.Queue() #待提交LLM的节点队列 - #自检线程-------- + #自检线程----------- self.check_th = None #自检线程句柄 - #-----四队列----- + #单步执行控制------- + self.one_step_num = 0 #单步执行任务次数，默认1 + #-----四队列------- self.run_instr_lock = threading.Lock() # 线程锁 self.runing_instr = {} #执行中指令记录 @@ -264,18 +266,30 @@ class TaskObject: self.is_had_work = True return bsuccess - def no_work_to_do(self): #任务单步状态控制-- 非工作中 + def no_work_to_do(self): #任务单步状态控制-- 非工作中--2025-5-7增加了轮次，需要全面验证执行逻辑的完整和正确性 # 待执行instr-node instr_node_list = list(self.instr_node_queue.queue) # 待执行指令的node--线程不安全 llm_node_list = list(self.llm_node_queue.queue) # 待提交llm的node--线程不安全 if len(instr_node_list) == 0 and len(llm_node_list) == 0: #没有待办节点了 + #遍历在执行 for str_instr in self.doing_instr_list: if str_instr != "": return for str_llm in self.doing_llm_list: if str_llm != "": return - #没有在执行任务了 + + # 没有在执行任务了 + self.one_step_num -= 1 + if not self.one_step_num: # 执行轮次已结束 + pass + else:#发起新的一轮 + bok = self.put_one_task() + if not bok: + self.one_step_num = 0 + #结束轮次--推送前端 + + #推送是否有工作任务的状态到前端, with self.is_had_work_lock: if self.is_had_work: #如果已经是False那就不需要修改了 @@ -483,14 +497,20 @@ class TaskObject: return False,"当前的任务或节点状态不允许执行单步，请检查！" #web端提交任务单步--任务单步 - async def put_one_task(self): + async def put_one_task(self,step_num): if self.task_status == 1 and self.work_type == 0: bsuccess = self.had_work_to_do() if bsuccess: nodes = self.attack_tree.traverse_bfs() + b_putwork = False for node in nodes: - _,_ = await self.put_one_node(node) - return True,"已提交单步任务" + bput,_ = await self.put_one_node(node) + if bput: + b_putwork = True + if b_putwork: + return True,"已提交单步任务" + else: + return False,"该任务已经没有待提交任务" else: return False,"当前任务正在执行任务中，不需要提交单步任务！" else: diff --git a/web/API/task.py b/web/API/task.py index 94342c9..ecb0613 100644 --- a/web/API/task.py +++ b/web/API/task.py @@ -137,9 +137,12 @@ async def task_one_step(): ''' data = await request.get_json() task_id = data.get("cur_task_id") + step_num = data.get("step_num") if not task_id: return jsonify({'error': 'Missing task_id'}), 400 - bsuccess,error = await g_TaskM.task_one_step(task_id) + if not step_num: + step_num = 1 #默认一轮 + bsuccess,error = await g_TaskM.task_one_step(task_id,step_num) return jsonify({"bsuccess":bsuccess,"error":error}) @api.route('/task/nodestep',methods=['POST']) diff --git a/web/main/templates/footer.html b/web/main/templates/footer.html index b411273..4faf7b2 100644 --- a/web/main/templates/footer.html +++ b/web/main/templates/footer.html @@ -14,7 +14,7 @@