diff --git a/.idea/deployment.xml b/.idea/deployment.xml index 12a800f..6f5d31a 100644 --- a/.idea/deployment.xml +++ b/.idea/deployment.xml @@ -1,6 +1,6 @@ - + @@ -16,6 +16,13 @@ + + + + + + + diff --git a/.idea/misc.xml b/.idea/misc.xml index 9c8bfb5..ec4fd51 100644 --- a/.idea/misc.xml +++ b/.idea/misc.xml @@ -1,4 +1,4 @@ - + \ No newline at end of file diff --git a/.idea/zf_safe.iml b/.idea/zf_safe.iml index 186aec3..30b1a51 100644 --- a/.idea/zf_safe.iml +++ b/.idea/zf_safe.iml @@ -2,7 +2,7 @@ - + \ No newline at end of file diff --git a/config.yaml b/config.yaml index dffaf28..869517d 100644 --- a/config.yaml +++ b/config.yaml @@ -1,5 +1,5 @@ #工作模式 -App_Work_type: 0 #开发模式，只允许单步模式 +App_Work_type: 0 #0-开发模式，只允许单步模式 1-生产模式 包裹下的逻辑可以正常执行 #线程休眠的时间 sleep_time: 20 @@ -22,7 +22,7 @@ mysql: LLM_max_chain_count: 10 #为了避免推理链过长，造成推理效果变差，应该控制一个推理链的长度上限 #Node -max_do_sn: 10 #同一节点最多执行5次指令 +max_do_sn: 15 #同一节点最多执行5次指令 #用户初始密码 pw: zfkj_123!@# diff --git a/mycode/AttackMap.py b/mycode/AttackMap.py index 1f8c9f2..087ce45 100644 --- a/mycode/AttackMap.py +++ b/mycode/AttackMap.py @@ -233,20 +233,20 @@ class TreeNode: if pp_node_name in content: #节点名称在内容中就代表跟该节点相关 bfind = True self.parent_messages.append(msg) - else: + else:#当二级父节点出行后，后面数据应该都是二级内的数据 self.parent_messages.append(msg) #cmsg --取最后两轮 if len(tmp_cmsg) <=4: #cmsg全收 self.parent_messages.extend(tmp_cmsg) else: isart = len(tmp_cmsg) - 4 #正常应该都是两个两个 - if isart/2 !=0: + if isart % 2 != 0: print("c_msg数量不对称，需要检查逻辑！") for msg in tmp_cmsg[isart:]: self.parent_messages.append(msg) #添加子节点 - def add_child(self, child_node,messages=None): + def add_child(self, child_node): child_node.parent = self child_node.path = self.path + f"->{child_node.name}" #子节点的路径赋值 self.children.append(child_node) @@ -266,7 +266,7 @@ class TreeNode: #-------后期扩充逻辑，目前wokr_status的修改交给上层类对象------- def add_instr(self,instr,p_msg,c_msg): - if not self.cur_messages: #为空时赋值 + if not self.parent_messages: #为空时赋值 self.copy_messages(p_msg,c_msg) self._instr_queue.append(instr) diff --git a/mycode/LLMManager.py b/mycode/LLMManager.py index b43d634..12033f7 100644 --- a/mycode/LLMManager.py +++ b/mycode/LLMManager.py @@ -8,6 +8,7 @@ import json import re import os from openai import OpenAI +from openai import OpenAIError, APIConnectionError, APITimeoutError from myutils.ConfigManager import myCongif from myutils.MyTime import get_local_timestr from myutils.MyLogger_logger import LogHandler @@ -64,17 +65,18 @@ class LLMManager: 你是一位渗透测试专家，来指导本地程序进行渗透测试，由你负责动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步测试指令，推动测试前进，直至完成渗透测试。 **总体要求** 1.以测试目标为根节点，结合信息收集和测试反馈的结果，以新的测试点作为子节点，逐步规划和推进下一步测试，形成树型结构（测试树）； -2.每次规划测试指令时，只关注当前节点的测试推进，若涉及新增节点，原则上只在当前节点下新增，避免测试路径的混乱； -3.只有当完成当前节点的测试（信息收集或者漏洞验证）后，再结合返回的结果，完整规划下一步的测试点(新增子节点并生成测试指令)，若没有新的测试点，结束当前节点即可； -4.若一次性新增的节点过多，无法为每个节点都匹配测试指令，请优先保障新增测试节点的完整性，若有未生成测试指令的节点，必须返回未生成指令节点列表； +2.每次规划测试指令时，只关注当前节点的测试推进，若涉及新增节点，原则上只应在当前节点下新增，避免测试路径的混乱； +3.只有当收到当前节点的所有测试指令的结果，且没有新的测试指令需要执行时，再判断是否需要新增子节点进一步进行验证测试，若没有，则结束该路径的验证； +4.若一次性新增的节点过多，无法为每个节点都匹配测试指令，请优先保障新增测试节点的完整性，若有新增的节点未能匹配测试指令，必须返回未匹配指令的节点列表； 5.生成的指令有两类：节点指令和测试指令，指令之间必须以空行间隔，不能包含注释和说明； -6.本地程序会执行生成的指令，但不具备分析和判断能力，只会把执行结果返回提交，执行结果应尽量规避无效的信息； +6.本地程序会执行生成的指令，但不具备分析判断和保持会话能力，只会把执行结果返回提交； 7.若无需要处理的节点数据，节点指令可以不生成； +8.只有当漏洞验证成功，才更新该节点的漏洞信息； 8.若节点已完成测试，测试指令可以不生成。 **测试指令生成准则** 1.测试指令必须对应已有节点，或同时生成新增节点指令； 2.优先使用覆盖面广成功率高的指令；不要生成重复的指令； -3.若需要前后执行的指令，请生成对应的python指令； +3.若需要多条指令配合测试，请生成对应的python指令，完成闭环返回； 4.需要避免用户交互，必须要能返回。 **节点指令格式** - 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1，节点2\"}； @@ -87,7 +89,7 @@ class LLMManager: - [节点路径]为从根节点到目标节点的完整层级路径； **核心要求** - 指令之间必须要有一个空行； -- 需确保测试指令的节点路径和指令的目标节点一致； +- 需确保测试指令的节点路径和指令的目标节点一致,例如：针对子节点的测试指令，节点路径不能指向当前节点； - 根据反馈信息，测试目标有可能产生高危漏洞的，必须新增节点，并提供测试指令； **响应示例** {\"action\":\"add_node\", \"parent\": \"192.168.1.100\", \"nodes\": \"3306端口,22端口\"} @@ -111,49 +113,59 @@ mysql -u root -p 192.168.1.100 sendmessage = [] sendmessage.extend(node.parent_messages) sendmessage.extend(node.cur_messages) - #提交LLM - post_time = get_local_timestr() + #准备请求参数 + params = { + "model": self.model, + "messages": sendmessage, + } + # 某些模型额外的参数 if self.model == "o3-mini-2025-01-31": - response = self.client.chat.completions.create( - model=self.model, - reasoning_effort="high", - messages = sendmessage - ) - else: - response = self.client.chat.completions.create( - model=self.model, - messages= sendmessage - ) + params["reasoning_effort"] = "high" + + try: + # 调用 API + response = self.client.chat.completions.create(**params) + except APITimeoutError: + self.logger.error("OpenAI API 请求超时") + return False, "","","", f"调用超时（model={self.model})" + except APIConnectionError as e: + self.logger.error(f"网络连接错误: {e}") + return False, "","", "", "网络连接错误" + except OpenAIError as e: + # 包括 400/401/403/500 等各种 API 错误 + self.logger.error(f"OpenAI API 错误: {e}") + return False, "","", "", f"API错误: {e}" + except Exception as e: + # 兜底，防止意外 + self.logger.exception("调用 LLM 时出现未预期异常") + return False, "","", "", f"未知错误: {e}" #LLM返回结果处理 + choice = response.choices[0].message reasoning_content = "" content = "" #LLM返回处理 if self.model == "deepseek-reasoner": - #返回错误码：DS-https://api-docs.deepseek.com/zh-cn/quick_start/error_codes - reasoning_content = response.choices[0].message.reasoning_content #推理过程 - #print(reasoning_content) - content = response.choices[0].message.content #推理内容 - print(content) + reasoning_content = getattr(choice, "reasoning_content", "") + content = choice.content # 记录llm历史信息 node.cur_messages.append({'role': 'assistant', 'content': content}) elif self.model == "deepseek-chat": - content = response.choices[0].message - # 记录llm历史信息 + content = choice node.cur_messages.append(content) elif self.model == "o3-mini-2025-01-31": - reasoning_content = "" #gpt不返回推理内容 - content = response.choices[0].message.content - print(content) + content = choice.content # 记录llm历史信息 node.cur_messages.append({'role': 'assistant', 'content': content}) else: self.logger.error("处理到未预设的模型！") - return "","","","","" + return False,"","","","处理到未预设的模型！" + + print(content) #按格式规定对指令进行提取 node_cmds,commands = self.fetch_instruction(content) - return node_cmds,commands,reasoning_content, content, post_time + return True,node_cmds,commands,reasoning_content, content def fetch_instruction(self,response_text): ''' diff --git a/mycode/PythoncodeTool.py b/mycode/PythoncodeTool.py index 8128638..316d056 100644 --- a/mycode/PythoncodeTool.py +++ b/mycode/PythoncodeTool.py @@ -24,6 +24,14 @@ import random import tempfile import multiprocessing import textwrap +import smb +import pexpect +from itertools import product +from socket import create_connection +from cryptography import x509 +from cryptography.hazmat.backends import default_backend +from pymetasploit3.msfrpc import MsfRpcClient +from time import sleep from base64 import b64encode from datetime import datetime,timedelta from mycode.Result_merge import my_merge @@ -52,7 +60,7 @@ def _execute_dynamic(instruction_str): 'set': set, 'str': str, 'sum': sum, 'type': type, 'open': open, 'Exception': Exception, 'locals': locals, 'ConnectionResetError':ConnectionResetError,'BrokenPipeError':BrokenPipeError, - 'bytes':bytes, + 'bytes':bytes,'tuple':tuple, } # 构造安全的 globals safe_globals = { @@ -83,7 +91,15 @@ def _execute_dynamic(instruction_str): 'os':os, 'datetime':datetime, 'timedelta':timedelta, - 'b64encode':b64encode + 'b64encode':b64encode, + 'smb':smb, + 'pexpect':pexpect, + 'sleep':sleep, + 'MsfRpcClient':MsfRpcClient, + 'x509':x509, + 'default_backend':default_backend, + 'product':product, + 'create_connection':create_connection } safe_locals = {} try: @@ -176,9 +192,7 @@ class PythoncodeTool(): HIGH_RISK = { 'eval', # eval(...) 'exec', # exec(...) - 'os.system', # os.system(...) 'subprocess.call', # subprocess.call(...) - 'subprocess.Popen' # subprocess.Popen(...) } try: tree = ast.parse(code) diff --git a/mycode/TaskManager.py b/mycode/TaskManager.py index d060622..d9f4636 100644 --- a/mycode/TaskManager.py +++ b/mycode/TaskManager.py @@ -193,6 +193,9 @@ class TaskManager: if node: work_status = node.get_work_status() if work_status == 0 or work_status == 3: + if work_status == 0: + if not node.parent_messages: #如果messages为空--且不会是根节点 + node.copy_messages(node.parent.parent_messages,node.parent.cur_messages) bsuccess,error = node.updatemsg(newtype,newcontent,0) #取的第一条，也就修改第一条 return bsuccess,error else: diff --git a/mycode/TaskObject.py b/mycode/TaskObject.py index f864a29..9d50261 100644 --- a/mycode/TaskObject.py +++ b/mycode/TaskObject.py @@ -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 diff --git a/pipfile b/pipfile index 372e846..98beab2 100644 --- a/pipfile +++ b/pipfile @@ -1,29 +1,33 @@ # -i https://pypi.tuna.tsinghua.edu.cn/simple/ -pip install openai -pip install mysql-connector-python -pip install pymysql +pip install openai -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install mysql-connector-python -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install pymysql -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install pysmb -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install msgpack -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install pymetasploit3 -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install cryptography -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install loguru -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install paramiko -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install impacket -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install beautifulsoup4 -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install cryptography -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install psycopg2 -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install dirsearch -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install pexpect -i https://pypi.tuna.tsinghua.edu.cn/simple/ + + apt install sublist3r apt install gobuster apt install jq +apt install libpq-dev python3-dev +apt install sshpass + #smuggler git clone https://github.com/defparam/smuggler.git -pip install beautifulsoup4 -pip install cryptography -pip install loguru -i https://pypi.tuna.tsinghua.edu.cn/simple/ -pip install paramiko -i https://pypi.tuna.tsinghua.edu.cn/simple/ -pip install impacket -i https://pypi.tuna.tsinghua.edu.cn/simple/ - -sudo apt-get install libpq-dev python3-dev -pip install psycopg2 -i https://pypi.tuna.tsinghua.edu.cn/simple/ - cd /usr/share/wordlists/ gzip -d rockyou.txt.gz -pip install dirsearch -i https://pypi.tuna.tsinghua.edu.cn/simple/ - -apt install sshpass - #----gittools---- # 克隆 GitTools 仓库 git clone https://github.com/internetwache/GitTools.git diff --git a/test.py b/test.py index 3630ea3..94e329f 100644 --- a/test.py +++ b/test.py @@ -76,6 +76,15 @@ class Mytest: print(cur.fetchall()) # 应该显示 ('character_set_client', 'utf8') cnx.close() + def tmp_test(self): + list_a = [0,1,2,3,4,5,6,7,8,9] + + isart = len(list_a) - 4 # 正常应该都是两个两个 + if isart % 2 != 0: + print("c_msg数量不对称，需要检查逻辑！") + for msg in list_a[isart:]: + print(msg) + if __name__ == "__main__": # 示例使用 @@ -92,8 +101,35 @@ if __name__ == "__main__": mytest.dynamic_fun() elif test_type == 1: # # 获取所有自定义函数详情 HIGH_RISK_FUNCTIONS = ['eval', 'exec', 'os.system', 'subprocess.call', 'subprocess.Popen'] - str_instr = ''' -hydra -L /mnt/zfsafe/tools/../payload/users -P /mnt/zfsafe/tools/../payload/passwords -t 6 -f -I -s 5900 -e ns 192.168.204.137 vnc -o hydra_result.txt + str_instr = '''python-code + +import ssl +from socket import create_connection + +def dynamic_fun(): + try: + # 强制使用CBC模式弱加密套件 + context = ssl.SSLContext(ssl.PROTOCOL_TLSv1) + context.set_ciphers('AES128-SHA') + + # 构造异常填充测试数据 + sock = create_connection(('58.216.217.70', 443)) + ssock = context.wrap_socket(sock, server_hostname='58.216.217.70') + + # 发送包含异常填充的测试请求 + ssock.send(b"GET / HTTP/1.1\\r\\nHost: 58.216.217.70\\r\\n" + b"Cookie: test=AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA\\r\\n\\r\\n") + response = ssock.recv(2048) + + # 检测异常响应模式 + if b"HTTP/1.1 200 OK" in response: + return (True, "服务器接受异常填充数据") + return (False, "未检测到典型漏洞特征") + + except ssl.SSLError as e: + return (False, f"加密错误: {repr(e)}") + except Exception as e: + return (False, f"验证失败: {str(e)}") ''' #str_instr = str_instr.strip() + " --max-time 10" dedented_code = textwrap.dedent(str_instr.strip()) @@ -126,7 +162,7 @@ hydra -L /mnt/zfsafe/tools/../payload/users -P /mnt/zfsafe/tools/../payload/pass g_PKM.WriteData(task.attack_tree,str(task.task_id)) elif test_type ==3: #测试指令入节点 strinstr = ''' -```python-[目标系统->192.168.204.137->端口扫描->80-HTTP服务检测->目录扫描->phpMyAdmin访问测试->默认凭证登录尝试->常用凭证爆破->字典暴力破解->版本漏洞检测] import requests def dynamic_fun(): try: # 检测phpMyAdmin版本信息 r = requests.get('http://192.168.204.137/phpMyAdmin/README', timeout=5) if 'phpMyAdmin' in r.text and 'Version' in r.text: return (True, f"版本信息泄露:{r.text.split('Version')[1].split('\\n')[0].strip()}") # 检测ChangeLog文件泄露 r = requests.get('http://192.168.204.137/phpMyAdmin/ChangeLog', timeout=5) if 'phpMyAdmin ChangeLog' in r.text: return (True, "存在ChangeLog文件泄露风险") return (True, "未获取到有效版本信息") except Exception as e: return (False, f"版本检测异常:{str(e)}") ``` +) ''' strNodes = "执行系统命令探测，权限提升尝试，横向移动测试" nodes = strNodes.split(', ') @@ -135,21 +171,23 @@ hydra -L /mnt/zfsafe/tools/../payload/users -P /mnt/zfsafe/tools/../payload/pass print(node_name) elif test_type == 4: # 修改Messages - attact_tree = g_PKM.ReadData("6") + attact_tree = g_PKM.ReadData("27") # 创建一个新的节点 from mycode.AttackMap import TreeNode testnode = TreeNode("test", 0) LLM.build_initial_prompt(testnode) # 新的Message - systems = testnode.messages[0]["content"] + systems = testnode.parent_messages[0]["content"] # print(systems) # 遍历node，查看有instr的ndoe nodes = attact_tree.traverse_bfs() for node in nodes: - node.messages[0]["content"] = systems - g_PKM.WriteData(attact_tree, "6") + node.parent_messages[0]["content"] = systems + g_PKM.WriteData(attact_tree, "27") print("完成Messgae更新") elif test_type ==5: mytest.do_test() + elif test_type == 6: + mytest.tmp_test() else: pass diff --git a/tools/CurlTool.py b/tools/CurlTool.py index 15b0af1..5a2a03b 100644 --- a/tools/CurlTool.py +++ b/tools/CurlTool.py @@ -1,7 +1,8 @@ import requests -import shlex +import pexpect import re import json +import subprocess from bs4 import BeautifulSoup from tools.ToolBase import ToolBase @@ -13,12 +14,12 @@ class CurlTool(ToolBase): # self.verify_ssl = True def get_time_out(self): - return 60*5 + return 60 def validate_instruction(self, instruction_old): #instruction = instruction_old #指令过滤 - timeout = 0 #curl指令遇到不返回的情况 curl --path-as-is -i http://192.168.204.137:8180/webapps/../conf/tomcat-users.xml + timeout = self.get_time_out() #curl指令遇到不返回的情况 curl --path-as-is -i http://192.168.204.137:8180/webapps/../conf/tomcat-users.xml #添加-i 返回信息头 if 'base64 -d' in instruction_old: return instruction_old @@ -44,6 +45,77 @@ class CurlTool(ToolBase): final_instruction = ' '.join(curl_parts) return final_instruction, timeout + def do_worker_pexpect(self,str_instruction,timeout,ext_params): + try: + result = "" + exc_do = pexpect.spawn('bash',['-c',str_instruction],timeout=timeout)#spawn 第一个参数是可执行文件 + while True: + index = exc_do.expect([ + pexpect.TIMEOUT, + pexpect.EOF, + ]) + try: + response = exc_do.before.decode('utf-8', errors='replace') + except Exception as e: + response = f"无法解码响应: {str(e)}" + result += response + + if index == 0: + result += f"\n执行超时{timeout}秒" + break + elif index == 1: + break + else: + print("遇到其他输出！") + break + print(result) + return result + except Exception as e: + return f"执行错误: {str(e)}" + + def do_worker_subprocess(self,str_instruction,timeout,ext_params): + try: + # 执行命令，捕获输出为字节形式 + if timeout: + result = subprocess.run(str_instruction, shell=True,timeout=timeout, capture_output=True) + else: + result = subprocess.run(str_instruction, shell=True, capture_output=True) + + if result.returncode == 0: + # 命令成功，处理 stdout + try: + response = result.stdout.decode('utf-8', errors='replace') + except Exception as e: + response = f"无法解码响应: {str(e)}" + return response + else: + # 命令失败，处理 stderr + try: + error = result.stderr.decode('utf-8', errors='replace') + except Exception as e: + error = f"无法解码错误信息: {str(e)}" + return error + + except Exception as e: + return (False, f"命令执行失败: {str(e)}") + + def execute_instruction(self, instruction_old): + ext_params = self.create_extparams() + # 第一步：验证指令合法性 + instruction,time_out = self.validate_instruction(instruction_old) + if not instruction: + return False, instruction_old, "该指令暂不执行！","",ext_params + # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ + + # 第二步：执行指令---需要对ftp指令进行区分判断 + #output = self.do_worker_pexpect(instruction, time_out, ext_params) + output = self.do_worker_subprocess(instruction, time_out, ext_params) + + # 第三步：分析执行结果 + analysis = self.analyze_result(output,instruction,"","") + + return True, instruction, analysis,output,ext_params + def get_ssl_info(self,stderr,stdout): # -------------------------- # 解释信息的安全意义： diff --git a/tools/FtpTool.py b/tools/FtpTool.py index 613d8b0..bd85d94 100644 --- a/tools/FtpTool.py +++ b/tools/FtpTool.py @@ -5,6 +5,7 @@ import os import ipaddress import subprocess import tempfile +import pexpect from tools.ToolBase import ToolBase class FtpTool(ToolBase): @@ -56,33 +57,40 @@ class FtpTool(ToolBase): return new_instr,timeout - def do_worker_subprocess(self,str_instruction,timeout,ext_params): - output = "" - stdout = "" - stderr = "" + def do_worker_pexpect(self,str_instruction,timeout,ext_params): try: - if timeout == 0: - result = subprocess.run(str_instruction, shell=True, capture_output=True, text=True) - elif timeout > 0: - result = subprocess.run(str_instruction, shell=True, capture_output=True, text=True, timeout=timeout) - else: - print("timeout参数错误") - stderr = result.stderr - stdout = result.stdout - except subprocess.TimeoutExpired as e: - stdout = e.stdout if e.stdout is not None else "" - stderr = e.stderr if e.stderr is not None else "" - ext_params.is_user = True # 对于超时的也需要人工进行确认，是否是预期的超时 + result = "" + exc_do = pexpect.spawn('bash',['-c',str_instruction],timeout=timeout,encoding='utf-8')#spawn 第一个参数是可执行文件 + while True: + index = exc_do.expect([ + pexpect.TIMEOUT, + pexpect.EOF, + '\):', + 'Password:', + 'ftp>' + ]) + result += str(exc_do.before) + if index == 0: + result += f"\n执行超时{timeout}秒" + break + elif index == 1: + break + elif index==2: + result += "):\n" + exc_do.sendline('') # 输入空密码后不知道会有多少种情况，密码不对，密码对 + continue + elif index ==3:#针对要输入密码的情况，暂时智能输入个空字符 + result += "Password:\n" + exc_do.sendline('') # 输入空密码后不知道会有多少种情况，密码不对，密码对 + continue + elif index == 4: + break + else: + print("遇到其他输出！") + break + return result except Exception as e: - ext_params.is_user = True - return False, str_instruction, f"执行失败：{str(e)}", "", ext_params # 执行失败，提交给人工确认指令的正确性 - - output = stdout - if stderr: - output += stderr - if isinstance(output, bytes): # 若是bytes则转成str - output = output.decode('utf-8', errors='ignore') - return output + return f"执行错误: {str(e)}" def do_worker_script(self,str_instruction,timeout,ext_params): # 创建临时文件保存输出 @@ -130,7 +138,8 @@ class FtpTool(ToolBase): # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ # 第二步：执行指令---需要对ftp指令进行区分判断 - output = self.do_worker_script(instruction, time_out, ext_params) + #output = self.do_worker_script(instruction, time_out, ext_params) + output = self.do_worker_pexpect(instruction, time_out, ext_params) # 第三步：分析执行结果 analysis = self.analyze_result(output,instruction,"","") diff --git a/tools/NcTool.py b/tools/NcTool.py index 0b69d32..d3c651b 100644 --- a/tools/NcTool.py +++ b/tools/NcTool.py @@ -1,3 +1,4 @@ +import pexpect from tools.ToolBase import ToolBase class NcTool(ToolBase): @@ -8,6 +9,50 @@ class NcTool(ToolBase): instruction = f"bash -c \"{instruction}\"" return instruction,timeout + def do_worker_pexpect(self,str_instruction,timeout,ext_params): + try: + result = "" + exc_do = pexpect.spawn('bash',['-c',str_instruction],timeout=timeout,encoding='utf-8')#spawn 第一个参数是可执行文件 + while True: + index = exc_do.expect([ + pexpect.TIMEOUT, + pexpect.EOF, + 'Password:' + ]) + result += str(exc_do.before) + if index == 0: + result += f"\n执行超时{timeout}秒" + break + elif index == 1: + break + elif index == 2: + result += "Password:\n" + exc_do.sendline('') # 输入空密码后不知道会有多少种情况，密码不对，密码对 + continue + else: + print("遇到其他输出！") + break + return result + except Exception as e: + return f"执行错误: {str(e)}" + + def execute_instruction(self, instruction_old): + ext_params = self.create_extparams() + # 第一步：验证指令合法性 + instruction,time_out = self.validate_instruction(instruction_old) + if not instruction: + return False, instruction_old, "该指令暂不执行！","",ext_params + # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ + + # 第二步：执行指令---需要对ftp指令进行区分判断 + #output = self.do_worker_script(instruction, time_out, ext_params) + output = self.do_worker_pexpect(instruction, time_out, ext_params) + + # 第三步：分析执行结果 + analysis = self.analyze_result(output,instruction,"","") + + return True, instruction, analysis,output,ext_params + def analyze_result(self, result,instruction,stderr,stdout): #指令结果分析 diff --git a/tools/NmapTool.py b/tools/NmapTool.py index 45c3a55..f3276f0 100644 --- a/tools/NmapTool.py +++ b/tools/NmapTool.py @@ -1,6 +1,7 @@ # Nmap工具类 import re import json +import pexpect from typing import List, Dict, Any from tools.ToolBase import ToolBase @@ -8,11 +9,48 @@ from tools.ToolBase import ToolBase class NmapTool(ToolBase): def validate_instruction(self, instruction): #nmap过滤 - timeout = 0 + timeout = 60*15 if "&& nc " in instruction or "&& ftp " in instruction: timeout = 60 return instruction,timeout + def do_worker_pexpect(self,str_instruction,timeout,ext_params): + try: + result = "" + exc_do = pexpect.spawn('bash',['-c',str_instruction],timeout=timeout,encoding='utf-8')#spawn 第一个参数是可执行文件 + index = exc_do.expect([ + pexpect.TIMEOUT, + pexpect.EOF + ]) + result += str(exc_do.before) + if index == 0: + result += f"\n执行超时{timeout}秒" + elif index ==1: + pass + else: + print("遇到其他输出！") + pass + return result + except Exception as e: + return f"执行错误: {str(e)}" + + def execute_instruction(self, instruction_old): + ext_params = self.create_extparams() + # 第一步：验证指令合法性 + instruction,time_out = self.validate_instruction(instruction_old) + if not instruction: + return False, instruction_old, "该指令暂不执行！","",ext_params + # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ + + # 第二步：执行指令---需要对ftp指令进行区分判断 + #output = self.do_worker_script(instruction, time_out, ext_params) + output = self.do_worker_pexpect(instruction, time_out, ext_params) + + # 第三步：分析执行结果 + analysis = self.analyze_result(output,instruction,"","") + + return True, instruction, analysis,output,ext_params + def parse_nmap_output(self,nmap_output: str) -> Dict[str, Any]: """ 分析 nmap 扫描输出，提取常见信息： diff --git a/tools/OpensslTool.py b/tools/OpensslTool.py index 741d1ce..cd1c691 100644 --- a/tools/OpensslTool.py +++ b/tools/OpensslTool.py @@ -4,13 +4,53 @@ from cryptography.hazmat.backends import default_backend from cryptography.x509 import NameOID import re import json +import pexpect class OpensslTool(ToolBase): def validate_instruction(self, instruction): #指令过滤 - timeout = 0 + timeout = 60*5 return instruction,timeout + def do_worker_pexpect(self, str_instruction, timeout, ext_params): + try: + result = "" + exc_do = pexpect.spawn('bash', ['-c', str_instruction], timeout=timeout, + encoding='utf-8') # spawn 第一个参数是可执行文件 + index = exc_do.expect([ + pexpect.TIMEOUT, + pexpect.EOF + ]) + result += str(exc_do.before) + if index == 0: + result += f"\n执行超时{timeout}秒" + elif index == 1: + pass + else: + print("遇到其他输出！") + pass + return result + except Exception as e: + return f"执行错误: {str(e)}" + + def execute_instruction(self, instruction_old): + ext_params = self.create_extparams() + # 第一步：验证指令合法性 + instruction,time_out = self.validate_instruction(instruction_old) + if not instruction: + return False, instruction_old, "该指令暂不执行！","",ext_params + # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ + + # 第二步：执行指令---需要对ftp指令进行区分判断 + #output = self.do_worker_script(instruction, time_out, ext_params) + output = self.do_worker_pexpect(instruction, time_out, ext_params) + + # 第三步：分析执行结果 + analysis = self.analyze_result(output,instruction,"","") + + return True, instruction, analysis,output,ext_params + + def parse_name(self,name): """解析X509名称对象为结构化字典""" return { @@ -63,6 +103,7 @@ class OpensslTool(ToolBase): def analyze_result(self, result,instruction,stderr,stdout): #指令结果分析 - result = self.parse_ssl_info(stdout) - result = json.dumps(result,ensure_ascii=False) + if len(result) > 3000: + result = self.parse_ssl_info(stdout) + result = json.dumps(result,ensure_ascii=False) return result \ No newline at end of file diff --git a/tools/OtherTool.py b/tools/OtherTool.py index d4dd407..1662425 100644 --- a/tools/OtherTool.py +++ b/tools/OtherTool.py @@ -4,7 +4,7 @@ import os import shlex import subprocess import tempfile - +import pexpect class OtherTool(ToolBase): def validate_instruction(self, instruction): @@ -54,6 +54,27 @@ class OtherTool(ToolBase): pass # 文件可能未创建 return output + def do_worker_pexpect(self, str_instruction, timeout, ext_params): + try: + result = "" + exc_do = pexpect.spawn('bash', ['-c', str_instruction], timeout=timeout, + encoding='utf-8') # spawn 第一个参数是可执行文件 + index = exc_do.expect([ + pexpect.TIMEOUT, + pexpect.EOF + ]) + result += str(exc_do.before) + if index == 0: + result += f"\n执行超时{timeout}秒" + elif index == 1: + pass + else: + print("遇到其他输出！") + pass + return result + except Exception as e: + return f"执行错误: {str(e)}" + def execute_instruction(self, instruction_old): ext_params = self.create_extparams() # 第一步：验证指令合法性 @@ -63,7 +84,8 @@ class OtherTool(ToolBase): # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ # 第二步：执行指令---需要对ftp指令进行区分判断 - output = self.do_worker_script(instruction, time_out, ext_params) + #output = self.do_worker_script(instruction, time_out, ext_params) + output = self.do_worker_pexpect(instruction, time_out, ext_params) # 第三步：分析执行结果 analysis = self.analyze_result(output,instruction,"","") diff --git a/tools/SshTool.py b/tools/SshTool.py index 8fe8fc5..452976d 100644 --- a/tools/SshTool.py +++ b/tools/SshTool.py @@ -3,7 +3,7 @@ from tools.ToolBase import ToolBase class SshTool(ToolBase): def validate_instruction(self, instruction): #指令过滤 - timeout = 0 + timeout = 60 return instruction,timeout def analyze_result(self, result,instruction,stderr,stdout): diff --git a/web/API/system.py b/web/API/system.py index 90c1078..308c37b 100644 --- a/web/API/system.py +++ b/web/API/system.py @@ -1,14 +1,17 @@ from . import api +from web.common.utils import login_required from mycode.DBManager import app_DBM from myutils.ReManager import mReM from quart import Quart, render_template, redirect, url_for, request,jsonify @api.route('/system/getinfo',methods=['GET']) +@login_required async def get_system_info(): data = app_DBM.getsystem_info() return jsonify({"local_ip":data[0],"version":data[1]}) @api.route('/system/updateip',methods=['POST']) +@login_required async def update_local_ip(): data = await request.get_json() local_ip = data.get("local_ip") diff --git a/web/API/task.py b/web/API/task.py index 71d64b9..efbb73d 100644 --- a/web/API/task.py +++ b/web/API/task.py @@ -3,6 +3,7 @@ from quart import Quart, render_template, redirect, url_for, request,jsonify from mycode.TargetManager import g_TM from mycode.DBManager import app_DBM from mycode.TaskManager import g_TaskM +from web.common.utils import login_required def is_valid_target(test_target: str) -> bool: @@ -15,12 +16,15 @@ def is_valid_target(test_target: str) -> bool: return False @api.route('/task/start',methods=['POST']) +@login_required async def start_task(): #开始任务 data = await request.get_json() test_target = data.get("testTarget") cookie_info = data.get("cookieInfo") llm_type = data.get("curmodel") # //0-腾讯云，1-DS，2-2233.ai,3-GPT 目前只有1-2，2025-4-4 work_type = data.get("workType") #0-人工，1-自动 + if llm_type == 2: + return jsonify({"error": "O3余额不足，请更换模型！"}), 400 #新增任务处理 bok,_,_ = g_TM.validate_and_extract(test_target) if not bok: @@ -38,6 +42,7 @@ async def start_task(): #开始任务 return redirect(url_for('main.get_html', html='task_manager.html')) @api.route('/task/taskover',methods=['POST']) +@login_required async def over_task(): data = await request.get_json() task_id = data.get("cur_task_id") @@ -47,6 +52,7 @@ async def over_task(): return jsonify({"bsuccess": bsuccess, "error": error}) @api.route('/task/deltask',methods=['POST']) +@login_required async def del_task(): data = await request.get_json() task_id = data.get("task_id") @@ -57,6 +63,7 @@ async def del_task(): @api.route('/task/getlist',methods=['GET']) +@login_required async def get_task_list(): #task_list = app_DBM.get_task_list() #从内存取--2025-4-6 task_list = g_TaskM.get_task_list() @@ -66,6 +73,7 @@ async def get_task_list(): return jsonify({"error":"查询任务数据出错！"}),500 @api.route('/task/getinstr',methods=['POST']) +@login_required async def get_instr(): data = await request.get_json() task_id = data.get("cur_task_id") @@ -76,6 +84,7 @@ async def get_instr(): return jsonify({"instrs":instrs}) @api.route('/task/getvul',methods=['POST']) +@login_required async def get_vul(): data = await request.get_json() task_id = data.get("cur_task_id") @@ -88,6 +97,7 @@ async def get_vul(): return jsonify({"vuls":vuls}) @api.route('/task/gettree',methods=['POST']) +@login_required async def get_tree(): data = await request.get_json() task_id = data.get("task_id") @@ -97,6 +107,7 @@ async def get_tree(): return jsonify({"tree":tree_dict}) @api.route('/task/gethistree',methods=['POST']) +@login_required async def get_his_tree(): data = await request.get_json() task_id = data.get("task_id") @@ -106,6 +117,7 @@ async def get_his_tree(): return jsonify({"tree":tree_dict}) @api.route('/task/taskcontrol',methods=['POST']) +@login_required async def task_status_control(): '''控制任务状态 1.对于执行时间长的指令，如何处理？强制停止的话，要有个执行中指令的缓存，强制停止该指令返回到待执行，执行完成，该指令到执行完成； @@ -121,6 +133,7 @@ async def task_status_control(): return jsonify({'error': strerror}), 400 @api.route('/task/taskstep',methods=['POST']) +@login_required async def task_one_step(): '''单步推进任务--也就是待处理node 返回bsuccess,error 1.执行单步的前提条件是，工作线程都要在工作； @@ -134,6 +147,7 @@ async def task_one_step(): return jsonify({"bsuccess":bsuccess,"error":error}) @api.route('/task/nodestep',methods=['POST']) +@login_required async def node_one_step(): data = await request.get_json() task_id = data.get("task_id") @@ -144,6 +158,7 @@ async def node_one_step(): return jsonify({"bsuccess":bsuccess,"error":error}) @api.route('/task/taskworktype',methods=['POST']) +@login_required async def task_work_type_control(): data = await request.get_json() task_id = data.get("cur_task_id") @@ -154,6 +169,7 @@ async def task_work_type_control(): return jsonify({"bsuccess": bsuccess}) @api.route('/task/nodecontrol',methods=['POST']) +@login_required async def node_work_status_control(): data = await request.get_json() task_id = data.get("task_id") @@ -167,6 +183,7 @@ async def node_work_status_control(): return jsonify({"newbwork":newbwork}) @api.route('/task/nodegetinstr',methods=['POST']) +@login_required async def node_get_instr(): data = await request.get_json() task_id = data.get("task_id") @@ -179,6 +196,7 @@ async def node_get_instr(): return jsonify({"doneInstrs":doneInstrs,"todoInstrs":todoInstrs}) @api.route('/task/hisnodegetinstr',methods=['POST']) +@login_required async def his_node_get_instr(): data = await request.get_json() task_id = data.get("task_id") @@ -191,6 +209,7 @@ async def his_node_get_instr(): return jsonify({"doneInstrs":doneInstrs}) @api.route('/task/nodegetmsg',methods=['POST']) +@login_required async def node_get_msg(): data = await request.get_json() task_id = data.get("task_id") @@ -201,6 +220,7 @@ async def node_get_msg(): return jsonify({"submitted": submitted, "pending": pending}) @api.route('/task/nodeupdatemsg',methods=['POST']) +@login_required async def node_update_msg(): data = await request.get_json() task_id = data.get("task_id") @@ -213,6 +233,7 @@ async def node_update_msg(): return jsonify({"bsuccess":bsuccess,"error":error}) @api.route('/task/delnodeinstr',methods=['POST']) +@login_required async def node_del_instr(): data = await request.get_json() task_id = data.get("task_id") @@ -225,6 +246,7 @@ async def node_del_instr(): @api.route('/task/histasks',methods=['POST']) +@login_required async def get_his_task(): data = await request.get_json() target_name = data.get("target_name") diff --git a/web/API/user.py b/web/API/user.py index 9b4a928..9154884 100644 --- a/web/API/user.py +++ b/web/API/user.py @@ -76,30 +76,6 @@ async def user_logout(): session.clear() return redirect(url_for('main.login')) -@api.route('/user/adduser',methods=['POST']) -@login_required -async def user_adduser(): #新增用户 - username = (await request.form)['username'] - people = (await request.form)['people'] - tellnum = (await request.form)['tellnum'] - strsql = f"select username from user where username = '{username}';" - password = myCongif.get_data('pw') - data = app_DBM.do_select(strsql) - if data: - reStatus = 0 - reMsg = '用户名重复，请重新输入！' - else: - strsql = (f"INSERT INTO user (username ,password ,status,people,tellnum ) VALUES " - f"('{username}','{password}',1,'{people}','{tellnum}');") - ret = app_DBM.do_sql(strsql) - if ret == True: - reStatus = 1 - reMsg = '添加用户成功' - else: - reStatus = 0 - reMsg = '添加用户异常，请联系管理员处理！' - return jsonify({'status':reStatus,'msg':reMsg}) - @api.route('/user/passwd',methods=['POST']) @login_required async def user_change_passwd(): #修改密码 diff --git a/web/API/wsm.py b/web/API/wsm.py index f824c3b..73551c1 100644 --- a/web/API/wsm.py +++ b/web/API/wsm.py @@ -2,9 +2,11 @@ import json from . import api from quart import Quart, websocket, jsonify from mycode.WebSocketManager import g_WSM +from web.common.utils import login_required # WebSocket 路由，端口默认与 HTTP 同端口，例如 5000（开发时） @api.websocket("/ws") +@login_required async def ws(): """ WebSocket 连接入口： diff --git a/web/main/static/resources/css/node_tree.css b/web/main/static/resources/css/node_tree.css index 0082ce5..ec4e89f 100644 --- a/web/main/static/resources/css/node_tree.css +++ b/web/main/static/resources/css/node_tree.css @@ -99,24 +99,25 @@ transform: translateY(-1px); } .tree-node.selected { - border-color: #1890ff; - background-color: #bae7ff; + border: 2px solid #1890ff !important; + background-color: #bae7ff ; } /* 针对漏洞级别的样式 */ .tree-node.vul-low { border-color: #87d068; - background-color: #f6ffed; + background-color: #f6ffed !important; } .tree-node.vul-medium { border-color: #faad14; - background-color: #fff7e6; + background-color: #fff7e6 !important; } .tree-node.vul-high { border-color: #ff4d4f; - background-color: #fff1f0; + background-color: #fff1f0 !important; } .tree-node.no-work { - border-color: #151515; + border: 2px solid #151515; + /*border-color: #151515;*/ background-color: #c4c4c4; } /* 收缩/展开时隐藏子节点 ul */ diff --git a/web/main/static/resources/scripts/node_tree.js b/web/main/static/resources/scripts/node_tree.js index 838d953..447e0db 100644 --- a/web/main/static/resources/scripts/node_tree.js +++ b/web/main/static/resources/scripts/node_tree.js @@ -33,7 +33,6 @@ } // 根据漏洞级别添加样式 if (nodeData.node_vulgrade) { - nodeSpan.classList.remove("no-work"); if (nodeData.node_vulgrade === "低危") { nodeSpan.classList.add("vul-low"); } else if (nodeData.node_vulgrade === "中危") { @@ -204,7 +203,13 @@ if(node_path === selectedNodeData.node_path){ //只有是当前选中节点才更新数据 selectedNodeData.workstatus = node_workstatus; strnew = getWorkStatus_Str(node_workstatus); - document.getElementById("node_workstatus").textContent = strnew; + work_status_el = document.getElementById("node_workstatus") + work_status_el.textContent = strnew; + work_status_el.classList.toggle('cmd-none',str_workStatus==="无待执行任务"); + work_status_el.classList.toggle('cmd-pending',str_workStatus==="待执行指令中"); + work_status_el.classList.toggle('cmd-running',str_workStatus==="指令执行中"); + work_status_el.classList.toggle('cmd-waiting-llm',str_workStatus==="待提交llm中"); + work_status_el.classList.toggle('cmd-submitting-llm',str_workStatus==="提交llm中"); } } } else { @@ -220,12 +225,13 @@ document.getElementById("node_vulLevel").textContent = vulLevel; str_workStatus = getWorkStatus_Str(Number(workStatus)); - document.getElementById("node_workstatus").textContent = str_workStatus; - document.getElementById("node_workstatus").classList.toggle('cmd-none',str_workStatus==="无待执行任务"); - document.getElementById("node_workstatus").classList.toggle('cmd-pending',str_workStatus==="待执行指令中"); - document.getElementById("node_workstatus").classList.toggle('cmd-running',str_workStatus==="指令执行中"); - document.getElementById("node_workstatus").classList.toggle('cmd-waiting-llm',str_workStatus==="待提交llm中"); - document.getElementById("node_workstatus").classList.toggle('cmd-submitting-llm',str_workStatus==="提交llm中"); + work_status_el = document.getElementById("node_workstatus") + work_status_el.textContent = str_workStatus; + work_status_el.classList.toggle('cmd-none',str_workStatus==="无待执行任务"); + work_status_el.classList.toggle('cmd-pending',str_workStatus==="待执行指令中"); + work_status_el.classList.toggle('cmd-running',str_workStatus==="指令执行中"); + work_status_el.classList.toggle('cmd-waiting-llm',str_workStatus==="待提交llm中"); + work_status_el.classList.toggle('cmd-submitting-llm',str_workStatus==="提交llm中"); if(nodebwork==="true"){ document.getElementById("node_bwork").textContent = "执行中"; diff --git a/web/main/static/resources/scripts/task_modal.js b/web/main/static/resources/scripts/task_modal.js index 9aa9097..105b517 100644 --- a/web/main/static/resources/scripts/task_modal.js +++ b/web/main/static/resources/scripts/task_modal.js @@ -269,12 +269,6 @@ const instrCanvas = new bootstrap.Offcanvas(document.getElementById('instrCanvas')); instrCanvas.show(); - // const modalEl = document.getElementById("instrModal"); - // // 假设用 Bootstrap 5 的 Modal 组件 - // const instrModal = new bootstrap.Modal(modalEl, {keyboard: false}); - // 显示对话框 - //instrModal.show(); - // 在打开 modal 时，先更新提示内容，将 loadingMsg 显示“请稍后，数据获取中…” const loadingMsg = document.getElementById("loadingMsg"); if (loadingMsg) { diff --git a/web/main/templates/assets_manager.html b/web/main/templates/assets_manager.html index 15aa44c..5e77f6b 100644 --- a/web/main/templates/assets_manager.html +++ b/web/main/templates/assets_manager.html @@ -8,7 +8,7 @@ {% block content %} -

