|
|
@ -40,7 +40,11 @@ class ControlCenter: |
|
|
|
pass |
|
|
|
|
|
|
|
def get_user_init_info(self): |
|
|
|
'''开始任务初,获取用户设定的基础信息''' |
|
|
|
'''开始任务初,获取用户设定的基础信息,初始信息可以分为两块: |
|
|
|
1.提交llm的补充信息 和 保留在本地的信息(如工具补充参数等,cookie) |
|
|
|
2.用户可以设置全局,和指定节点(端口) |
|
|
|
3.补充测试节点 |
|
|
|
''' |
|
|
|
# ?包括是否对目标进行初始化的信息收集 |
|
|
|
return {"已知信息":"无"} |
|
|
|
|
|
|
@ -82,10 +86,14 @@ class ControlCenter: |
|
|
|
'''获取该节点的llm提交数据,会清空type和res_quere''' |
|
|
|
llm_type = node.llm_type |
|
|
|
node.llm_type = 0 #回归0 |
|
|
|
res_list = node.res_quere[:] #复制独立副本 |
|
|
|
node.res_quere.clear() #清空待处理数据 |
|
|
|
res_list = node.res_quere[:] #浅拷贝,复制第一层 |
|
|
|
node.res_quere.clear() #清空待处理数据,相当于把原应用关系接触 |
|
|
|
return llm_type,res_list |
|
|
|
|
|
|
|
def restore_one_llm_work(self,node,llm_type,res_list): |
|
|
|
node.llm_type = llm_type |
|
|
|
node.res_quere = res_list |
|
|
|
|
|
|
|
#llm请求提交线程 |
|
|
|
def th_llm_worker(self):#LLM没有修改的全局变量,应该可以共用一个client |
|
|
|
''' |
|
|
@ -104,7 +112,7 @@ class ControlCenter: |
|
|
|
self.get_llm_instruction(node,th_DBM) |
|
|
|
#释放锁 |
|
|
|
|
|
|
|
# 暂存状态--测试时使用 |
|
|
|
# 暂存状态--测试时使用--限制条件llm工作线程只能未1个 |
|
|
|
with open("attack_tree", 'wb') as f: |
|
|
|
pickle.dump(self.attack_tree, f) |
|
|
|
|
|
|
@ -137,8 +145,7 @@ class ControlCenter: |
|
|
|
# 构造本次提交的prompt |
|
|
|
ext_Prompt = f''' |
|
|
|
上一步结果:{res_str} |
|
|
|
任务:生成下一步渗透测试指令或判断是否完成该节点测试 |
|
|
|
请确保生成的指令满足以下 |
|
|
|
任务:生成下一步渗透测试指令或判断是否完成该节点测试。 |
|
|
|
''' |
|
|
|
elif llm_type ==2: #llm返回的指令存在问题,需要再次请求返回 |
|
|
|
ext_Prompt = f''' |
|
|
@ -146,25 +153,17 @@ class ControlCenter: |
|
|
|
错误信息:{res_str} |
|
|
|
任务:请按格式要求重新生成该节点上一次返回中生成的所有指令。 |
|
|
|
''' |
|
|
|
elif llm_type ==3: #已生成节点,但未生成测试指令 |
|
|
|
ext_Prompt = f''' |
|
|
|
反馈类型:需要继续补充信息 |
|
|
|
缺失信息:{res_str} |
|
|
|
任务: |
|
|
|
1.请生成这些节点的测试指令; |
|
|
|
2.这些节点的父节点为当前节点,请正确生成这些节点的节点路径; |
|
|
|
3.若还有节点未能生成测试指令,必须返回未生成指令的节点列表。 |
|
|
|
''' |
|
|
|
elif llm_type ==4: #未生成节点列表 |
|
|
|
ext_Prompt = f''' |
|
|
|
反馈类型:需要继续补充信息 |
|
|
|
缺失信息:{res_str} |
|
|
|
任务: |
|
|
|
1.请生成这些节点的新增节点指令,并生成对应的测试指令; |
|
|
|
2.这些节点的父节点为当前节点,请正确生成这些节点的节点路径; |
|
|
|
3.若节点未能全部新增,必须返回未新增的节点列表 |
|
|
|
4.若有未生成指令的节点,必须返回未生成指令的节点列表。 |
|
|
|
''' |
|
|
|
# ''' |
|
|
|
# elif llm_type ==4: #未生成节点列表 |
|
|
|
# ext_Prompt = f''' |
|
|
|
# 反馈类型:需要继续补充信息 |
|
|
|
# 缺失信息:{res_str} |
|
|
|
# 任务: |
|
|
|
# 1.请生成这些节点的新增节点指令,并生成对应的测试指令; |
|
|
|
# 2.这些节点的父节点为当前节点,请正确生成这些节点的节点路径; |
|
|
|
# 3.若节点未能全部新增,必须返回未新增的节点列表 |
|
|
|
# 4.若有未生成指令的节点,必须返回未生成指令的节点列表。 |
|
|
|
# ''' |
|
|
|
elif llm_type ==5: |
|
|
|
ext_Prompt = f''' |
|
|
|
反馈类型:测试指令格式错误 |
|
|
@ -186,10 +185,10 @@ class ControlCenter: |
|
|
|
2.LLM的回复开始反复时(有点难判断) |
|
|
|
''' |
|
|
|
# 更新tree |
|
|
|
bok = self.tree_manager(node_cmds, node,commands) |
|
|
|
bok,new_commands = self.tree_manager(node_cmds, node,commands,DBM) |
|
|
|
# 分析指令入对应节点 |
|
|
|
if bok: #节点指令若存在错误,测试指令都不处理,需要LLM重新生成 |
|
|
|
node_list = self.instr_in_node(commands, node) |
|
|
|
node_list = self.instr_in_node(new_commands, node) |
|
|
|
# 插入TM的node_queue中,交TM线程处理---除了LLM在不同的请求返回针对同一节点的测试指令,正常业务不会产生两次进队列 |
|
|
|
for node in node_list: |
|
|
|
self.TM.node_queue.put(node) |
|
|
@ -229,27 +228,29 @@ class ControlCenter: |
|
|
|
self.put_one_llm_work(strerror,node,2) |
|
|
|
return False |
|
|
|
|
|
|
|
def tree_manager(self,node_cmds,node,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 |
|
|
|
return True,commands |
|
|
|
#对节点指令进行校验 |
|
|
|
if not self.verify_node_cmds(node_cmds,node): |
|
|
|
return False #节点指令存在问题,终止执行 |
|
|
|
#执行节点操作 |
|
|
|
return False,commands #节点指令存在问题,终止执行 |
|
|
|
|
|
|
|
#执行节点操作---先执行add_node,怕返回顺序不一直 |
|
|
|
residue_node_cmds = [] |
|
|
|
for node_json in node_cmds: |
|
|
|
action = node_json["action"] |
|
|
|
if action == "add_node": # 新增节点 |
|
|
|
parent_node_name = node_json["parent"] |
|
|
|
# 新增节点原则上应该都是当前节点增加子节点 |
|
|
|
if node.name == parent_node_name: |
|
|
|
if node.name == parent_node_name or parent_node_name.endswith(node.name): |
|
|
|
status = node_json["status"] |
|
|
|
node_names = node_json["nodes"].split(',') |
|
|
|
for node_name in node_names: |
|
|
|
#判重 |
|
|
|
# 判重---遇到过补充未生成指令的节点时,返回了新增这些节点的指令 |
|
|
|
bfind = False |
|
|
|
for node_child in node.children: |
|
|
|
if node_child.name == node_name: |
|
|
@ -261,17 +262,32 @@ class ControlCenter: |
|
|
|
node.add_child(new_node) # message的传递待验证 |
|
|
|
else: |
|
|
|
self.logger.error(f"添加子节点时,遇到父节点名称不一致的,需要介入!!{node_json}") # 丢弃该节点 |
|
|
|
elif action == "update_status": |
|
|
|
else:#其他指令添加到list |
|
|
|
residue_node_cmds.append(node_json) |
|
|
|
|
|
|
|
#执行剩余的节点指令--不分先后 |
|
|
|
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 "vulnerability" in node_json: |
|
|
|
vul_type = json.dumps(node_json["vulnerability"]) |
|
|
|
if node.name == 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.name = 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: |
|
|
|
self.logger.error(f"遇到不是修改本节点状态的,需要介入!!{node_json}") |
|
|
|
str_user = f"遇到不是修改本节点状态的,需要介入!!{node_json}" |
|
|
|
self.logger.error(str_user) |
|
|
|
self.need_user_know(str_user,node) |
|
|
|
elif action == "no_instruction": |
|
|
|
#返回的未生成指令的数据进行校验:1.要有数据;2.节点不是当前节点就是子节点 |
|
|
|
nodes = [] |
|
|
@ -285,29 +301,68 @@ class ControlCenter: |
|
|
|
break |
|
|
|
if bcommand: #如果存在测试指令,则不把该节点放入补充信息llm任务 |
|
|
|
continue |
|
|
|
#验证对应节点是否已经创建---本节点或子节点,其他节点不处理(更狠一点就是本节点都不行) |
|
|
|
if node_name == node.name: |
|
|
|
nodes.append(node_name) |
|
|
|
# str_add = "请生成测试指令" |
|
|
|
# self.put_one_llm_work(str_add,node,1) |
|
|
|
else: |
|
|
|
for child_node in node.children: |
|
|
|
if child_node.name == node_name: |
|
|
|
nodes.append(node_name) |
|
|
|
# str_add = "无" |
|
|
|
# self.put_one_llm_work(str_add, child_node, 1) |
|
|
|
break |
|
|
|
if nodes: #找到对应的节点才返回 |
|
|
|
str_nodes = ",".join(nodes) |
|
|
|
str_add = {"已新增但未生成测试指令的节点":str_nodes} |
|
|
|
# 提交一个错误反馈任务--但继续后续工作 |
|
|
|
self.put_one_llm_work(str_add, node, 3) |
|
|
|
self.logger.debug(f"已新增但未生成指令的节点有:{nodes}") |
|
|
|
elif action == "no_create": |
|
|
|
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.put_one_llm_work(str_add, node, 4) |
|
|
|
self.logger.debug(f"未新增的节点有:{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 |
|
|
|
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要更新 |
|
|
|
#把返回的测试指令进行追加 |
|
|
|
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 instr_in_node(self,commands,node): |
|
|
|
node_list = [] #一次返回的测试指令 |
|
|
@ -334,6 +389,10 @@ class ControlCenter: |
|
|
|
# 3.独立队列处理 |
|
|
|
return node_list |
|
|
|
|
|
|
|
#需要用户确认的信息--待完善 |
|
|
|
def need_user_know(self,strinfo,node): |
|
|
|
pass |
|
|
|
|
|
|
|
#待修改 |
|
|
|
def is_user_instr(self,instr): |
|
|
|
''' |
|
|
@ -380,3 +439,5 @@ class ControlCenter: |
|
|
|
self.brun =False |
|
|
|
for th in self.llmth_list: |
|
|
|
th.jion() |
|
|
|
|
|
|
|
|
|
|
|