''' 渗透测试任务管理类 一次任务的闭合性要检查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