You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
578 lines
30 KiB
578 lines
30 KiB
3 weeks ago
|
'''
|
||
|
渗透测试任务管理类 一次任务的闭合性要检查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
|
||
|
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.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 = [] #线程句柄list
|
||
|
self.instr_node_queue = queue.Queue() #待执行指令的节点队列
|
||
|
# self.long_instr_num = 0 #耗时指令数量
|
||
|
# self.long_time_instr = ['nikto'] #耗时操作不计入批量执行的数量,不加不减
|
||
|
self.lock = threading.Lock() #线程锁
|
||
|
self.node_num = 0 #在处理Node线程的处理
|
||
|
#llm执行相关--------
|
||
|
#self.max_thread_num = 1 # 控制最大并发指令数量 --- 多线程的话节点树需要加锁
|
||
|
self.llmth_list = [] # llm线程list
|
||
|
self.llm_node_queue = queue.Queue() #待提交LLM的节点队列
|
||
|
#自检线程--------
|
||
|
self.check_th = None #自检线程句柄
|
||
|
|
||
|
#---------------三个线程------------
|
||
|
#测试指令执行线程
|
||
|
def do_worker_th(self):
|
||
|
#线程的dbm需要一个线程一个
|
||
|
th_DBM = DBManager()
|
||
|
th_DBM.connect()
|
||
|
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
|
||
|
start_time = get_local_timestr() # 指令执行开始时间
|
||
|
instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction)
|
||
|
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):
|
||
|
'''
|
||
|
几个规则--TM的work线程同
|
||
|
1.线程获取一个节点后,其他线程不能再获取这个节点(遇到被执行的节点,直接放弃执行)--- 加了没办法保存中间结果进行测试
|
||
|
2.llm返回的指令,只可能是该节点,或者是子节点的,不符合这条规则的都不处理,避免llm处理混乱。
|
||
|
:return:
|
||
|
'''
|
||
|
# 线程的dbm需要一个线程一个
|
||
|
th_DBM = DBManager()
|
||
|
th_DBM.connect()
|
||
|
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)
|
||
|
# 提交llm请求返回数据--并对返回数据进行处理,节点指令直接执行,测试指令根据工作模式执行
|
||
|
node_cmds, commands,reasoning_content, content, post_time = self.LLM.get_llm_instruction(prompt,llm_node) # message要更新
|
||
|
# 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 = 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)
|
||
|
#一个节点完成,节点树持久化---待验证是否有局部更新持久化的方案
|
||
|
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)
|
||
|
|
||
|
#自检线程
|
||
|
def th_check(self):
|
||
|
print("自检线程待实现中!")
|
||
|
|
||
|
#------------入两个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 put_node_instrlist(self, commands, node): #如果当前节点没有进一般指令返回,需要修改节点执行状态
|
||
|
node_list = [] #一次返回的测试指令
|
||
|
for command in commands:
|
||
|
# 使用正则匹配方括号中的node_path(非贪婪模式)
|
||
|
match = re.search(r'\[(.*?)\]', command)
|
||
|
if match:
|
||
|
node_path = match.group(1)
|
||
|
#'''强制约束,不是本节点或者是子节点的指令不处理'''
|
||
|
find_node = self.attack_tree.find_node_by_nodepath_parent(node_path,node)
|
||
|
if find_node:
|
||
|
instruction = re.sub(r'\[.*?\]', "", command,count=1,flags=re.DOTALL)
|
||
|
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 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
|
||
|
|
||
|
#对节点指令进行校验
|
||
|
bok,strerror = self.CCM.verify_node_cmds(node_cmds)
|
||
|
if not bok: #节点指令存在问题,则不进行后续处理,提交一个错误反馈任务
|
||
|
# 提交llm待处理任务
|
||
|
self.put_node_reslist(node, strerror, 2)
|
||
|
return False,commands
|
||
|
|
||
|
#先执行add_node操作
|
||
|
residue_node_cmds = []
|
||
|
badd_node = False
|
||
|
for node_json in node_cmds:
|
||
|
action = node_json["action"]
|
||
|
if action == "add_node": # 新增节点
|
||
|
badd_node = True
|
||
|
parent_node_name = node_json["parent"]
|
||
|
status = node_json["status"]
|
||
|
node_names = node_json["nodes"].split(',')
|
||
|
# 新增节点原则上应该都是当前节点增加子节点
|
||
|
if node.name == parent_node_name or parent_node_name.endswith(node.name): #2233ai,节点名称字段会返回整个路径
|
||
|
for node_name in node_names:
|
||
|
# 判重---遇到过补充未生成指令的节点时,返回了新增这些节点的指令
|
||
|
bfind = False
|
||
|
for node_child in node.children:
|
||
|
if node_child.name == node_name:
|
||
|
bfind = True
|
||
|
break
|
||
|
if not bfind:
|
||
|
# 添加节点
|
||
|
new_node = TreeNode(node_name, node.task_id, status)
|
||
|
node.add_child(new_node) # message的传递待验证
|
||
|
elif node.parent.name == parent_node_name or parent_node_name.endswith(node.parent.name):
|
||
|
#是添加当前节点的平级节点(当前节点的父节点下添加子节点) --使用2233ai-o3时遇到的情况
|
||
|
for node_name in node_names:
|
||
|
# 判重---遇到过补充未生成指令的节点时,返回了新增这些节点的指令
|
||
|
bfind = False
|
||
|
for node_child in node.parent.children:
|
||
|
if node_child.name == node_name:
|
||
|
bfind = True
|
||
|
break
|
||
|
if not bfind:
|
||
|
# 添加节点
|
||
|
new_node = TreeNode(node_name, node.task_id, status)
|
||
|
node.parent.add_child(new_node)
|
||
|
else:
|
||
|
self.logger.error(f"添加子节点时,遇到父节点名称不一致的,需要介入!!{node_json}") # 丢弃该节点
|
||
|
else:#其他指令添加到list
|
||
|
residue_node_cmds.append(node_json)
|
||
|
|
||
|
if badd_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_node_cmds:
|
||
|
action = node_json["action"]
|
||
|
if action == "update_status":
|
||
|
node_name = node_json["node"]
|
||
|
status = node_json["status"]
|
||
|
vul_type = "未发现"
|
||
|
if node.name == node_name or node_name.endswith(node_name):
|
||
|
node.status = status
|
||
|
if "vulnerability" in node_json:
|
||
|
#{\"name\":\"漏洞名称\",\"risk\":\"风险等级(低危/中危/高危)\",\"info\":\"补充信息(没有可为空)\"}};
|
||
|
#vul_type = json.dumps(node_json["vulnerability"],ensure_ascii=False) #json转字符串
|
||
|
try:
|
||
|
node.vul_type = node_json["vulnerability"]["name"]
|
||
|
node.vul_grade = node_json["vulnerability"]["risk"]
|
||
|
node.vul_info = node_json["vulnerability"]["info"]
|
||
|
except:
|
||
|
self.logger.error("漏洞信息错误")
|
||
|
#node.vul_type = vul_type
|
||
|
else:
|
||
|
str_user = f"遇到不是修改本节点状态的,需要介入!!{node_json}"
|
||
|
self.logger.error(str_user)
|
||
|
#self.need_user_know(str_user,node)
|
||
|
elif action == "no_instruction":
|
||
|
#返回的未生成指令的数据进行校验:1.要有数据;2.节点不是当前节点就是子节点
|
||
|
nodes = []
|
||
|
node_names = node_json["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
|
||
|
#验证对应节点是否已经创建---本节点或子节点,其他节点不处理(更狠一点就是本节点都不行)
|
||
|
bfind = False
|
||
|
if node_name == node.name:
|
||
|
bfind = True
|
||
|
nodes.append(node_name)
|
||
|
else:
|
||
|
for child_node in node.children:
|
||
|
if child_node.name == node_name:
|
||
|
bfind = True
|
||
|
nodes.append(node_name)
|
||
|
break
|
||
|
if not bfind:
|
||
|
self.logger.debug(f"没有找到该节点{node_name}")
|
||
|
if nodes: #阻塞式,在当前节点提交补充信息,完善节点指令 -- 优势是省token
|
||
|
new_commands = self.get_other_instruction(nodes,DBM,node)
|
||
|
commands.extend(new_commands)
|
||
|
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
|
||
|
|
||
|
#阻塞轮询补充指令
|
||
|
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)
|
||
|
w_th.start()
|
||
|
self.workth_list.append(w_th)
|
||
|
#启动llm提交线程--llm暂时单线程,多线程处理时attack_tree需要加锁
|
||
|
l_th = threading.Thread(target=self.th_llm_worker)
|
||
|
l_th.start()
|
||
|
self.llmth_list.append(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()
|
||
|
#结束任务需要收尾处理#?
|
||
|
|
||
|
def do_work(self,taks_id,work_type):
|
||
|
'''
|
||
|
手动控制程序
|
||
|
:param taks_id:
|
||
|
:param work_type:
|
||
|
:return:
|
||
|
'''
|
||
|
pass
|
||
|
|
||
|
|
||
|
if __name__ == "__main__":
|
||
|
pass
|