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.
179 lines
7.0 KiB
179 lines
7.0 KiB
#自动化测试逻辑规则控制
|
|
#统一控制规则 和 渗透测试树的维护
|
|
import json
|
|
import re
|
|
import queue
|
|
from mycode.AttackMap import AttackTree
|
|
from mycode.AttackMap import TreeNode
|
|
from mycode.LLMManager import LLMManager
|
|
from myutils.ConfigManager import myCongif #单一实例
|
|
from myutils.MyLogger_logger import LogHandler
|
|
|
|
class ControlCenter:
|
|
def __init__(self,DBM):
|
|
self.logger = LogHandler.get_logger("ControlCenter")
|
|
self.task_id = None
|
|
self.target = None
|
|
self.attack_tree = None
|
|
self.DBM = DBM
|
|
#LLM对象
|
|
self.LLM = LLMManager(myCongif.get_data("LLM_type"))
|
|
|
|
def __del__(self):
|
|
self.task_id = None
|
|
self.target = None
|
|
self.attack_tree = None
|
|
self.DBM = None
|
|
|
|
def init_cc_data(self):
|
|
#一次任务一次数据
|
|
pass
|
|
|
|
def get_user_init_info(self):
|
|
'''开始任务初，获取用户设定的基础信息'''
|
|
# ?包括是否对目标进行初始化的信息收集
|
|
return "无"
|
|
|
|
def start_do(self,target,task_id):
|
|
'''一个新任务的开始'''
|
|
self.task_id = task_id
|
|
self.target = target
|
|
#创建测试树
|
|
if self.attack_tree:
|
|
self.attack_tree = None #释放
|
|
root_node = TreeNode(target)
|
|
self.attack_tree = AttackTree(root_node)#创建测试树，同时更新根节点相关内容
|
|
#获取初始指令
|
|
know_info = self.get_user_init_info()
|
|
prompt = self.LLM.build_initial_prompt(target,know_info,root_node)
|
|
node_cmds, commands = self.LLM.get_llm_instruction(prompt,self.DBM,root_node)
|
|
# 更新tree
|
|
self.tree_manager(node_cmds,root_node)
|
|
# 分析指令入对应节点
|
|
node_list = self.instr_in_node(commands,root_node)
|
|
return node_list
|
|
|
|
#约束1：一个节点只能同时提交一次，未测试的节点不要重复
|
|
def get_llm_instruction(self,node):
|
|
#拼接结果字符串
|
|
json_strings = []
|
|
while True:
|
|
try:
|
|
item = node.res_quere.get_nowait()
|
|
json_str = json.dumps(item,ensure_ascii=False)
|
|
json_strings.append(json_str)
|
|
except queue.Empty:
|
|
break
|
|
res_str = ','.join(json_strings)
|
|
|
|
#构造本次提交的prompt
|
|
user_Prompt = f'''
|
|
当前分支路径：{node.path}
|
|
当前节点信息：
|
|
- 节点名称：{node.name}
|
|
- 节点状态：{node.status}
|
|
- 漏洞类型：{node.vul_type}
|
|
上一步结果：{res_str}
|
|
任务：生成下一步渗透测试指令或结束该节点渗透测试。
|
|
'''
|
|
node_cmds,commands = self.LLM.get_llm_instruction(user_Prompt,self.DBM,node)
|
|
#更新tree
|
|
self.tree_manager(node_cmds,node)
|
|
#分析指令入对应节点
|
|
node_list = self.instr_in_node(commands,node)
|
|
return node_list
|
|
|
|
def tree_manager(self,node_cmds,node):
|
|
'''更新渗透测试树'''
|
|
try:
|
|
node_jsons = json.loads(node_cmds)
|
|
for node_json in node_jsons:
|
|
if node_json["action"] == "add_node": #新增节点
|
|
parent_node_name = node_json["parent"]
|
|
node_name = node_json["node"]
|
|
status = node_json["status"]
|
|
#新增节点原则上应该都是当前节点增加子节点
|
|
if node.name == parent_node_name:
|
|
new_node = TreeNode(node_name,status)
|
|
node.add_child(new_node) #message的传递待验证
|
|
else:
|
|
self.logger.error("添加子节点时，遇到父节点名称不一致的，需要介入！！")
|
|
elif node_json["action"] == "update_status":
|
|
node_name = node_json["node"]
|
|
status = node_json["status"]
|
|
vul_type = node_json["vulnerability"]
|
|
if node.name == node_name:
|
|
node.status = status
|
|
node.vul_type = vul_type
|
|
else:
|
|
self.logger.error("遇到不是修改本节点状态的，需要介入！！")
|
|
else:
|
|
self.logger.error("node更新JSON遇到异常参数！")
|
|
except json.JSONDecodeError as e:
|
|
self.logger.error(f"JSON 解析失败：{e.msg}")
|
|
self.logger.error(f"错误位置：第{e.lineno}行，列{e.colno}")
|
|
except TypeError as e:
|
|
self.logger.error(f"输入类型错误：{str(e)}")
|
|
|
|
def instr_in_node(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(node_path)
|
|
if find_node:
|
|
instruction = re.sub(r'\[.*?\]', "", command,count=1,flags=re.DOTALL)
|
|
find_node.instr_queue.put(instruction)
|
|
#入输出队列
|
|
if find_node not in node_list:
|
|
node_list.append(find_node)
|
|
else:
|
|
self.logger.error(f"基于节点路径没有找到对应的节点{node_path}")
|
|
else:
|
|
self.logger.error(f"得到的指令格式不符合规范：{command}")
|
|
return node_list
|
|
|
|
|
|
def is_user_instr(self,instr):
|
|
'''
|
|
过滤需要人工确认或手动执行的指令 ---- 待完善
|
|
:param instr:
|
|
:return:
|
|
'''
|
|
#if instr.startswith("curl") or instr.startswith("http") or instr.startswith("wget"):
|
|
if instr.startswith("http") or instr.startswith("wget") or instr.startswith("ssh"):
|
|
return True
|
|
|
|
#指令入队列，待修改
|
|
def instr_in_quere(self,instr_list):
|
|
'''
|
|
对于运行需要较长时间的不强求同一批次返回给LLM
|
|
:param instr_list:
|
|
:return:
|
|
'''
|
|
for instr in instr_list:
|
|
if self.is_user_instr(instr):
|
|
self.user_instr.put(instr)
|
|
print(f"需要人工确认的指令{instr}")
|
|
else:
|
|
matched =False
|
|
for prefix in self.long_time_instr:
|
|
if instr.startswith(prefix):
|
|
matched =True
|
|
if not matched:
|
|
with self.lock:
|
|
self.batch_num += 1 #非耗时指令+1
|
|
print(f"&&&&&&当前batch_num:{self.batch_num}")
|
|
else:
|
|
with self.lock:
|
|
self.long_instr_num +=1 #耗时指令数量+1
|
|
# 指令入队列
|
|
self.instr_queue.put(instr)
|
|
|
|
def stop_do(self):
|
|
#清空数据
|
|
self.task_id = None
|
|
self.target = None
|
|
self.attack_tree = None
|