|
|
@ -52,6 +52,7 @@ class TaskObject: |
|
|
|
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() |
|
|
@ -237,7 +238,6 @@ class TaskObject: |
|
|
|
llm_node = self.llmth_node_list[th_index] |
|
|
|
#开始处理 |
|
|
|
bnode_work = True |
|
|
|
tmp_commands = [] |
|
|
|
# {llm_node.status} --- 暂时固化为未完成 |
|
|
|
user_Prompt = f''' |
|
|
|
当前分支路径:{llm_node.path} |
|
|
@ -246,42 +246,37 @@ class TaskObject: |
|
|
|
- 节点状态:未完成 |
|
|
|
- 漏洞类型:{llm_node.vul_type} |
|
|
|
''' |
|
|
|
while True: |
|
|
|
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"] |
|
|
|
#判断执行次数 |
|
|
|
llm_node.llm_sn += 1 |
|
|
|
if llm_node.llm_sn == max_llm_sn: #提交次数达到上限后,提示LLM结束该节点任务 |
|
|
|
llm_type = 10 |
|
|
|
#该节点的剩余任务不执行,若有--暂定 |
|
|
|
llm_node.clear_res() |
|
|
|
|
|
|
|
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) |
|
|
|
fake_prompt = self.DataFilter.filter_prompt(prompt) #目标脱敏 |
|
|
|
|
|
|
|
self.doing_llm_list[th_index] = prompt |
|
|
|
self.doing_llm_list[th_index] = prompt #执行更新 |
|
|
|
# 提交llm请求返回数据--并对返回数据进行处理,节点指令直接执行,测试指令根据工作模式执行 |
|
|
|
post_time = get_local_timestr() |
|
|
|
bsuccess,node_cmds, commands,reasoning_content, content = self.LLM.get_llm_instruction(fake_prompt,llm_node,self.DataFilter) # message要更新 --llm_node只使用messages,都是脱敏后的数据 |
|
|
|
if not bsuccess: |
|
|
|
self.logger.error(f"模型接口调用出错:{content}") |
|
|
|
continue #丢弃 --若需要再次尝试,把llm_data再入队列 |
|
|
|
# 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("数据库连接失败!") |
|
|
|
''' |
|
|
|
对于LLM返回的错误处理机制 |
|
|
|
1.验证节点是否都有测试指令返回 |
|
|
|
2.LLM的回复开始反复时(有点难判断) |
|
|
|
''' |
|
|
|
# 更新tree |
|
|
|
bok, new_commands,iadd_node = self.tree_manager(node_cmds, llm_node, commands, th_DBM) |
|
|
|
# 分析指令入对应节点 |
|
|
@ -298,7 +293,7 @@ class TaskObject: |
|
|
|
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: |
|
|
@ -425,10 +420,11 @@ class TaskObject: |
|
|
|
|
|
|
|
#修改节点的执行状态,并需要基于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: #llm执行完成后会发送单独的指令更新树,所以不发送1更新节点了 |
|
|
|
if work_status != 1 and old_work_status != 0: #llm执行完成后会发送单独的指令更新树,所以不发送1更新节点了 |
|
|
|
#判断是否是web端最新获取数据的task |
|
|
|
if self.taskM.web_cur_task == self.task_id: |
|
|
|
idatatype = 1 |
|
|
@ -476,25 +472,36 @@ class TaskObject: |
|
|
|
return find_node |
|
|
|
|
|
|
|
def replace_error_instr(self,command): |
|
|
|
command = command.replace("| |","||") |
|
|
|
command = command.replace("& &","&&") |
|
|
|
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) |
|
|
|
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 = None |
|
|
|
if node_name == node.name: #指令是当前节点的 |
|
|
|
#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: #暂时只找一层 |
|
|
@ -502,21 +509,20 @@ class TaskObject: |
|
|
|
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(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) #待执行指令 |
|
|
|
else:#如果还没找到就暂时放弃 |
|
|
|
self.logger.error(f"没有找到指令对应的节点:{node_path},当前节点{node.path}")#丢弃该指令 |
|
|
|
else: |
|
|
|
self.logger.error(f"得到的指令格式不符合规范:{command}")#丢弃该指令--- |
|
|
|
# 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,但有可能当前节点没有其他指令,不会触发提交,另外就算提交,会不会产生预设范围外的返回,不确定; |
|
|
@ -617,10 +623,14 @@ class TaskObject: |
|
|
|
|
|
|
|
#获取本次的提交提示词 |
|
|
|
def get_llm_prompt(self,llm_type,str_res,user_Prompt): |
|
|
|
if llm_type == 0: |
|
|
|
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} |
|
|
|
任务:请开始对该目标的渗透测试工作。 |
|
|
|
已知信息:{str_res} |
|
|
|
任务:请根据当前节点信息生成下一步指令。 |
|
|
|
''' |
|
|
|
elif llm_type == 1: # 提交指令执行结果 --- 正常提交 |
|
|
|
# 构造本次提交的prompt |
|
|
@ -630,7 +640,7 @@ class TaskObject: |
|
|
|
''' |
|
|
|
elif llm_type == 2: # llm返回的指令存在问题,需要再次请求返回 |
|
|
|
ext_Prompt = f''' |
|
|
|
反馈类型:节点指令格式错误 |
|
|
|
反馈类型:返回的信息不符合规则 |
|
|
|
错误信息:{str_res} |
|
|
|
任务:请按格式要求重新生成该节点上一次返回中生成的所有指令。 |
|
|
|
''' |
|
|
@ -669,23 +679,200 @@ class TaskObject: |
|
|
|
parent_node.add_child(new_node) |
|
|
|
#existing_names.add(child_name) # 更新集合 -- 已经去重过了,不需要在添加到比对 |
|
|
|
|
|
|
|
#处理节点指令 |
|
|
|
def tree_manager(self,node_cmds,node,commands,DBM): |
|
|
|
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 |
|
|
|
|
|
|
|
#对节点指令进行校验 |
|
|
|
bok,strerror = g_CV.verify_node_cmds(node_cmds) |
|
|
|
if not bok: #节点指令存在问题,则不进行后续处理,提交一个错误反馈任务 |
|
|
|
# 提交llm待处理任务 |
|
|
|
node.step_num += 1 #单步次数还原一个--暂时只针对有效的返回才算一次 |
|
|
|
self.put_node_reslist(node, strerror, 2) |
|
|
|
return False,commands,0 |
|
|
|
#message_调整传递时机后,可以先执行添加节点 |
|
|
|
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:#如果有没有添加的节点,默认在当前节点下添加 -- 一般不会有,还没遇到 |
|
|
@ -699,7 +886,7 @@ class TaskObject: |
|
|
|
action = node_json["action"] |
|
|
|
if action == "add_node": # 新增节点 |
|
|
|
if node.cur_layer >= self.max_layer: |
|
|
|
continue #节点层级达到上限后不允许再添加子节点-- 平级的一样 |
|
|
|
continue # 节点层级达到上限后不允许再添加子节点-- 平级的一样 |
|
|
|
parent_node_name = node_json["parent"] |
|
|
|
# status = "未完成" #2025-4-11修改MGS-节点指令格式,取消了status |
|
|
|
add_node_names = node_json["nodes"].replace(',', ',').split(',') |
|
|
@ -707,8 +894,9 @@ class TaskObject: |
|
|
|
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): # 添加当前节点的平级节点 |
|
|
|
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("遇到一次添加平级节点") |
|
|
@ -721,46 +909,48 @@ class TaskObject: |
|
|
|
self.logger.debug("遇到一次添加子节点的子节点") |
|
|
|
break |
|
|
|
if not badd: |
|
|
|
self.logger.error(f"添加子节点失败!父节点不是当前节点,不是当前节点的父节点,不是当前节点的子节点,需要介入!!{node_json}---当前节点为:{node.path}") # 丢弃该节点 |
|
|
|
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就可以不用了 |
|
|
|
# 2025-5-12 是否采用本地校验节点是否有指令,如果使用,则no_instruction就可以不用了 |
|
|
|
for add_node in all_add_node: |
|
|
|
bcommand = False |
|
|
|
for com in commands: |
|
|
|
if add_node in com: |
|
|
|
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): #正常应该是当前节点漏洞信息--暂时只考虑只会有一个漏洞 |
|
|
|
if node.name == node_name or node_name.endswith(node.name): # 正常应该是当前节点漏洞信息--暂时只考虑只会有一个漏洞 |
|
|
|
vul_node = node |
|
|
|
else: #匹配子节点 |
|
|
|
else: # 匹配子节点 |
|
|
|
for child in node.children: |
|
|
|
if child.name == node_name or node_name.endswith(child.name): |
|
|
|
vul_node = node |
|
|
|
break |
|
|
|
if vul_node: #找到对应了漏洞节点 |
|
|
|
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, |
|
|
|
# 保存到数据库 --- 数据库有记录多个,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("漏洞信息错误") |
|
|
@ -775,65 +965,17 @@ class TaskObject: |
|
|
|
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 get_other_instruction(self,nodes,DBM,cur_node): |
|
|
|
res_str = ','.join(nodes) |
|
|
|
new_commands = [] |
|
|
|
no_instr_nodes = nodes |
|
|
|
ierror = 0 |
|
|
|
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.这些节点的父节点为当前节点,请正确生成这些节点的节点路径; |
|
|
|
''' |
|
|
|
fake_prompt = self.DataFilter.filter_prompt(user_Prompt) |
|
|
|
#正常不应该会有node_cmds |
|
|
|
bsuccess,node_cmds, commands, reasoning_content, content = self.LLM.get_llm_instruction(fake_prompt, |
|
|
|
cur_node,self.DataFilter) # 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 |
|
|
|
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: |
|
|
|
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 |
|
|
|
return True, commands, len(add_node_names) |
|
|
|
|
|
|
|
#-----------------任务的启停-------------------- |
|
|
|
def init_task(self,task_id,attack_tree = None): |
|
|
@ -849,11 +991,11 @@ class TaskObject: |
|
|
|
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 |
|
|
|
# 插入一个user消息 |
|
|
|
# 提交第一个llm任务,开始工作 |
|
|
|
know_info = f"本测试主机的IP地址为:{self.local_ip}" |
|
|
|
self.put_node_reslist(root_node, know_info, 0) # 入待提交list,若是人工模式则不入待办MQ |
|
|
|
#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) #新增一个-1的状态值 |
|
|
|
|
|
|
|
# 初始保存个attack_tree文件 |
|
|
|
g_PKM.WriteData(self.attack_tree, str(self.task_id)) |
|
|
|
|
|
|
|