zfsafe/mycode/TaskObject.py

'''
渗透测试任务管理类  一次任务的闭合性要检查2025-3-10  一次任务后要清理LLM和InstrM的数据
'''
#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.DataFilterManager import DataFilterManager
from myutils.ReadWriteLock import ReadWriteLock
import asyncio
import queue
import time
import os
import re
import threading
import json
import textwrap

class TaskObject:

    def __init__(self,test_target,cookie_info,work_type,llm_type,num_threads,local_ip,fake_target,taskM,target_id=0,safe_rank=0):
        #功能类相关
        self.taskM = taskM
        self.logger = LogHandler().get_logger("TaskObject")
        self.InstrM = g_instrM  # 类对象渗透，要约束只读取信息，且不允许有类全局对象--持续检查
        self.PythonM = taskM.PythonM    #python进程池往上提一层
        self.DataFilter = DataFilterManager(test_target,fake_target)
        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.max_layer = myCongif.get_data("max_node_layer")
        self.sleep_time = myCongif.get_data("sleep_time")
        self.target = test_target
        self.target_id = target_id  #巡检任务有值--对应巡检目标
        self.cookie = cookie_info
        self.work_type = work_type  #工作模式 0-人工，1-自动
        self.task_id = None
        self.task_status = 0         #0-初始状态，1-执行中，2-暂停中，3-已完成，4-已结束
        self.local_ip = local_ip
        self.attack_tree = None     #任务节点树
        self.attack_tree_lock = None    #把节点的Lock外置
        self.safe_rank = safe_rank          #安全级别 0-9   #？暂时还没实现更新逻辑
        self.is_had_work = False
        self.is_had_work_lock = threading.Lock()
        self.cur_stage = 0      #0--信息收集阶段，1--渗透测试阶段

        #读写锁
        self.rwlock = ReadWriteLock()
        #指令执行相关-------
        self.max_thread_num = num_threads   #指令执行线程数量
        self.workth_list = [None] * num_threads   #线程句柄list
        self.workth_node_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.llmth_node_list = [None] * self.llm_max_nums  # llm线程--node
        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 get_instr_node(self,index):
        self.rwlock.acquire_read()
        try:
            self.workth_node_list[index] = self.instr_node_queue.get(block=False)
        except queue.Empty:
            self.workth_node_list[index] = None
        finally:
            self.rwlock.release_read()

    def put_instr_node(self,node):
        #self.instr_node_queue.put(node)
        self.rwlock.acquire_read()
        try:
            self.instr_node_queue.put(node)
        finally:
            self.rwlock.release_read()

    #读取待办llm节点
    def get_llm_node(self,index):
        self.rwlock.acquire_read()
        try:
            self.llmth_node_list[index] = self.llm_node_queue.get(block=False)
        except queue.Empty:
            self.llmth_node_list[index] = None
        finally:
            self.rwlock.release_read()

    def put_llm_node(self,node):
        # self.llm_node_queue.put(node)
        self.rwlock.acquire_read()
        try:
            self.llm_node_queue.put(node)
        finally:
            self.rwlock.release_read()


    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 smart_truncate(self,s: str,max_length: int = 5000,ellipsis: str = "...") -> str:
        """截取字符串前 max_length 个字符，并添加截断标记（确保总长度不超过限制）"""
        if not s:
            return s
        #一些无效信息的删除
        s = s.replace("pydev debugger: bytes arguments were passed to a new process creation function. Breakpoints may not work correctly.","")

        # 计算保留长度（考虑截断标记占位）
        truncated = s[:max_length - len(ellipsis)] if len(s) > max_length else s
        return truncated + ellipsis if len(s) > max_length else s

    def do_instruction(self,instruction):
        instruction = textwrap.dedent(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()  # 指令执行开始时间

        # 2025-4-27要重新定义bool值作用，初步想法True-作为指令已处理（执行或不执行），False-作为异常，指令还没有处理-可以回Q

        if instruction.startswith("python-code"):  # python代码--超过子进程数会阻塞等待，但不开始超时记时
            bsuccess, instr, reslut, source_result, ext_params = self.PythonM.execute_instruction(instruction)
        else:  # shell指令
            bsuccess, instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction)

        end_time = get_local_timestr()  # 指令执行结束时间
        # 只取结果的5000长度
        reslut = self.smart_truncate(reslut, 4000)
        source_result = self.smart_truncate(source_result)
        return start_time,end_time,bsuccess,instr,reslut,source_result,ext_params

    def do_worker_th(self,index):
        #线程的dbm需要一个线程一个
        th_DBM = DBManager()
        th_DBM.connect()
        th_index = index
        bnode_work = False
        while self.brun:
            if self.task_status == 1:
                self.get_instr_node(th_index)   #获取一个待办节点
                if self.workth_node_list[th_index]:
                    work_node = self.workth_node_list[th_index]
                    llm_type = 1
                    # 开始执行指令
                    bnode_work = True
                    results = []
                    while True: #遍历节点待执行指令，执行
                        instruction = work_node.get_instr()
                        if not instruction:
                            break
                        self.doing_instr_list[th_index] = instruction

                        start_time, end_time, bsuccess, instr, reslut, source_result, ext_params = self.do_instruction(instruction)

                        # 入数据库 -- 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("数据库连接失败！！")
                        # 暂存结果
                        fake_inst,fake_resul = self.DataFilter.filter_result(instr,reslut)
                        oneres = {'执行指令': fake_inst, '结果': fake_resul}
                        results.append(oneres)      #结果入队列后，指令就不能回退

                    #节点执行完成后置空
                    self.doing_instr_list[th_index] = ""
                    #指令都执行结束后，入节点待提交队列
                    str_res = json.dumps(results, ensure_ascii=False)
                    # 提交llm待处理任务 --更新节点work_status
                    self.put_node_reslist(work_node, str_res, llm_type)
                    # 保存记录
                    g_PKM.WriteData(self.attack_tree,str(self.task_id))
                else:
                    if bnode_work:
                        bnode_work = False
                        self.no_work_to_do()    #判断是否需要把当前任务的无工作状态推送到前端
                    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
        bnode_work = False
        max_llm_sn = myCongif.get_data("max_llm_sn")
        while self.brun:
            if self.task_status == 1:
                self.get_llm_node(th_index)
                if self.llmth_node_list[th_index]:
                    llm_node = self.llmth_node_list[th_index]
                    #开始处理
                    bnode_work = True
                    # {llm_node.status} --- 暂时固化为未完成
                    user_Prompt = f'''
当前分支路径：{llm_node.path}
当前节点信息：
- 节点名称：{llm_node.name}
- 节点状态：未完成
- 漏洞类型：{llm_node.vul_type}
'''
                    ido_count = 0  # 用于控制线程重复执行的上限
                    tmp_commands = []
                    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"]
                        #判断执行次数
                        if llm_type !=2:
                            llm_node.llm_sn += 1
                            if llm_node.llm_sn >= max_llm_sn + 1:   #提交次数达到上限后，提示LLM结束该节点任务
                                llm_type = 10
                                #该节点的剩余任务不执行,若有--暂定
                                llm_node.clear_res()
                        #获取提示词
                        prompt = self.get_llm_prompt(llm_type,str_res,user_Prompt)
                        self.doing_llm_list[th_index] = prompt  #执行更新
                        # 提交llm请求返回数据--并对返回数据进行处理，节点指令直接执行，测试指令根据工作模式执行
                        bsuccess, node_cmds, commands, strerror,reasoning_content,content =  self.node_get_llm_instruction(llm_node,prompt) #如果API调用出问题，会一直循环尝试
                        if not bsuccess:#返回的内容不符合格式约定 --重试3次还不符合--丢弃
                            continue    #下一个待办事项
                        #返回结果符合要求
                        # LLM记录存数据库 ---若有指令错误的情况，中间的请求修正的MSG会丢失数据库记录
                        if th_DBM.ok:
                            post_time = get_local_timestr()
                            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("数据库连接失败！")
                        # 更新tree
                        bok, new_commands,iadd_node = self.tree_manager(node_cmds, llm_node, commands, th_DBM)
                        # 分析指令入对应节点
                        if bok:  # 节点指令若存在错误，测试指令都不处理，需要LLM重新生成
                            tmp_commands.extend(new_commands)
                    #节点的待提交任务都完成后，统一处理指令
                    self.put_node_instrlist(tmp_commands, llm_node)
                    #一个节点完成，节点树持久化---待验证是否有局部更新持久化的方案
                    g_PKM.WriteData(self.attack_tree,str(self.task_id))
                    #推送前端刷新数据--执行一个节点就刷新一次
                    if 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
                    # 一个节点执行完成后--提交中llm置空
                    self.doing_llm_list[th_index] = ""
                else:
                    if bnode_work:
                        bnode_work = False
                        self.no_work_to_do()  # 判断是否需要把当前任务的无工作状态推送到前端
                    time.sleep(self.sleep_time)
            else:
                time.sleep(self.sleep_time)

    def had_work_to_do(self):   #任务单步状态控制 -- 置工作中 --约束：只有任务单步会调用
        bsuccess = False
        with self.is_had_work_lock:
            if not self.is_had_work:
                bsuccess = True  # 这不return 是怕不释放锁
                self.is_had_work = True
        return bsuccess

    def is_work_doing(self):
        '''检查是否还有在执行的事项'''
        #获取队列快照
        self.rwlock.acquire_write()
        try:
            instr_node_list = list(self.instr_node_queue.queue)  # 待执行指令的node
            llm_node_list = list(self.llm_node_queue.queue)  # 待提交llm的node
            wth_n_list = self.workth_node_list.copy()
            lth_n_list = self.llmth_node_list.copy()
        finally:
            self.rwlock.release_write()
        #只针对这个快照数据进行判断
        #先判断正在执行节点是否有值 ---
        for do_node in wth_n_list:
            if do_node:
                return True
        for do_node in lth_n_list:
            if do_node:
                return  True
        #再判断待办节点
        if len(instr_node_list) == 0 and len(llm_node_list) == 0:
            return False
        return True

    #是否需要线程锁-待定
    def no_work_to_do(self):    #任务单步状态控制-- 非工作中--2025-5-7增加了轮次，需要全面验证执行逻辑的完整和正确性
        #是否还有在执行事项
        bworking  = self.is_work_doing()
        if not  bworking:
            # 没有在执行任务了
            #推送是否有工作任务的状态到前端,
            with self.is_had_work_lock:
                if self.is_had_work:    #如果已经是False那就不需要修改了
                    #推送到前端
                    if self.task_id == self.taskM.web_cur_task:
                        #把是否有任务在执行的状态推送到前端
                        idatatype = 3
                        strdata = "单步任务执行完成！"
                        asyncio.run(g_WSM.send_data(idatatype, strdata))
                    self.is_had_work = False

    #自检线程 --1.输出执行状态。2.需要自检和修复
    def th_check(self):
        icount = 0
        while self.brun:
            try:
                bworking = False
                cur_time = get_local_timestr()
                print(f"----------{self.task_id}-当前时间程序运行情况:{cur_time}")
                #执行中instr-node
                index = 0
                for w_th in self.workth_list:
                    if not w_th.is_alive():#线程
                        print(f"Work线程-{index}已处于异常状态，需要重新创建一个工作线程")
                    else:
                        if self.doing_instr_list[index]:
                            bworking =True
                            print(f"Work线程-{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:
                        if self.doing_llm_list[index]:
                            bworking = True
                            print(f"LLM线程-{index}-在执行指令：{self.doing_llm_list[index]}")
                    index += 1

                #判断任务是否完成
                bover = self.is_over()
                if bover and not bworking:
                    self.over_task()    #自己停自己--确保没有阻塞操作
                    # 通知TM 开始下一个任务
                    self.taskM.start_next_task(self.task_id,self.target_id)
                    break   #退出循环--结束线程
                #处理点修复操作
                icount +=1
                if icount == 5:
                    pass

                #休眠60
                time.sleep(60)
            except Exception as e:
                print(f"*********自检线程异常退出:{str(e)}")
                break

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

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

    async def put_instr_mq_async(self,node):
        #这里不做状态的判断，调用前处理
        self.put_instr_node(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.put_llm_node(node)
        await self.update_node_work_status_async(node,4)  #提交中

    #修改节点的执行状态，并需要基于websocket推送到前端显示 同步线程调用
    def update_node_work_status(self,node,work_status):
        old_work_status = node.get_work_status()
        #更新状态
        bchange = node.update_work_status(work_status)  #1,3会返回Flase
        #基于websocket推送到前端
        if work_status != 1 and old_work_status != 0:    #llm执行完成后会发送单独的指令更新树，所以不发送1更新节点了
            #判断是否是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))

    #web端调用
    async def update_node_work_status_async(self,node,work_status):
        #更新状态
        bchange = node.update_work_status(work_status)
        #基于websocket推送到前端
        if work_status != 1:
            #判断是否是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 node.bwork:  #节点是工作状态
            if self.work_type == 1: #自动模式
                self.put_llm_mq(node)   #变4
            elif self.work_type == 0 and node.step_num > 0: #人工模式，但是单步步次还没有执行完成
                node.step_num -= 1
                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 replace_error_instr(self,command):
        content = command["content"]
        content = content.replace("| |","||")
        content = content.replace("& &","&&")
        command["content"] = content
        return command

    def put_node_instrlist(self, commands, node):       #如果当前节点没有进一般指令返回，需要修改节点执行状态
        if not node:
            return
        '''
        - dash指令：{\"action\":\"dash\",\"path\":\"节点路径\",\"content\":\"指令内容\"}
        - python指令：{\"action\":\"python\",\"path\":\"节点路径\",\"content\":\"指令内容\"}
        '''
        node_list = []  #有待办指令的节点
        for command in commands:
            command = self.replace_error_instr(command)
            # 使用正则匹配方括号中的node_path（非贪婪模式）
            #match = re.search(r'\[(.*?)\]', command)
            com_type = command["action"]
            com_path = command["path"]
            content = command["content"]
            node_name = com_path.split("->")[-1]
            if com_type == "python":    #对于python代码，加个python头
                content = "python-code " + content
            #寻找指令的对应节点#'''强制约束，不是本节点或者是子节点的指令不处理'''
            find_node = None
            if node_name == node.name:  #指令是当前节点的
                find_node = node
            else:
                if self.cur_stage == 0: #信息收集阶段，不是当前节点的指令，也作为当前节点
                    find_node = node
                else:
                    for child_node in node.children:    #暂时只找一层
                        if node_name == child_node.name:
                            if child_node.do_sn == 0:   #只有没执行过指令的子节点，才允许添加指令 2025-5-9 新增限制 避免父节点和子节点同时有指令执行，提交llm后，父节点返回子节点指令。
                                find_node = child_node
                            break
            # find_node = self.attack_tree.find_node_by_nodepath_parent(node_path,node,iadd_node,commands)
            # if not find_node:#对于没有基于节点路径找到对应节点--增加通过节点名称匹配的机制 2025-4-13日添加
            #     find_node = self.find_node_by_child_node_name(node, node_name)  # 递归找子节点

            if find_node:
                find_node.add_instr(content,node.parent_messages,node.cur_messages) #2025-4-23调整为第一添加指令时传递Msg
                #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"没有找到指令对应的节点：{com_path},当前节点{node.path}")#丢弃该指令

                #这里对于丢弃指令，有几种方案：
                # 1.直接丢弃不处理，但需要考虑会不会产生节点缺失指令的问题，需要有机制验证节点；------ 需要有个独立线程，节点要加锁--首选待改进方案
                # 2.入当前节点的res_queue,但有可能当前节点没有其他指令，不会触发提交，另外就算提交，会不会产生预设范围外的返回，不确定；
                # 3.独立队列处理

        #判断当前节点是否有指令
        if node not in node_list:
            #修改该节点状态为0--无待执行任务
            self.update_node_work_status(node,0)
        #判断是否入instr待办队列---放循环外面是等指令都添加完成后提交待办
        for node in node_list:
            if node.bwork:
                if self.work_type == 1: #是自动执行模式
                    self.put_instr_mq(node)     #2-执行中
                elif self.work_type == 0 and node.step_num > 0:  #人工模式 且剩余步次大于0
                    node.step_num -= 1
                    self.put_instr_mq(node)     #2-执行中

    def put_work_node(self,work_type):
        '''遍历节点需要处理的任务，提交mq,load_task-触发--线程安全
        work_type 0-人工，1-自动
        '''
        instr_status = None
        llm_status = None
        if work_type == 0:  #人工
            instr_status = (2,)
            llm_status = (4,)
        else:   #自动
            instr_status = (1,2)
            llm_status = (3,4)

        nodes = self.attack_tree.traverse_bfs()
        for node in nodes:
            if not node.bwork:
                continue
            node_work_status = node.get_work_status()
            if node_work_status in instr_status:   #有待执行指令
                if node.is_instr_empty():#
                    node.update_work_status(-1) #置0 -1作为额外的条件参数
                else:
                    self.put_instr_node(node)   #instr_status都提交执行
                    node.update_work_status(-2)# 置2 --执行中
                    #llm-list不处理，正常应该为空
            elif node_work_status in llm_status:
                if node.is_llm_empty():#数据有问题，放弃掉
                    node.update_work_status(-1)
                else:
                    self.put_llm_node(node)
                    node.update_work_status(-3) #置4
            else:
                pass

    #web端提交单步任务--节点单步--start
    async def put_one_node(self,node,step_num=1):
        #提交某个节点的代表任务
        if self.task_status ==1 and self.work_type==0 and node.bwork:
            iwork_status = node.get_work_status()
            if iwork_status in (1,3):
                node.step_num = step_num - 1   #单步步次赋值 -1 规则提交时-1，执行结束连续判断再提交
                if iwork_status == 1:
                    await self.put_instr_mq_async(node)
                else:
                    await self.put_llm_mq_async(node)
                return True,"已提交单步任务"
            else:
                error = ""
                if iwork_status == 0:
                    error = "该节点没有待办任务"
                elif iwork_status in (2,4):
                    error = "该节点已经在执行事项中"
                else:
                    self.logger.error("noed的工作状态值存在问题，需要人工介入！！")
                return False,error
        else:
            return False,"当前的任务或节点状态不允许执行单步，请检查！"

    #web端提交任务单步--任务单步---2025-5-8-增加约束，只允许单步启动时调用
    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:
                    bput,_ = await self.put_one_node(node,step_num) #错误信息有丢失
                    if bput:
                        b_putwork = True
                if b_putwork:
                    return True,"已提交单步任务"
                else:   #所有节点都没有提交任务
                    self.over_task()
                    return False,"该任务已经没有待提交任务"
            else:
                return False,"当前任务正在执行任务中，不需要提交单步任务！"
        else:
            return False,"当前的任务状态不允许执行单步，请检查！"

    #获取本次的提交提示词
    def get_llm_prompt(self,llm_type,str_res,user_Prompt):
        if llm_type == -2:  #调用llm api失败的情况，prompt已是组装好
            return str_res
        elif llm_type == -1:
            ext_Prompt = "请针对当前节点收集渗透测试所需的相关信息"
        elif 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}
任务：请根据格式要求，重新生成该测试指令。
'''
        elif llm_type == 10:
            max_do_sn = myCongif.get_data("max_do_sn")
            ext_Prompt = f'''
上一步结果：{str_res}
任务：当前节点已执行超过{max_do_sn}次测试指令，若无新发现请结束该节点的测试，若有新发现请生成子节点，在子节点推进下一步测试。
'''
        else:
            self.logger.debug("意外的类型参数")
            return ""

        user_Prompt = user_Prompt + ext_Prompt
        return user_Prompt

    #添加子节点
    def add_children_node(self,parent_node,children_names,cur_message=None,status="未完成"):
        existing_names = {node.name for node in parent_node.children}  # 现有子节点名称集合
        unique_names = list(set(children_names))  # 去重
        layer_num = parent_node.cur_layer + 1
        for child_name in unique_names:
            if child_name not in existing_names:
                # 添加节点
                new_node = TreeNode(child_name, parent_node.task_id,layer_num,status)
                parent_node.add_child(new_node)
                #existing_names.add(child_name)  # 更新集合 -- 已经去重过了，不需要在添加到比对

    def dedupe_ports(self,IPS):
        new_ips = []
        for ip_entry in IPS:
            seen_ports= set()
            filtered_ports = []
            for port_info in ip_entry.get("Ports", []):
                port = port_info.get("Port")
                if port not in seen_ports:
                    seen_ports.add(port)
                    filtered_ports.append(port_info)
            # 构造新的 IP 条目，只替换 Ports 字段
            new_entry = ip_entry.copy()
            new_entry["Ports"] = filtered_ports
            new_ips.append(new_entry)
        return new_ips

    #更新资产数据库
    def add_update_assets(self,URL,IPS,DBM):
        url_id = 0
        # URL信息入库或更新
        if URL:  # required_keys = {"Domain", "Subdomains", "Reqistrar", "Email", "Creation_date", "Expiration_date"}
            Domain = URL["Domain"]
            Subdomains = URL["Subdomains"]
            Registrant = URL["Registrant"]
            Registrar = URL["Registrar"]
            Email = URL["Email"]
            Creation_date = URL["Creation_date"]
            Expiration_date = URL["Expiration_date"]
            url_id = DBM.add_or_update_URL_asset(Domain, Subdomains, Registrant, Email, Creation_date,
                                                 Expiration_date,Registrar)
        # URL-关联的-IP信息入库或更新
        to_url_ip=[]
        for IP_data in IPS:
            #[\"IP\":\"192.168.1.100\",\"IPtype\":\"IPv4/IPv6\",\"Ports\":[端口信息json]]
            ip_address = IP_data["IP"]
            IPtype = IP_data["IPtype"]
            Ports = IP_data["Ports"]
            #IP入库或更新
            ip_id,scan_count = DBM.add_or_update_IP_asset(ip_address,IPtype)
            if url_id:#有url就需要入库关联表 --统一处理
                to_url_ip.append(ip_id)
            #把IP资产和task进行关联
            DBM.add_task_to_ip(self.task_id,ip_id)
            #Ports端口数据入库过更新
            DBM.update_port(ip_id,scan_count,Ports)
        #统一维护url-ip的对应关系
        if to_url_ip:
            DBM.update_url_to_ip(url_id,to_url_ip)

    def update_attack_tree(self,URL,IPS,cur_node):
        #self.add_children_node(node, add_node_names)
        # if URL: #说明检测目标是url 正常情况只有从url->IP,无法从ip->url
        #     for IP_data in IPS: #IPS不应该会重复
        #         #增加IP子节点
        #         new_node = TreeNode(IP_data["IP"], cur_node.task_id, 0, "已完成")  #IP不管如何算0层
        #         cur_node.add_child(new_node)
        # else:#说明目标是IP地址
        #     for IP_data in IPS:#IP应该只有一个
        #         pass
        for IP_data in IPS:
            ip_node = cur_node  #如果没有URL，IP就是当前根节点
            if URL:
                # 增加IP子节点
                ip_node = TreeNode(IP_data["IP"], cur_node.task_id, 0, "已完成")  # IP不管如何算0层
                cur_node.add_child(ip_node)
            #port端口节点
            ports = IP_data["Ports"]
            for port_data in ports:
                port = port_data["Port"]
                service = port_data["Service"]
                version = port_data["Version"]
                protocol = port_data["Protocol"]    #TCP/UDP
                status = port_data["Status"]    #open/closed/filtered
                #添加port节点
                port_node = TreeNode(str(port), ip_node.task_id, 1, "未完成")  # port为1层
                # parent_messages 初始化
                self.LLM.build_init_attact_prompt(port_node)
                ip_node.add_child(port_node)
                #添加已知信息
                know_info = f"端口号：{port},服务：{service},版本：{version},协议：{protocol},端口状态：{status}"
                self.put_node_reslist(port_node, know_info, 0)  #llm_type -0 info->accack的过渡，需要重新构建prompt
        #保存节点树--th_llm 节点任务执行完后会报错

    #提交prompt 拉取llm返回结果，再做一层封装 --- 信息的脱敏和还原都放这里面
    def node_get_llm_instruction(self,cur_node,prompt):
        i_instr_error = 0
        while True:
            fake_prompt = self.DataFilter.filter_prompt(prompt)  # 目标脱敏
            # 构造与更新message
            message = {"role": "user", "content": fake_prompt}
            cur_node.cur_messages.append(message)  # 更新节点message
            sendmessage = []
            sendmessage.extend(cur_node.parent_messages)
            sendmessage.extend(cur_node.cur_messages)
            ido_count = 0
            while True: #
                iresult,reasoning_content,content,error = self.LLM.get_llm_instruction(sendmessage)
                if iresult == -1:   #调用LLMapi 出错
                    if ido_count < 3:
                        ido_count += 1
                        self.logger.error(f"模型调用出错-{error}，休眠10秒后重试")
                        time.sleep(10)
                    else:  # 调用llmapi 已达错误上限--休眠半个小时候再试
                        ido_count = 0
                        self.logger.error(f"模型调用连续出错-{error}，休眠30分钟后重试")
                        time.sleep(60 * 30)
                    continue
                #llm调用成功后，记录llm历史信息
                cur_node.cur_messages.append({'role': 'assistant', 'content': content})
                print(content)
                break

            # 目标还原
            real_con = self.DataFilter.filter_instruction(content)

            # 解析验证返回内容
            bsuccess, node_cmds, commands, strerror = self.LLM.fetch_instruction(real_con)
            if not bsuccess:  # 返回的内容不符合格式约定--重新提交
                if i_instr_error >=3:
                    return bsuccess, node_cmds, commands, strerror,reasoning_content,real_con
                i_instr_error += 1
                user_Prompt = f'''
当前分支路径：{cur_node.path}
当前节点信息：
- 节点名称：{cur_node.name}
- 节点状态：未完成
- 漏洞类型：{cur_node.vul_type}
'''
                prompt = self.get_llm_prompt(2,strerror,user_Prompt)
                # one_llm = {'llm_type': 2, 'result': strerror}
                # cur_node.add_res(one_llm, 1)  # 入节点结果队列的第一个
                continue
            else:
                return bsuccess, node_cmds, commands, strerror,reasoning_content,real_con

    #阻塞轮询补充指令
    def get_other_instruction(self,nodes,DBM,cur_node):
        new_commands = []
        res_str = ','.join(nodes)
        no_instr_nodes = nodes
        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.这些节点的父节点为当前节点，请正确生成这些节点的节点路径；
        '''
            bsuccess, node_cmds, commands, strerror,reasoning_content,content = self.node_get_llm_instruction(cur_node, user_Prompt)
            #正常不应该会有node_cmds
            if not bsuccess:    #失败就结束
                break
            res_str = ""
            # LLM记录存数据库
            cur_node.llm_sn += 1
            post_time = get_local_timestr()
            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)
            #再验证是否还有缺少的
            tmp_nodes = []
            for no_instr_node in no_instr_nodes:
                bcommand = False
                for com in commands:
                    if no_instr_node in com["path"]:
                        bcommand = True
                        break
                if bcommand:  # 如果存在测试指令，则不把该节点放入补充信息llm任务---尝试不对比是否有返回指令，DS会一直返回指令，还返回on_instruction
                    continue
                # 没有对应指令
                tmp_nodes.append(no_instr_node)
            res_str = ','.join(tmp_nodes)
            no_instr_nodes = tmp_nodes

        self.logger.debug("未添加指令的节点，都已完成指令的添加！")
        return new_commands

    # 处理节点指令
    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

        # message_调整传递时机后，可以先执行添加节点
        # #对节点数据进行初步验证
        # 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 --- 还没处理

        residue_cmd_no_add = []
        all_add_node = []
        add_node_names = []
        for node_json in node_cmds:
            action = node_json["action"]
            if action == "add_node":  # 新增节点
                if node.cur_layer >= self.max_layer:
                    continue  # 节点层级达到上限后不允许再添加子节点-- 平级的一样
                parent_node_name = node_json["parent"]
                # status = "未完成"  #2025-4-11修改MGS-节点指令格式，取消了status
                add_node_names = node_json["nodes"].replace('，', ',').split(',')
                # 新增节点原则上应该都是当前节点增加子节点
                if node.name == parent_node_name or parent_node_name.endswith(node.name):  # 2233ai，节点名称字段会返回整个路径
                    # 添加当前节点的子节点 -- 这是标准情况
                    self.add_children_node(node, add_node_names)
                    all_add_node.extend(add_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, add_node_names)
                    self.logger.debug("遇到一次添加平级节点")
                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, add_node_names)
                            self.logger.debug("遇到一次添加子节点的子节点")
                            break
                    if not badd:
                        self.logger.error(
                            f"添加子节点失败！父节点不是当前节点，不是当前节点的父节点，不是当前节点的子节点，需要介入！！{node_json}---当前节点为：{node.path}")  # 丢弃该节点
            else:  # 未处理的节点指令添加到list
                residue_cmd_no_add.append(node_json)

        no_instr_nodes = []
        # 2025-5-12 是否采用本地校验节点是否有指令，如果使用，则no_instruction就可以不用了
        for add_node in all_add_node:
            bcommand = False
            for com in commands:
                if add_node in com["path"]:  # 这里要不要改成匹配尾部字串?
                    bcommand = True
                    break
            if bcommand:  # 如果存在测试指令，则不把该节点放入补充信息llm任务---尝试不对比是否有返回指令，DS会一直返回指令，还返回on_instruction
                continue
            # 没有对应指令
            no_instr_nodes.append(add_node)
        if no_instr_nodes:  # 阻塞式，在当前节点提交补充信息，完善节点指令 -- 优势是省token
            new_commands = self.get_other_instruction(no_instr_nodes, DBM, node)
            commands.extend(new_commands)

        # 执行剩余的节点指令--不分先后
        for node_json in residue_cmd_no_add:
            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"]
                        # 保存到数据库 --- 数据库有记录多个，tree只保留最新一个
                        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}--当前节点{node.path}"
                    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}--当前节点{node.path}"
                    self.logger.error(str_user)
            elif action == "asset":
                self.cur_stage = 1  # 进入渗透测试阶段
                print(f"{node_json}")
                IPS = self.dedupe_ports(node_json["IPS"])  # ips的端口做去重，ds用到端口重复的情况
                URL = node_json["URL"]
                self.add_update_assets(URL, IPS, DBM)  # 网络资产信息入库或更新
                # 节点树增加节点，并初始化信息采集的数据 ---- 推进到下一阶段
                self.update_attack_tree(URL, IPS, node)
            else:
                self.logger.error("****不应该执行到这！程序逻辑存在问题！")
        return True, commands, len(add_node_names)

    #-----------------任务的启停--------------------
    def init_task(self,task_id,attack_tree = None):
        self.task_id = task_id
        # 初始化节点树
        if attack_tree:  # 有值的情况是load
            self.attack_tree = attack_tree
            # 加载未完成的任务-并恢复工作
            self.put_work_node(self.work_type)

        else:  # 无值的情况是new_create
            root_node = TreeNode(self.target, self.task_id,0)  # 根节点
            self.attack_tree = AttackTree(root_node)  # 创建测试树，同时更新根节点相关内容
            #self.LLM.build_initial_prompt(root_node)  # 对根节点初始化system-msg
            self.LLM.build_init_info_prompt(root_node)  #新建任务先开始信息收集工作 2025-5-15
            know_info = ""
            self.put_node_reslist(root_node,know_info,-1)   #llm_type新增一个-1的状态值

            # 初始保存个attack_tree文件
            g_PKM.WriteData(self.attack_tree, str(self.task_id))

    def is_task_stop(self):
        #检查任务是否处于停止状态--防止停止后，线程还没停，又启动工作线程，造成混乱
        #工作线程
        for work_th in self.workth_list:
            if work_th:
                if work_th.is_alive():
                    self.logger.debug(f"{self.task_id}有存活工作线程")
                    return False
        #llm线程
        for llm_th in self.llmth_list:
            if llm_th:
                if llm_th.is_alive():
                    self.logger.debug(f"{self.task_id}有存活LLM线程")
                    return False
        #自检线程
        if self.check_th:
            if self.check_th.is_alive():
                self.logger.debug(f"{self.task_id}有存活自检线程")
                return False

        return True

    def update_task_work_type(self,new_work_type):
        self.work_type = new_work_type
        #更新数据库
        app_DBM.update_task_work_type(self.task_id,new_work_type)

    def update_task_status(self,new_status):
        self.task_status = new_status
        #更新数据库
        app_DBM.update_task_status(self.task_id,new_status)

    def start_task(self):
        '''
        启动该测试任务，
        :return int str
        '''
        if self.is_task_stop():
            if self.is_over():  #只是判断了待办list，在执行的会漏判断
                return 0, "该任务所有未暂停节点的待办事项都已经结束"
            #更新状态标识
            self.update_task_status(1)
            self.brun = True #线程正常启动
            #启动指令工作线程
            for i in range(self.max_thread_num):
                w_th = threading.Thread(target=self.do_worker_th,args=(i,),name=f"{self.task_id}-w_th-{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,),name=f"{self.task_id}-l_th-{i}")
                l_th.start()
                self.llmth_list[j]=l_th
            #启动自检线程
            self.check_th = threading.Thread(target=self.th_check)
            self.check_th.start()

            return 1,"启动任务成功！"   #就算启动了一部分线程，也要认为是成功
        else:
            return -1,"该任务的工作线程未全面停止，不能重新启动工作，请稍后，若半个小时候还不能启动，请联系技术支持！"

    def stop_task(self): #--暂停task
        if self.brun and self.task_status == 1: #不是正常工作状态就不执行停止工作了
            #停止线程
            self.brun = False
            if self.target_id == 0:
                self.update_task_status(2)  #置状态2--暂停状态
            else:
                self.logger.debug("stop_task不应该处理到巡检任务！")

    def over_task(self):    #结束任务
        if self.brun and self.task_status ==1:
            # 停止线程
            self.brun = False
            if self.target_id == 0:
                self.update_task_status(3)  # 置状态3--完成状态
            else:
                self.update_task_status(4)  #巡检任务直接结束
                # 结束任务需要收尾处理#？
            self.InstrM.init_data() #pass


    def is_over(self):
        '''
        判断当前任务是否所有节点都已结束
        :return: False-非暂停节点有代表事项，True-非粘贴节点事项都已完成
        '''
        b_over = True
        nodes = self.attack_tree.traverse_bfs()
        for node in nodes:
            if node.bwork:
                work_status = node.get_work_status()
                if work_status != 0:
                    b_over = False  # 有一个有任务就不继续进行判断
                    break
        return  b_over


    def test(self,task_id):
        root_node = TreeNode(self.target, task_id,0)  # 根节点
        self.attack_tree = AttackTree(root_node)  # 创建测试树，同时更新根节点相关内容
        self.LLM.build_initial_prompt(root_node)  # 对根节点初始化system-msg
        # 初始保存个attack_tree文件
        g_PKM.WriteData(self.attack_tree, str(task_id))




if __name__ == "__main__":
    pass