功能建设中，在二期实现。。。

+

功能规划中，在二期实现。。。

{% endblock %} diff --git a/web/main/templates/header.html b/web/main/templates/header.html index b859eee..4d9a248 100644 --- a/web/main/templates/header.html +++ b/web/main/templates/header.html @@ -29,7 +29,7 @@ diff --git a/web/main/templates/index.html b/web/main/templates/index.html index 0db168f..b3001fd 100644 --- a/web/main/templates/index.html +++ b/web/main/templates/index.html @@ -85,7 +85,7 @@
- +
****本工具仅限于在得到授权的前提下使用，若目标重要性很高，请使用单步模式，确认测试指令的影响后执行！ +
{% endblock %} diff --git a/web/main/templates/task_manager.html b/web/main/templates/task_manager.html index ff7d866..3b134aa 100644 --- a/web/main/templates/task_manager.html +++ b/web/main/templates/task_manager.html @@ -111,6 +111,21 @@ overflow-y: auto; } +.modal-dialog { + max-width: calc(100vw - 10rem); + margin: 1rem auto; +} +.table { + width: 100%; + table-layout: fixed; +} +.table th, +.table td { + word-wrap: break-word; + white-space: normal; +} + + .disabled-btn { /* 禁用状态样式 */ background-color: #cccccc; /* 灰色背景 */ diff --git a/web/main/templates/task_manager_modal.html b/web/main/templates/task_manager_modal.html index 51b58ab..d9ffbb7 100644 --- a/web/main/templates/task_manager_modal.html +++ b/web/main/templates/task_manager_modal.html @@ -59,19 +59,20 @@ role="tabpanel" aria-labelledby="doneInstrTab" > - +
+
- - - - + + + + -
序号执行指令执行时间执行结果序号执行指令执行时间执行结果
+