@ -98,12 +98,14 @@ class TaskObject:
th_DBM = DBManager ( )
th_DBM . connect ( )
th_index = index
bnode_work = False
while self . brun :
if self . task_status == 1 :
try :
llm_type = 1
work_node = self . instr_node_queue . get ( block = False ) #正常一个队列中一个节点只会出现一次，进而也就只有一个线程在处理
# 开始执行指令
bnode_work = True
results = [ ]
while True :
instruction = work_node . get_instr ( )
@ -137,8 +139,6 @@ class TaskObject:
ext_params , work_node . path )
else :
self . logger . error ( " 数据库连接失败！！ " )
if work_node . do_sn > = 10 : #该节点执行指令已超过10条
llm_type = 10
#暂存结果
oneres = { ' 执行指令 ' : instr , ' 结果 ' : reslut }
results . append ( oneres )
@ -147,11 +147,15 @@ class TaskObject:
#指令都执行结束后，入节点待提交队列
str_res = json . dumps ( results , ensure_ascii = False )
# 提交llm待处理任务 --更新节点work_status
if work_node . do_sn > = myCongif . get_data ( " max_do_sn " ) : # 该节点执行指令已超过10条
llm_type = 10
self . put_node_reslist ( work_node , str_res , llm_type )
# 保存记录
g_PKM . WriteData ( self . attack_tree , str ( self . task_id ) )
except queue . Empty :
self . no_work_to_do ( ) #判断是否需要把当前任务的无工作状态推送到前端
if bnode_work :
bnode_work = False
self . no_work_to_do ( ) #判断是否需要把当前任务的无工作状态推送到前端
time . sleep ( self . sleep_time )
else : #不是工作状态则休眠
time . sleep ( self . sleep_time )
@ -168,11 +172,13 @@ class TaskObject:
th_DBM = DBManager ( )
th_DBM . connect ( )
th_index = index
bnode_work = False
while self . brun :
if self . task_status == 1 :
try :
llm_node = self . llm_node_queue . get ( block = False ) #获取一个待处理节点
#开始处理
bnode_work = True
tmp_commands = [ ]
# {llm_node.status} --- 暂时固化为未完成
user_Prompt = f '''
@ -192,7 +198,11 @@ class TaskObject:
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要更新
post_time = get_local_timestr ( )
bsuccess , node_cmds , commands , reasoning_content , content = self . LLM . get_llm_instruction ( prompt , llm_node ) # message要更新
if not bsuccess :
self . logger . error ( f " 模型接口调用出错： { content } " )
continue #丢弃 --若需要再次尝试，把llm_data再入队列
# LLM记录存数据库
if th_DBM . ok :
llm_node . llm_sn + = 1
@ -210,16 +220,24 @@ class TaskObject:
bok , new_commands , iadd_node = self . tree_manager ( node_cmds , llm_node , commands , th_DBM )
# 分析指令入对应节点
if bok : # 节点指令若存在错误，测试指令都不处理，需要LLM重新生成
tmp_commands . extend ( new_commands )
self . doing_llm_list [ th_index ] = " "
#tmp_commands.extend(new_commands)
# 测试指令入节点待处理队列 --同时修改节点的work_status
self . put_node_instrlist ( new_commands , llm_node , iadd_node )
#测试指令入节点待处理队列 --同时修改节点的work_status
self . put_node_instrlist ( tmp_commands , llm_node , iadd_node )
self . doing_llm_list [ th_index ] = " "
#一个节点完成，节点树持久化---待验证是否有局部更新持久化的方案
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
except queue . Empty :
#self.logger.debug("llm队列中暂时无新的提交任务！")
self . no_work_to_do ( ) # 判断是否需要把当前任务的无工作状态推送到前端
if bnode_work :
bnode_work = False
self . no_work_to_do ( ) # 判断是否需要把当前任务的无工作状态推送到前端
time . sleep ( self . sleep_time )
else :
time . sleep ( self . sleep_time )
@ -228,8 +246,8 @@ class TaskObject:
bsuccess = False
with self . is_had_work_lock :
if not self . is_had_work :
bsuccess = True # 这不return 是怕不释放锁
self . is_had_work = True
bsuccess = True #这不return 是怕不释放锁
return bsuccess
def no_work_to_do ( self ) : #任务单步状态控制-- 非工作中
@ -247,16 +265,17 @@ class TaskObject:
#推送是否有工作任务的状态到前端,
with self . is_had_work_lock :
if self . is_had_work : #如果已经是False那就不需要修改了
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 :
cur_time = get_local_timestr ( )
@ -277,6 +296,10 @@ class TaskObject:
else :
print ( f " LLM线程- { index } 在执行指令： { self . doing_llm_list [ index ] } " )
index + = 1
#处理点修复操作
icount + = 1
if icount == 5 :
pass
#休眠60
time . sleep ( 60 )
@ -347,35 +370,46 @@ class TaskObject:
return command
def put_node_instrlist ( self , commands , node , iadd_node ) : #如果当前节点没有进一般指令返回，需要修改节点执行状态
node_list = [ ] #一次返回的测试指令
if not node :
return
node_list = [ ] #有待办指令的节点
for command in commands :
command = self . replace_error_instr ( command )
# 使用正则匹配方括号中的node_path（非贪婪模式）
match = re . search ( r ' \ [(.*?) \ ] ' , command )
if match :
node_path = match . group ( 1 )
node_name = node_path . split ( " -> " ) [ - 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 )
find_node = None
if node_name == node . name :
find_node = node
else :
for child_node in node . children : #暂时只找一层
if node_name == child_node . name :
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 ( instruction , node . parent_messages , node . cur_messages )
find_node . add_instr ( instruction , 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 ) #待执行
self . update_node_work_status ( find_node , 1 ) #待执行指令
else : #如果还没找到就暂时放弃
self . logger . error ( f " 基于节点路径 没有找到对应的节点{ node_path } ,父节点都没找到！ 当前节点 { node . path } " ) #丢弃该指令
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--无待执行任务
@ -408,9 +442,9 @@ class TaskObject:
async def put_one_node ( self , node ) :
#提交某个节点的代表任务
if self . task_status == 1 and self . work_type == 0 and node . bwork :
node_status = node . get_work_status ( )
if node_status == 2 or node_status == 4 :
return False , "
# node_status = node.get_work_status( )
# if node_status == 2 or node_status == 4 :
# return False," 当前节点正在执行任务，请稍后点击单步！"
if not node . is_instr_empty ( ) : #待执行指令有值
if not node . is_llm_empty ( ) :
self . logger . error ( f " { node . path } 即存在待执行指令，还存在待提交的llm，需要人工介入！！ " )
@ -469,7 +503,7 @@ class TaskObject:
# 构造本次提交的prompt
ext_Prompt = f '''
上一步结果 ： { str_res }
任务 ： 生成下一步渗透测试 指令或判断是否完成该节点测试 。
任务 ： 请 生成下一步指令。
'''
elif llm_type == 2 : # llm返回的指令存在问题，需要再次请求返回
ext_Prompt = f '''
@ -488,9 +522,10 @@ class TaskObject:
任务 ： 请根据格式要求 ， 重新生成该测试指令 。
'''
elif llm_type == 10 :
max_do_sn = myCongif . get_data ( " max_do_sn " )
ext_Prompt = f '''
上一步结果 ： { str_res }
任务 ： 当前节点已执行10 次测试指令 ， 若无新发现请结束该节点的测试 ， 若有新发现请生成子节点 ， 在子节点推进下一步测试 。
任务 ： 当前节点已执行超过 { max_do_sn } 次测试指令 ， 若无新发现请结束该节点的测试 ， 若有新发现请生成子节点 ， 在子节点推进下一步测试 。
'''
else :
self . logger . debug ( " 意外的类型参数 " )
@ -500,14 +535,14 @@ class TaskObject:
return user_Prompt
#添加子节点
def add_children_node ( self , parent_node , children_names , cur_message , status = " 未完成 " ) :
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 ) ) # 去重
for child_name in unique_names :
if child_name not in existing_names :
# 添加节点
new_node = TreeNode ( child_name , parent_node . task_id , status )
parent_node . add_child ( new_node , cur_message ) # message的传递待验证
parent_node . add_child ( new_node )
#existing_names.add(child_name) # 更新集合 -- 已经去重过了，不需要在添加到比对
#处理节点指令
@ -525,11 +560,44 @@ class TaskObject:
# 提交llm待处理任务
self . put_node_reslist ( node , strerror , 2 )
return False , 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 = [ ]
add_node_names = [ ]
for node_json in node_cmds :
action = node_json [ " action " ]
if action == " add_node " : # 新增节点
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 )
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 )
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 )
break
if not badd :
self . logger . error ( f " 添加子节点失败！父节点不是当前节点，不是当前节点的父节点，不是当前节点的子节点，需要介入！！ { node_json } ---当前节点为： { node . path } " ) # 丢弃该节点
else : # 未处理的节点指令添加到list
residue_cmd_no_add . append ( node_json )
#处理on_instruction
residue_node_cmds = [ ]
no_instr_nodes = [ ]
#如果有on_instruction，先补全指令保障信息链的完整
for node_cmd in node_cmds :
for node_cmd in residue_cmd_no_add :
action = node_cmd [ " action " ]
if action == " no_instruction " :
node_names = node_cmd [ " nodes " ] . replace ( ' ， ' , ' , ' ) . split ( ' , ' )
@ -542,57 +610,22 @@ class TaskObject:
break
if bcommand : # 如果存在测试指令，则不把该节点放入补充信息llm任务---尝试不对比是否有返回指令，DS会一直返回指令，还返回on_instruction
continue
no_instr_nodes . append ( node_name )
#判断是否有对应节点---原则上只允许同批次add的节点没有添加指令的情况
if node_name in add_node_names :
no_instr_nodes . append ( node_name )
else :
self . logger . error ( " 遇到一次不在add_node中，但在no_instr_nodes中的数据 " )
#粗暴的做法，添加在当前节点下
self . add_children_node ( node , [ node_name ] )
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 " ] . replace ( ' ， ' , ' , ' ) . split ( ' , ' )
# 新增节点原则上应该都是当前节点增加子节点
if node . name == parent_node_name or parent_node_name . endswith ( node . name ) : #2233ai，节点名称字段会返回整个路径
#添加当前节点的子节点 -- 这是标准情况
self . add_children_node ( node , node_names )
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 )
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 )
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重新调整了节点指令格式定义
for node_json in residue_node_cmds : #2025-4-11重新调整了节点指令格式定义
action = node_json [ " action " ]
if action == " find_vul " :
node_name = node_json [ " node " ]
@ -616,31 +649,24 @@ class TaskObject:
self . logger . error ( " 漏洞信息错误 " )
continue
else :
str_user = f " 遇到不是修改本节点状态 的，需要介入！！ { node_json } "
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 } "
str_user = f " 遇到不是修改本节点状态的，需要介入！！ { node_json } --当前节点 { node . path }
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
return True , commands , len ( add_node_names )
#阻塞轮询补充指令
def get_other_instruction ( self , nodes , DBM , cur_node ) :
res_str = ' , ' . join ( nodes )
new_commands = [ ]
ierror = 0
while res_str :
self . logger . debug ( f " 开始针对f { res_str } 这些节点请求测试指令 " )
user_Prompt = f '''
@ -655,10 +681,15 @@ class TaskObject:
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 ,
bsuccess , node_cmds , commands , reasoning_content , content , post_time = self . LLM . get_llm_instruction ( user_Prompt ,
cur_node ) # message要更新
if not bsuccess :
self . logger . error ( f " 模型接口调用出错： { content } " )
ierror + = 1
if ierror == 3 : #重试3词
break
continue # res_str没有调整，重复使用
res_str = " "
# 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 )
@ -676,6 +707,8 @@ class TaskObject:
tmp_nodes . append ( node_name )
res_str = ' , ' . join ( tmp_nodes )
break
else : #其他节点指令不处理
self . logger . error ( f " 遇到一次no_instruction补充指令时返回了其他节点指令 { node_cmds } " )
self . logger . debug ( " 未添加指令的节点，都已完成指令的添加！ " )
return new_commands
