diff --git a/mycode/CommandVerify.py b/mycode/CommandVerify.py new file mode 100644 index 0000000..841c884 --- /dev/null +++ b/mycode/CommandVerify.py @@ -0,0 +1,96 @@ +#对llm返回的指令进行校验 +import re +class CommandVerify: + def __init__(self): + pass + + #验证节点指令的结构完整性--主要是判断JSON元素是否完整 + def verify_node_cmds(self,node_cmds): + ''' + 验证节点指令的合规性,持续维护 + :param node_cmds: + :param node: + :return: Flase 存在问题, True 合规 + ''' + strerror = "" + for node_json in node_cmds: + if "action" not in node_json: + 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: + strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点!"} + break + elif action == "update_status": + if "status" not in node_json or "node" not in node_json: + strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点!"} + break + elif action =="no_instruction" or action=="no_create": + if "nodes" not in node_json: + strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点!"} + break + else: + strerror = {"节点指令错误": f"{node_json}不可识别的action值!"} + break + if not strerror: + return True,strerror + return False,strerror + + # 验证节点数据的合规性 + def verify_node_data(self,node_cmds): + add_nodes = [] + no_instr_nodes = [] + for node_cmd in node_cmds: + do_type = node_cmd["action"] + 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("遇到未知节点处理类型") + #核对指令是否有缺失 + 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: + """获取 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: + """获取 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: + """获取 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: + """获取 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: + """集合差集:list_a - list_b""" + return list(set(list_a) - set(list_b)) + + def _difference_b_simple(list_a: list, list_b: list) -> list: + """集合差集:list_b - list_a""" + return list(set(list_b) - set(list_a)) + +g_CV = CommandVerify() \ No newline at end of file diff --git a/mycode/ControlCenter.py b/mycode/ControlCenter.py index 0397190..b01bcde 100644 --- a/mycode/ControlCenter.py +++ b/mycode/ControlCenter.py @@ -33,38 +33,6 @@ class ControlCenter: # ?包括是否对目标进行初始化的信息收集 return {"已知信息":"无"} - def verify_node_cmds(self,node_cmds): - ''' - 验证节点指令的合规性,持续维护 - :param node_cmds: - :param node: - :return: Flase 存在问题, True 合规 - ''' - strerror = "" - for node_json in node_cmds: - if "action" not in node_json: - 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: - strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点!"} - break - elif action == "update_status": - if "status" not in node_json or "node" not in node_json: - strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点!"} - break - elif action =="no_instruction" or action=="no_create": - if "nodes" not in node_json: - strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点!"} - break - else: - strerror = {"节点指令错误": f"{node_json}不可识别的action值!"} - break - if not strerror: - return True,strerror - return False,strerror def restore_one_llm_work(self,node,llm_type,res_list): node.llm_type = llm_type diff --git a/mycode/DBManager.py b/mycode/DBManager.py index 49b167e..7555610 100644 --- a/mycode/DBManager.py +++ b/mycode/DBManager.py @@ -23,7 +23,6 @@ class DBManager: self.passwd = myCongif.get_data('mysql.passwd') self.database = myCongif.get_data('mysql.database') self.connection = None - self.cursor = None elif self.itype ==1: self.dbfile = myCongif.get_data("sqlite") if not os.path.exists(self.dbfile): @@ -35,9 +34,7 @@ class DBManager: def __del__(self): if self.ok: - self.cursor.close() self.connection.close() - self.cursor = None self.connection = None self.logger.debug("DBManager销毁") @@ -46,10 +43,8 @@ class DBManager: if self.itype ==0: self.connection = pymysql.connect(host=self.host, port=self.port, user=self.user, passwd=self.passwd, db=self.database,charset='utf8') - self.cursor = self.connection.cursor() elif self.itype ==1: self.connection = sqlite3.connect(self.dbfile) - self.cursor = self.connection.cursor() self.ok = True self.logger.debug("服务器端数据库连接成功") return True @@ -73,16 +68,17 @@ class DBManager: data = None if self.Retest_conn(): try: - self.cursor.execute(strsql) self.connection.commit() # select要commit提交事务,是存在获取不到最新数据的问题(innoDB事务机制) + with self.connection.cursor() as cursor: + cursor.execute(strsql) + if itype == 1: + data = cursor.fetchone() + else: + data = cursor.fetchall() except Exception as e: self.logger.error("do_select异常报错:%s" % str(e)) self.lock.release() return None - if itype == 1: - data = self.cursor.fetchone() - else: - data = self.cursor.fetchall() self.lock.release() return data @@ -92,47 +88,53 @@ class DBManager: self.lock.acquire() if self.Retest_conn(): try: - # self.conn.begin() - if data: - iret = self.cursor.executemany(strsql, data) #批量执行sql语句 - else: - iret = self.cursor.execute(strsql) - self.connection.commit() - bok = True + with self.connection.cursor() as cursor: + # self.conn.begin() + if data: + iret = cursor.executemany(strsql, data) #批量执行sql语句 + else: + iret = cursor.execute(strsql) + self.connection.commit() + bok = True except Exception as e: self.logger.error("执行数据库语句%s出错:%s" % (strsql, str(e))) self.connection.rollback() self.lock.release() return bok - def safe_do_sql(self,strsql,params): + def safe_do_sql(self,strsql,params,itype=0): bok = False + task_id = 0 self.lock.acquire() if self.Retest_conn(): try: - self.cursor.execute(strsql, params) - self.connection.commit() - bok = True + with self.connection.cursor() as cursor: + cursor.execute(strsql, params) + self.connection.commit() + if itype ==1: + task_id = cursor.lastrowid + bok = True except Exception as e: self.logger.error("执行数据库语句%s出错:%s" % (strsql, str(e))) self.connection.rollback() self.lock.release() - return bok + return bok,task_id def safe_do_select(self,strsql,params,itype=0): results = [] + self.lock.acquire() if self.Retest_conn(): - cursor = self.connection.cursor() #每次查询使用新的游标 ---待验证(用于处理出现一次查询不到数据的问题) + self.connection.commit() try: - cursor.execute(strsql, params) # 执行参数化查询 - if itype ==0: - results = cursor.fetchall() # 获取所有结果 - elif itype ==1: - results = cursor.fetchone() #获得一条记录 + with self.connection.cursor() as cursor: + cursor.execute(strsql, params) # 执行参数化查询 + if itype ==0: + results = cursor.fetchall() # 获取所有结果 + elif itype ==1: + results = cursor.fetchone() #获得一条记录 except Exception as e: print(f"查询出错: {e}") - finally: - cursor.close() + self.lock.release() return results def is_json(self,s:str) -> bool: @@ -174,10 +176,15 @@ class DBManager: sql = "INSERT INTO task (task_name,task_target,start_time,task_status,safe_rank,work_type,cookie_info,llm_type) " \ "VALUES (%s,%s,%s,%s,%s,%s,%s,%s)" params = (test_target,test_target,start_time,1,0,work_type,cookie_info,llm_type) - self.safe_do_sql(sql,params) - task_id = self.cursor.lastrowid + bok,task_id = self.safe_do_sql(sql,params,1) return task_id + def over_task(self,task_id): + strsql = "update task set task_status=2 where ID=%s;" + params = (task_id) + 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 = "" @@ -212,7 +219,8 @@ class DBManager: """ params = (task_id, instruction, str_result, do_sn,start_time,end_time,source_result,ext_params['is_user'], ext_params['is_vulnerability'],node_path) - return self.safe_do_sql(sql,params) + bok,_ = self.safe_do_sql(sql,params) + return bok #llm数据入库 def insert_llm(self,task_id,prompt,reasoning_content,content,post_time,llm_sn,path): @@ -248,7 +256,8 @@ class DBManager: str_content = str_content.encode('utf-8').decode('unicode_escape') params = (task_id,llm_sn,prompt,str_reasoning,str_content,post_time,path) - return self.safe_do_sql(sql,params) + bok,_=self.safe_do_sql(sql,params) + return bok #获取任务的测试指令执行情况 def get_task_instrs(self,task_id,nodename): diff --git a/mycode/LLMManager.py b/mycode/LLMManager.py index 480eb89..ee1a0c6 100644 --- a/mycode/LLMManager.py +++ b/mycode/LLMManager.py @@ -58,37 +58,32 @@ class LLMManager: def build_initial_prompt(self,node): if not node: return - #根节点初始化message + #根节点初始化message----后续有可能需要为每个LLM生成不同的system msg node.messages = [{"role": "system", "content":''' 你是一位渗透测试专家,基于你的专业知识和经验,负责动态控制整个渗透测试过程,根据当前测试状态和返回结果,决定下一步测试指令,推动测试前进,直至完成渗透测试。本地程序会执行你的指令并返回结果。 **总体要求** -1.以目标系统IP为根节点,每个渗透测试点(如端口、服务、漏洞点)作为子节点,形成树型结构(测试树); -2.测试树整体数据由本地程序存储,你只需要关注当前节点的测试推进、状态更新(未完成/已完成)及完整新增子节点; -3.返回两类指令:节点指令和测试指令,以空行间隔,不要包含注释和说明; -4.若一次性新增的节点过多,无法为每个节点都匹配测试指令,请优先保障新增中高危节点的完整性,若有未生成测试指令的节点,必须返回未生成指令节点列表; -5.若无需要处理的节点数据,节点指令可以不生成,但测试指令必须对应已有节点; -**决策流程** -1. 若当前节点是IP且未进行端口扫描,则对当前节点执行端口扫描; -2. 若端口扫描发现开放端口,对可能存在中高危以上风险的端口新增节点并提供测试指令; -3. 若当前节点是端口且未进行服务扫描,则执行服务扫描; -4. 若服务扫描发现服务版本或漏洞,则新增漏洞测试节点并提供测试指令; -5. 若漏洞验证成功,则根据结果决定是否需要进一步测试,若需要进一步测试,则必须将进一步测试的内容作为子节点并提供测试指令; -6. 当当前节点没有新的测试指令时,更新状态为“已完成”。 +1.以测试目标为根节点,每个渗透测试点和方向(如端口、服务、漏洞点)作为子节点,形成树型结构(测试树),层层递进; +2.测试树整体数据由本地程序存储,你只需要关注当前节点的测试推进、状态更新(未完成/已完成)及是否有子节点新增; +3.返回的指令有两类:节点指令和测试指令,指令之间必须以空行间隔,不要包含注释和说明; +4.测试时必须为每个不同的测试点和测试方向,新增节点,同时生成对应的测试指令; +5.若一次性新增的节点过多,无法为每个节点都匹配测试指令,请优先保障新增测试节点的完整性,若有未生成测试指令的节点,必须返回未生成指令节点列表; +6.若漏洞验证成功,则根据结果决定是否需要进一步测试,若需要进一步测试,则为测试内容新增子节点并提供测试指令; +7.当当前节点没有新的测试指令时,更新状态为“已完成”; +8.若无需要处理的节点数据,节点指令可以不生成。 **测试指令生成准则** -1.明确每个测试指令的测试目标,并优先尝试最简单、最直接的办法,不要在同一个请求生成测试效果覆盖的指令; -2.使用递进逻辑组织指令:先尝试基础测试方法,根据执行结果决定是否进行更深入的测试; +1.明确每个测试指令的测试目标,并优先尝试最简单、最直接的办法,不要生成测试效果覆盖的指令; +2.使用递进逻辑组织指令:先尝试基础测试方法,根据执行结果决定是否进行更深入的测试; **节点指令格式** -- 新增节点:{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1,节点2\", \"status\": \"未完成\"}; +- 新增节点:{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1,节点2\"}; - 未生成指令节点列表:{\"action\": \"no_instruction\", \"nodes\": \"节点1,节点2\"}; -- 完成测试未发现漏洞:{\"action\": \"update_status\", \"node\": \"节点\", \"status\": \"已完成\"}; -- 完成测试且发现漏洞:{\"action\": \"update_status\", \"node\": \"节点\", \"status\": \"已完成\",\"vulnerability\": {\"name\":\"漏洞名称\",\"risk\":\"风险等级(低危/中危/高危)\",\"info\":\"补充信息(没有可为空)\"}}; +- 漏洞验证成功:{\"action\": \"find_vul\", \"node\": \"节点\",\"vulnerability\": {\"name\":\"漏洞名称\",\"risk\":\"风险等级(低危/中危/高危)\",\"info\":\"补充信息(没有可为空)\"}}; +- 完成测试:{\"action\": \"end_work\", \"node\": \"节点\"}; **测试指令格式** - shell指令:```bash-[节点路径]指令内容```包裹,需要避免用户交互,若涉及到多步指令,请生成python代码; - python指令:```python-[节点路径]指令内容```包裹,主函数名为dynamic_fun,需包含错误处理,必须返回一个tuple(status, output); -- [节点路径]为从根节点到目标节点的完整层级描述。 +- [节点路径]为从根节点到目标节点的完整层级路径,且需要与指令的目标节点一致。 **核心要求** -- 优先保障新增中高危测试节点的完整性。 - 指令之间必须要有一个空行。 **响应示例** {\"action\":\"add_node\", \"parent\": \"192.168.1.100\", \"nodes\": \"3306端口,22端口\", \"status\": \"未完成\"} diff --git a/mycode/TaskManager.py b/mycode/TaskManager.py index 31a6dee..fa5b3ba 100644 --- a/mycode/TaskManager.py +++ b/mycode/TaskManager.py @@ -67,6 +67,18 @@ class TaskManager: else: return False + def over_task(self,task_id): + task = self.tasks[task_id] + if task: + task.brun = False + #修改数据库数据 + bsuccess = app_DBM.over_task(task_id) + if bsuccess: + del self.tasks[task_id] #删除缓存 + return bsuccess,"" + else: + return False,"没有找到对应的任务" + #控制task启停----线程不停 def control_taks(self,task_id): task = self.tasks[task_id] diff --git a/mycode/TaskObject.py b/mycode/TaskObject.py index a735f27..407e898 100644 --- a/mycode/TaskObject.py +++ b/mycode/TaskObject.py @@ -13,6 +13,7 @@ from myutils.MyLogger_logger import LogHandler from myutils.PickleManager import g_PKM from myutils.ConfigManager import myCongif from mycode.WebSocketManager import g_WSM +from mycode.CommandVerify import g_CV import asyncio import queue import time @@ -345,6 +346,19 @@ class TaskObject: user_Prompt = user_Prompt + ext_Prompt return user_Prompt + #添加子节点 + def add_children_node(self,parent_node,children_names,status="未完成"): + for child_name in children_names: + bfind = False + for node_child in parent_node.children: + if node_child.name == child_name: + bfind = True #有重复的了 + break + if not bfind: + # 添加节点 + new_node = TreeNode(child_name, parent_node.task_id, status) + parent_node.add_child(new_node) # message的传递待验证 + #处理节点指令 def tree_manager(self,node_cmds,node,commands,DBM): '''更新渗透测试树 @@ -355,11 +369,16 @@ class TaskObject: return True,commands #对节点指令进行校验 - bok,strerror = self.CCM.verify_node_cmds(node_cmds) + 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操作 residue_node_cmds = [] @@ -373,30 +392,10 @@ class TaskObject: node_names = node_json["nodes"].split(',') # 新增节点原则上应该都是当前节点增加子节点 if node.name == parent_node_name or parent_node_name.endswith(node.name): #2233ai,节点名称字段会返回整个路径 - for node_name in node_names: - # 判重---遇到过补充未生成指令的节点时,返回了新增这些节点的指令 - bfind = False - for node_child in node.children: - if node_child.name == node_name: - bfind = True - break - if not bfind: - # 添加节点 - new_node = TreeNode(node_name, node.task_id, status) - node.add_child(new_node) # message的传递待验证 + self.add_children_node(node, node_names,status) #添加子节点 elif node.parent.name == parent_node_name or parent_node_name.endswith(node.parent.name): #是添加当前节点的平级节点(当前节点的父节点下添加子节点) --使用2233ai-o3时遇到的情况 - for node_name in node_names: - # 判重---遇到过补充未生成指令的节点时,返回了新增这些节点的指令 - bfind = False - for node_child in node.parent.children: - if node_child.name == node_name: - bfind = True - break - if not bfind: - # 添加节点 - new_node = TreeNode(node_name, node.task_id, status) - node.parent.add_child(new_node) + self.add_children_node(node.parent,node_names,status) #添加当前节点的平级节点 else: self.logger.error(f"添加子节点时,遇到父节点名称不一致的,需要介入!!{node_json}") # 丢弃该节点 else:#其他指令添加到list diff --git a/web/API/task.py b/web/API/task.py index 52434dc..f7f0a88 100644 --- a/web/API/task.py +++ b/web/API/task.py @@ -37,6 +37,16 @@ async def start_task(): #开始任务 #跳转到任务管理页面 return redirect(url_for('main.get_html', html='task_manager.html')) +@api.route('/task/taskover',methods=['POST']) +async def over_task(): + data = await request.get_json() + task_id = data.get("cur_task_id") + if not task_id: + return jsonify({'error': 'Missing task_id'}), 400 + bsuccess,error = g_TaskM.over_task(task_id) + return jsonify({"bsuccess": bsuccess, "error": error}) + + @api.route('/task/getlist',methods=['GET']) async def get_task_list(): #task_list = app_DBM.get_task_list() #从内存取--2025-4-6 diff --git a/web/main/static/resources/scripts/node_tree.js b/web/main/static/resources/scripts/node_tree.js index 2145890..2250190 100644 --- a/web/main/static/resources/scripts/node_tree.js +++ b/web/main/static/resources/scripts/node_tree.js @@ -366,6 +366,8 @@ } } + + //----------------------查看指令modal---------------------------- let doneInstrs = []; // 已执行指令的所有数据 let todoInstrs = []; // 待执行指令的所有数据 diff --git a/web/main/static/resources/scripts/task_manager.js b/web/main/static/resources/scripts/task_manager.js index f92c720..c7d9ba6 100644 --- a/web/main/static/resources/scripts/task_manager.js +++ b/web/main/static/resources/scripts/task_manager.js @@ -237,15 +237,55 @@ async function controlTask(){ actionButton.textContent = "已完成"; } cur_task.taskStatus = newstatus; //光有个cur_taks也可以 - setSetpBtnStatus() + setSetpBtnStatus(); //更新task_list的显示 - updateTaskList() + updateTaskList(); } catch (error) { console.error("控制任务状态异常:", error); alert("控制任务状态异常:",error); } } +//结束任务-brun-false +document.getElementById("btnTaskOver").addEventListener("click",()=>{ + overTask(); +}) +async function overTask(){ + if(cur_task_id === 0){ + alert("请先选择一个任务!") + return + } + try { + if (confirm('确定要结束此任务吗?')){ + const res = await fetch("/api/task/taskover", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ cur_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(); + bsuccess = data.bsuccess; + error = data.error; + if(bsuccess){ + //更新页面 + task_list = [] + cur_task = null //当前选择的task--用于修改缓存时使用 + cur_task_id = 0 //当前选择的cur_task_id + //重新获取任务list + getTasklist(); + }else { + alert("结束任务失败:",error); + } + } + } catch (error) { + alert("结束任务失败:",error); + } +} + //修改了涉及到tasklist的展示内容,修改tasklist显示 function updateTaskList(){ //更新数据 diff --git a/web/main/templates/task_manager.html b/web/main/templates/task_manager.html index 4396a72..f07c8be 100644 --- a/web/main/templates/task_manager.html +++ b/web/main/templates/task_manager.html @@ -141,7 +141,7 @@
-
+
192.168.1.110 @@ -186,10 +186,11 @@
-
+
+