zfsafe/mycode/TaskObject.py

'''
渗透测试任务管理类  一次任务的闭合性要检查2025-3-10  一次任务后要清理LLM和InstrM的数据
'''
from mycode.TargetManager import TargetManager  # 从模块导入类
#from LLMManager import LLMManager        # 同理修正其他导入
from mycode.ControlCenter import ControlCenter         #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。
from mycode.InstructionManager import g_instrM
from mycode.DBManager import DBManager,app_DBM
from mycode.LLMManager import LLMManager
from mycode.AttackMap import AttackTree,TreeNode
from myutils.MyTime import get_local_timestr
from myutils.MyLogger_logger import LogHandler
from myutils.PickleManager import g_PKM
from myutils.ConfigManager import myCongif
from mycode.WebSocketManager import g_WSM
from mycode.CommandVerify import g_CV
from mycode.PythonTManager import PythonTManager
import asyncio
import queue
import time
import os
import re
import threading
import json

class TaskObject:

    def __init__(self,test_target,cookie_info,work_type,llm_type,num_threads,local_ip,taskM):
        #功能类相关
        self.taskM = taskM
        self.logger = LogHandler().get_logger("TaskObject")
        self.InstrM = g_instrM  # 类对象渗透，要约束只读取信息，且不允许有类全局对象--持续检查
        self.PythonM = PythonTManager(myCongif.get_data("Python_max_procs"))
        self.CCM = ControlCenter() #一个任务一个CCM
        self.LLM = LLMManager(llm_type)  # LLM对象调整为一个任务一个对象，这样可以为不同的任务选择不同的LLM
        #全局变量
        self.app_work_type = myCongif.get_data("App_Work_type") #app工作为0时，只允许单步模式工作，是附加规则，不影响正常逻辑处理
        self.brun = False           #任务的停止可以用该变量来控制
        self.sleep_time = myCongif.get_data("sleep_time")
        self.target = test_target
        self.cookie = cookie_info
        self.work_type = work_type  #工作模式 0-人工，1-自动
        self.task_id = None
        self.task_status = 0         #0-暂停，1-执行中，2-已完成
        self.local_ip = local_ip
        self.attack_tree = None     #任务节点树
        self.safe_rank = 0          #安全级别 0-9   #？暂时还没实现更新逻辑

        #指令执行相关-------
        self.max_thread_num = num_threads   #指令执行线程数量
        self.workth_list = [None] * num_threads   #线程句柄list
        self.doing_instr_list= [""] * num_threads
        self.instr_node_queue = queue.Queue()     #待执行指令的节点队列
        self.node_num = 0           #在处理Node线程的处理
        #llm执行相关--------
        self.llm_max_nums = myCongif.get_data("LLM_max_threads")  # 控制最大并发指令数量 --- 多线程的话节点树需要加锁
        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.run_instr_lock = threading.Lock()  # 线程锁
        self.runing_instr = {}  #执行中指令记录

    #---------------三个线程------------
    #测试指令执行线程
    def mill_instr_preprocess(self,instructions,str_split):
        new_instr = []
        instrs = instructions.split(str_split)
        index = 0
        for instr in instrs:
            if instr.strip().startswith("curl"):
                if " --max-time " not in instr:
                    out_time = g_instrM.get_tool_out_time("curl")
                    instr = instr.strip() + f" --max-time {str(out_time)}"
                    #instr = instr.strip() + " --max-time 10"
            new_instr.append(instr)
            index += 1
        new_star_instr = f"{str_split}".join(new_instr)
        print(new_star_instr)
        return new_star_instr

    def do_worker_th(self,index):
        #线程的dbm需要一个线程一个
        th_DBM = DBManager()
        th_DBM.connect()
        th_index = index
        while self.brun:
            if self.task_status == 1:
                try:
                    work_node = self.instr_node_queue.get(block=False)#正常一个队列中一个节点只会出现一次，进而也就只有一个线程在处理
                    # 开始执行指令
                    results = []
                    while True:
                        instruction = work_node.get_instr()
                        if not instruction:
                            break
                        instruction = instruction.strip()
                        #对多shell指令的情况进行处理--也有风险
                        if "python-code" not in instruction:
                            if "&&" in instruction:
                                instruction = self.mill_instr_preprocess(instruction, "&&")
                            elif "||" in instruction:
                                instruction = self.mill_instr_preprocess(instruction, "||")

                        start_time = get_local_timestr()  # 指令执行开始时间
                        self.doing_instr_list[th_index] = instruction
                        if instruction.startswith("python-code"):#python代码--超过子进程数会阻塞等待，但不开始超时记时
                            instr, reslut, source_result, ext_params = self.PythonM.execute_instruction(instruction)
                        else:#shell指令
                            instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction)
                        self.doing_instr_list[th_index] = ""
                        end_time = get_local_timestr()  # 指令执行结束时间

                        # 入数据库 -- bres True和False 都入数据库2025-3-10---加node_path(2025-3-18)#？
                        if th_DBM.ok:
                            work_node.do_sn += 1
                            th_DBM.insetr_result(self.task_id, instr, reslut, work_node.do_sn, start_time, end_time,
                                                 source_result,
                                                 ext_params, work_node.path)
                        else:
                            self.logger.error("数据库连接失败！！")
                        #暂存结果
                        oneres = {'执行指令': instr, '结果': reslut}
                        results.append(oneres)
                    #指令都执行结束后，入节点待提交队列
                    str_res = json.dumps(results, ensure_ascii=False)
                    # 提交llm待处理任务 --更新节点work_status
                    self.put_node_reslist(work_node, str_res, 1)
                    # 保存记录
                    g_PKM.WriteData(self.attack_tree,str(self.task_id))
                except queue.Empty:
                    time.sleep(self.sleep_time)
            else:#不是工作状态则休眠
                time.sleep(self.sleep_time)

    #llm请求提交线程
    def th_llm_worker(self,index):
        '''
        几个规则--TM的work线程同
        1.线程获取一个节点后，其他线程不能再获取这个节点（遇到被执行的节点，直接放弃执行）--- 加了没办法保存中间结果进行测试
        2.llm返回的指令，只可能是该节点，或者是子节点的，不符合这条规则的都不处理，避免llm处理混乱。
        :return:
        '''
        # 线程的dbm需要一个线程一个
        th_DBM = DBManager()
        th_DBM.connect()
        th_index = index
        while self.brun:
            if self.task_status == 1:
                try:
                    llm_node = self.llm_node_queue.get(block=False)  #获取一个待处理节点
                    #开始处理
                    tmp_commands = []
                    # {llm_node.status} --- 暂时固化为未完成
                    user_Prompt = f'''
当前分支路径：{llm_node.path}
当前节点信息：
- 节点名称：{llm_node.name}
- 节点状态：未完成
- 漏洞类型：{llm_node.vul_type}
                        '''
                    while True:
                        llm_data = llm_node.get_res()
                        if llm_data is None:
                            break
                        llm_type = llm_data["llm_type"]
                        str_res = llm_data["result"]
                        #获取提示词
                        prompt = self.get_llm_prompt(llm_type,str_res,user_Prompt)
                        self.doing_llm_list[th_index] = prompt
                        # 提交llm请求返回数据--并对返回数据进行处理，节点指令直接执行，测试指令根据工作模式执行
                        node_cmds, commands,reasoning_content, content, post_time = self.LLM.get_llm_instruction(prompt,llm_node)  # message要更新
                        self.doing_llm_list[th_index] = ""
                        # LLM记录存数据库
                        if th_DBM.ok:
                            llm_node.llm_sn += 1
                            bres = th_DBM.insert_llm(self.task_id, prompt, reasoning_content, content, post_time, llm_node.llm_sn,llm_node.path)
                            if not bres:
                                self.logger.error(f"{llm_node.name}-llm入库失败！")
                        else:
                            self.logger.error("数据库连接失败！")
                        '''
                        对于LLM返回的错误处理机制
                        1.验证节点是否都有测试指令返回
                        2.LLM的回复开始反复时（有点难判断）
                        '''
                        # 更新tree
                        bok, new_commands,iadd_node = self.tree_manager(node_cmds, llm_node, commands, th_DBM)
                        # 分析指令入对应节点
                        if bok:  # 节点指令若存在错误，测试指令都不处理，需要LLM重新生成
                            tmp_commands.extend(new_commands)
                    #测试指令入节点待处理队列 --同时修改节点的work_status
                    self.put_node_instrlist(tmp_commands, llm_node,iadd_node)
                    #一个节点完成，节点树持久化---待验证是否有局部更新持久化的方案
                    g_PKM.WriteData(self.attack_tree,str(self.task_id))
                except queue.Empty:
                    #self.logger.debug("llm队列中暂时无新的提交任务！")
                    time.sleep(self.sleep_time)
            else:
                time.sleep(self.sleep_time)

    #自检线程 --1.输出执行状态。2.需要自检和修复
    def th_check(self):
        while self.brun:
            try:
                cur_time = get_local_timestr()
                print(f"-----------当前时间程序运行情况:{cur_time}")
                # #待执行instr-node
                # instr_node_list = list(self.instr_node_queue.queue) #待执行指令的node--线程不安全
                # print(f"**当前待执行指令node的数量为:{len(instr_node_list)}")
                #执行中instr-node
                index = 0
                for w_th in self.workth_list:
                    if not w_th.is_alive():#线程
                        print(f"线程-{index}已处于异常状态，需要重新创建一个工作线程")
                    else:
                        print(f"线程-{index}在执行指令：{self.doing_instr_list[index]}")
                    index += 1

                index = 0
                for l_th in self.llmth_list:
                    if not l_th.is_alive():
                        print(f"LLM线程-{index}已处于异常状态，需要重新创建一个LLM线程")
                    else:
                        print(f"LLM线程-{index}在执行指令：{self.doing_llm_list[index]}")
                    index += 1
                #待提交llm-node
                # llm_node_list = list(self.llm_node_queue.queue) #待提交llm的node--线程不安全
                # print(f"**当前待提交llm的node的数量为:{len(llm_node_list)}")
                #休眠60
                time.sleep(60)
            except Exception as e:
                print(f"*********自检线程异常退出:{str(e)}")
                break

    #------------入两个nodeMQ-禁止直接调用入队列-----------
    def put_instr_mq(self,node):
        #这里不做状态的判断，调用前处理
        self.instr_node_queue.put(node)
        self.update_node_work_status(node,2)  #在执行--1.work_status不影响整个任务的执行，错了问题不大，2--attack_tree持久化需要出去lock信息。

    def put_llm_mq(self,node):
        #同instr_mq
        self.llm_node_queue.put(node)
        self.update_node_work_status(node,4)  #提交中

    async def put_instr_mq_async(self,node):
        #这里不做状态的判断，调用前处理
        self.instr_node_queue.put(node)
        await self.update_node_work_status_async(node,2)  #在执行--1.work_status不影响整个任务的执行，错了问题不大，2--attack_tree持久化需要出去lock信息。

    async def put_llm_mq_async(self,node):
        #同instr_mq
        self.llm_node_queue.put(node)
        await self.update_node_work_status_async(node,4)  #提交中

    async def update_node_work_status_async(self,node,work_status):
        #更新状态
        bchange = node.update_work_status(work_status)
        #基于websocket推送到前端
        if bchange:
            #判断是否是web端最新获取数据的task
            if self.taskM.web_cur_task == self.task_id:
                idatatype = 1
                strdata = {"node_path":node.path,"node_workstatus":work_status}
                await g_WSM.send_data(idatatype,strdata)

    #------------入Node的两个list--禁止直接调用入list-------
    def put_node_reslist(self, node, str_res, llm_type):
        # 推送llm提交任务到节点的待处理任务中，并根据工作模式判断是否如llm_node_quere
        one_llm = {'llm_type': llm_type, 'result': str_res}
        node.add_res(one_llm)  # 入节点结果队列
        self.update_node_work_status(node,3)    #待提交llm
        # 如果是自动执行的模式则入队列交给llm线程处理
        if self.app_work_type == 1:
            if self.work_type == 1 and node.bwork:
                self.put_llm_mq(node)   #变4

    #递归找节点
    def find_node_by_child_node_name(self,cur_node,node_name):
        find_node = None
        if cur_node.children:
            for child_node in cur_node.children:
                if child_node.name == node_name:
                    find_node = child_node
                    break
                else:
                    find_node = self.find_node_by_child_node_name(child_node,node_name)
                    if find_node:
                        break
        return find_node

    def put_node_instrlist(self, commands, node,iadd_node):       #如果当前节点没有进一般指令返回，需要修改节点执行状态
        node_list = []  #一次返回的测试指令
        for command in commands:
            # 使用正则匹配方括号中的node_path（非贪婪模式）
            match = re.search(r'\[(.*?)\]', command)
            if match:
                node_path = match.group(1)
                instruction = re.sub(r'\[.*?\]', "", command, count=1, flags=re.DOTALL)
                #'''强制约束，不是本节点或者是子节点的指令不处理'''
                find_node = self.attack_tree.find_node_by_nodepath_parent(node_path,node,iadd_node,commands)
                if not find_node:#对于没有基于节点路径找到对应节点--增加通过节点名称匹配的机制 2025-4-13日添加
                    node_name = node_path.split("->")[-1]
                    find_node = self.find_node_by_child_node_name(node,node_name)

                if find_node:
                    find_node.add_instr(instruction)
                    #DS-llm存在返回指令还会修改节点状态为已完成的问题，需要修正
                    find_node.status = "未完成"
                    if find_node not in node_list:
                        node_list.append(find_node)
                        self.update_node_work_status(find_node,1)   #待执行
                else:#如果还没找到就暂时放弃
                    self.logger.error(f"基于节点路径没有找到对应的节点{node_path},父节点都没找到！当前节点{node.path}")#丢弃该指令
            else:
                self.logger.error(f"得到的指令格式不符合规范：{command}")#丢弃该指令---
                #这里对于丢弃指令，有几种方案：
                # 1.直接丢弃不处理，但需要考虑会不会产生节点缺失指令的问题，需要有机制验证节点；------ 需要有个独立线程，节点要加锁--首选待改进方案
                # 2.入当前节点的res_queue,但有可能当前节点没有其他指令，不会触发提交，另外就算提交，会不会产生预设范围外的返回，不确定；
                # 3.独立队列处理
        #判断当前节点是否有指令
        if node not in node_list:
            #修改该节点状态为0--无待执行任务
            self.update_node_work_status(node,0)
        #入instr队列
        if self.app_work_type == 1:
            if self.work_type == 1: #是自动执行模式
                for node in node_list:
                    if node.bwork:
                        self.put_instr_mq(node)     #2-执行中

    def put_work_node(self):
        '''遍历节点需要处理的任务，提交mq'''
        nodes = self.attack_tree.traverse_bfs()
        for node in nodes:
            if not node.is_instr_empty():   #待执行指令有值
                if not node.is_llm_empty():
                    self.logger.error(f"{node.path}即存在待执行指令，还存在待提交的llm，需要人工介入！！")
                else:
                    if node.bwork:
                        self.put_instr_mq(node) #提交执行
            elif not node.is_llm_empty():   #待提交llm有值
                if not node.is_instr_empty():
                    self.logger.error(f"{node.path}即存在待执行指令，还存在待提交的llm，需要人工介入！！")
                else:
                    if node.bwork:
                        self.put_llm_mq(node)   #提交执行

    #web端提交单步任务--节点单步
    async def put_one_node(self,node):
        #提交某个节点的代表任务
        if self.task_status ==1 and self.work_type==0 and node.bwork:
            if not node.is_instr_empty():   #待执行指令有值
                if not node.is_llm_empty():
                    self.logger.error(f"{node.path}即存在待执行指令，还存在待提交的llm，需要人工介入！！")
                    return False,"该节点的待执行任务数据不正确，请联系管理员！"
                else:
                    if node.bwork:
                        await self.put_instr_mq_async(node) #提交执行
            elif not node.is_llm_empty():   #待提交llm有值
                if not node.is_instr_empty():
                    self.logger.error(f"{node.path}即存在待执行指令，还存在待提交的llm，需要人工介入！！")
                    return False, "该节点的待执行任务数据不正确，请联系管理员！"
                else:
                    if node.bwork:
                        await self.put_llm_mq_async(node)   #提交执行
            else:
                await self.update_node_work_status_async(node,0)  #只是修补措施，保障状态的一致性
                return False,"当前节点没有待执行任务！"
            return True,"已提交单步任务"
        else:
            return False,"当前的任务或节点状态不允许执行单步，请检查！"

    #web端提交任务单步--任务单步
    async def put_one_task(self):
        if self.task_status == 1 and self.work_type == 0:
            nodes = self.attack_tree.traverse_bfs()
            for node in nodes:
                await self.put_one_node(node)
            return True,"已提交单步任务"
        else:
            return False,"当前的任务状态不允许执行单步，请检查！"

    #修改节点的执行状态，并需要基于websocket推送到前端显示 同步线程调用
    def update_node_work_status(self,node,work_status):
        #更新状态
        bchange = node.update_work_status(work_status)
        #基于websocket推送到前端
        if bchange:
            #判断是否是web端最新获取数据的task
            if self.taskM.web_cur_task == self.task_id:
                idatatype = 1
                strdata = {"node_path":node.path,"node_workstatus":work_status}
                asyncio.run(g_WSM.send_data(idatatype,strdata))

    #获取本次的提交提示词
    def get_llm_prompt(self,llm_type,str_res,user_Prompt):
        if llm_type == 0:
            ext_Prompt = f'''
初始信息：{str_res}
任务：请开始对该目标的渗透测试工作。
'''
        elif llm_type == 1:  # 提交指令执行结果 --- 正常提交
            # 构造本次提交的prompt
            ext_Prompt = f'''
上一步结果：{str_res}
任务：生成下一步渗透测试指令或判断是否完成该节点测试。
'''
        elif llm_type == 2:  # llm返回的指令存在问题，需要再次请求返回
            ext_Prompt = f'''
反馈类型：节点指令格式错误
错误信息：{str_res}
任务：请按格式要求重新生成该节点上一次返回中生成的所有指令。
'''
        elif llm_type ==3:  #对节点没有指令的，请求指令
            ext_Prompt = f'''
任务：{str_res}。
'''
        elif llm_type == 5:
            ext_Prompt = f'''
反馈类型：测试指令格式错误
错误信息：{str_res}
任务：请根据格式要求，重新生成该测试指令。
'''
        else:
            self.logger.debug("意外的类型参数")
            return ""

        user_Prompt = user_Prompt + ext_Prompt
        return user_Prompt

    #添加子节点
    def add_children_node(self,parent_node,children_names,cur_node,status="未完成"):
        for child_name in children_names:
            bfind = False
            for node_child in parent_node.children:
                if node_child.name == child_name:
                    bfind = True    #有重复的了
                    break
            if not bfind:
                # 添加节点
                new_node = TreeNode(child_name, parent_node.task_id, status)
                parent_node.add_child(new_node,cur_node.messages)  # message的传递待验证

    #处理节点指令
    def tree_manager(self,node_cmds,node,commands,DBM):
        '''更新渗透测试树
        node_cmds是json-list
        2025-03-22添加commands参数，用于处理LLM对同一个节点返回了测试指令，但还返回了no_instruction节点指令
        '''
        if not node_cmds:    # or len(node_cmds)==0: 正常not判断就可以有没有节点指令
            return True,commands,0

        #对节点指令进行校验
        bok,strerror = g_CV.verify_node_cmds(node_cmds)
        if not bok: #节点指令存在问题，则不进行后续处理，提交一个错误反馈任务
            # 提交llm待处理任务
            self.put_node_reslist(node, strerror, 2)
            return False,commands,0

        residue_node_cmds = []
        no_instr_nodes = []
        #如果有on_instruction，先补全指令保障信息链的完整
        for node_cmd in node_cmds:
            action = node_cmd["action"]
            if action == "no_instruction":
                node_names = node_cmd["nodes"].split(',')
                for node_name in node_names:
                    # 先判断是否在测试指令中，若在则不提交llm任务，只能接受在一次返回中同一节点有多条测试指令，不允许分次返回
                    bcommand = False
                    for com in commands:
                        if node_name in com:
                            bcommand = True
                            break
                    if bcommand:  # 如果存在测试指令，则不把该节点放入补充信息llm任务---尝试不对比是否有返回指令，DS会一直返回指令，还返回on_instruction
                        continue
                    no_instr_nodes.append(node_name)
            else:#剩余的节点指令
                residue_node_cmds.append(node_cmd)

        if no_instr_nodes:  # 阻塞式，在当前节点提交补充信息，完善节点指令 -- 优势是省token
            new_commands = self.get_other_instruction(no_instr_nodes, DBM, node)
            commands.extend(new_commands)

        # #对节点数据进行初步验证
        # ad_instr_nodes, no_add_nodes = g_CV.verify_node_data(node_cmds)
        # if no_add_nodes:#如果有没有添加的节点，默认在当前节点下添加 -- 一般不会有，还没遇到
        #     self.add_children_node(node,no_add_nodes,node)
        # #ad_instr_nodes --- 还没处理

        #先执行add_node操作---2025-4-12-调整：message取当前节点，节点允许为子节点添加子节点
        iadd_node = 0
        residue_cmd_sno_add_= []
        for node_json in residue_node_cmds:
            action = node_json["action"]
            if action == "add_node":  # 新增节点
                parent_node_name = node_json["parent"]
                status = "未完成"  #2025-4-11修改MGS-节点指令格式，取消了status
                node_names = node_json["nodes"].split(',')
                # 新增节点原则上应该都是当前节点增加子节点
                if node.name == parent_node_name or parent_node_name.endswith(node.name): #2233ai，节点名称字段会返回整个路径
                    #添加当前节点的子节点 -- 这是标准情况
                    self.add_children_node(node, node_names,node)
                    iadd_node += len(node_names)  # 添加节点的数量---当前只记录给当前节点添加的子节点的数量
                elif node.parent.name == parent_node_name or parent_node_name.endswith(node.parent.name):#添加当前节点的平级节点
                    #是添加当前节点的平级节点（当前节点的父节点下添加子节点） --使用2233ai-o3时遇到的情况
                    self.add_children_node(node.parent,node_names,node)
                else:
                    badd = False
                    for child_node in node.children:#给子节点添加子节点
                        if parent_node_name == child_node.name or parent_node_name.endswith(child_node.name):
                            badd = True
                            self.add_children_node(child_node,node_names,node)
                            break
                    if not badd:
                        self.logger.error(f"添加子节点时，遇到父节点名称没有找到的，需要介入！！{node_json}")  # 丢弃该节点
            else:#未处理的节点指令添加到list
                residue_cmd_sno_add_.append(node_json)

        if iadd_node and self.taskM.web_cur_task == self.task_id: #如果新增了节点，且该节点树是当前查看的数据，需要通知前端更新数据
            idatatype = 2
            strdata = "update accack_tree!"
            asyncio.run(g_WSM.send_data(idatatype, strdata))
            #先取消当前task,已经通知前端重新获取,这样可以避免后端不必要的数据推送
            self.taskM.web_cur_task = 0

        #执行剩余的节点指令--不分先后
        for node_json in residue_cmd_sno_add_:#2025-4-11重新调整了节点指令格式定义
            action = node_json["action"]
            if action == "find_vul":
                node_name = node_json["node"]
                vul_node = None
                if node.name == node_name or node_name.endswith(node.name): #正常应该是当前节点漏洞信息--暂时只考虑只会有一个漏洞
                    vul_node = node
                else: #匹配子节点
                    for child in node.children:
                        if child.name == node_name or node_name.endswith(child.name):
                            vul_node = node
                            break
                if vul_node:    #找到对应了漏洞节点
                    try:
                        vul_node.vul_type = node_json["vulnerability"]["name"]
                        vul_node.vul_grade = node_json["vulnerability"]["risk"]
                        vul_node.vul_info = node_json["vulnerability"]["info"]
                        #保存到数据库
                        DBM.insert_taks_vul(self.task_id,vul_node.name,vul_node.path,vul_node.vul_type,vul_node.vul_grade,
                                            vul_node.vul_info)
                    except:
                        self.logger.error("漏洞信息错误")
                        continue
                else:
                    str_user = f"遇到不是修改本节点状态的，需要介入！！{node_json}"
                    self.logger.error(str_user)
            elif action == "end_work":
                node_name = node_json["node"]
                if node.name == node_name or node_name.endswith(node_name):  # 正常应该是当前节点
                    node.status = "已完成"
                else:
                    str_user = f"遇到不是修改本节点状态的，需要介入！！{node_json}"
                    self.logger.error(str_user)
            elif action == "no_create": #提交人工确认
                nodes = node_json["nodes"]
                if nodes:
                    str_add = {"未新增的节点": nodes}
                    self.logger.debug(str_add)
                    # 提交一个继续反馈任务--继续后续工作 2025-3-25不自动处理
                    # self.put_one_llm_work(str_add, node, 4)
                    # self.logger.debug(f"未新增的节点有：{nodes}")
            else:
                self.logger.error("****不应该执行到这！程序逻辑存在问题！")
        return True,commands,iadd_node

    #阻塞轮询补充指令
    def get_other_instruction(self,nodes,DBM,cur_node):
        res_str = ','.join(nodes)
        new_commands = []
        while res_str:
            self.logger.debug(f"开始针对f{res_str}这些节点请求测试指令")
            user_Prompt = f'''
当前分支路径：{cur_node.path}
当前节点信息：
- 节点名称：{cur_node.name}
- 节点状态：{cur_node.status}
- 漏洞类型：{cur_node.vul_type}
反馈类型：需要补充以下子节点的测试指令:{res_str}
任务：
1.请生成这些子节点的测试指令,注意不要生成重复的测试指令；
2.这些节点的父节点为当前节点，请正确生成这些节点的节点路径；
3.只有当还有节点未能生成测试指令或不完整时，才返回未生成指令的节点列表。
        '''
            res_str = ""
            #node_cmds, commands = self.LLM.get_llm_instruction(user_Prompt, DBM, cur_node)  # message要更新
            node_cmds, commands, reasoning_content, content, post_time = self.LLM.get_llm_instruction(user_Prompt,
                                                                                                      cur_node)  # message要更新
            # LLM记录存数据库
            cur_node.llm_sn += 1
            bres = DBM.insert_llm(self.task_id, user_Prompt, reasoning_content, content, post_time, cur_node.llm_sn,cur_node.path)
            if not bres:
                self.logger.error(f"{cur_node.name}-llm入库失败！")
            #把返回的测试指令进行追加
            new_commands.extend(commands)
            #判断是否还有未添加指令的节点
            for node_json in node_cmds: #正常应该只有一条no_instruction --暂时只处理
                if "no_instruction" in node_json and "nodes" in node_json:
                    tmp_nodes = []
                    node_names = node_json["nodes"].split(',')
                    for node_name in node_names:
                        if node_name in nodes:
                            tmp_nodes.append(node_name)
                    res_str = ','.join(tmp_nodes)
                    break
        self.logger.debug("未添加指令的节点，都已完成指令的添加！")
        return new_commands

    def start_task(self,task_id,task_status=1,attack_tree=None):
        self.task_id = task_id
        '''
        启动该测试任务
        '''
        #判断目标合法性
        # bok,target,type = self.TargetM.validate_and_extract(self.target)    #是否还需要判断待定?
        # if not bok:
        #     return False, "{target}检测目标不合规，请检查！"
        #初始化节点树
        if attack_tree:#有值的情况是load
            self.attack_tree =attack_tree
            #加载未完成的任务
            if self.work_type ==1:#自动模式
                #提交到mq,待线程执行
                self.put_work_node()
        else: #无值的情况是new_create
            root_node = TreeNode(self.target, self.task_id)  #根节点
            self.attack_tree = AttackTree(root_node)    #创建测试树，同时更新根节点相关内容
            self.LLM.build_initial_prompt(root_node)    #对根节点初始化system-msg
            #插入一个user消息
            # 提交第一个llm任务，开始工作
            know_info = f"本测试主机的IP地址为:{self.local_ip}"
            if self.cookie:
                know_info = know_info + f",本站点的cookie值为{self.cookie}"
            self.put_node_reslist(root_node,know_info,0)    #入待提交list
            #初始保存个attack_tree文件
            g_PKM.WriteData(self.attack_tree,str(self.task_id))
        #启动工作线程
        self.task_status = task_status
        self.brun = True #线程正常启动
        #启动指令工作线程
        for i in range(self.max_thread_num):
            w_th = threading.Thread(target=self.do_worker_th,args=(i,))
            w_th.start()
            self.workth_list[i] = w_th
        #启动llm提交线程--llm暂时单线程，多线程处理时attack_tree需要加锁
        for j in range(self.llm_max_nums):
            l_th = threading.Thread(target=self.th_llm_worker,args=(j,))
            l_th.start()
            self.llmth_list[j]=l_th
        #启动自检线程
        self.check_th = threading.Thread(target=self.th_check)
        self.check_th.start()

    def stop_task(self): #还未处理
        self.brun = False
        self.InstrM.init_data()
        #结束任务需要收尾处理#？


if __name__ == "__main__":
    pass