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