diff --git a/TaskManager.py b/TaskManager.py deleted file mode 100644 index 09b550e..0000000 --- a/TaskManager.py +++ /dev/null @@ -1,321 +0,0 @@ -''' -渗透测试任务管理类 一次任务的闭合性要检查2025-3-10 一次任务后要清理LLM和InstrM的数据 -''' -from mycode.TargetManager import TargetManager # 从模块导入类 -#from LLMManager import LLMManager # 同理修正其他导入 -from mycode.ControlCenter import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。 -from myutils.FileManager import FileManager -from mycode.InstructionManager import InstructionManager -from mycode.DBManager import DBManager -from myutils.MyTime import get_local_timestr -from myutils.MyLogger_logger import LogHandler -import pickle -import queue -import time -import os -import threading - -class TaskManager: - - def __init__(self): - self.TargetM = TargetManager() - self.logger = LogHandler().get_logger("TaskManager") - # 生成功能对象 - self.DBM = DBManager() #主进程一个DBM - if not self.DBM.connect(): - self.logger.error("数据库连接失败！终止工作！") - return - self.CCM = ControlCenter(self.DBM,self) - self.InstrM = InstructionManager(self) # 类对象渗透，要约束只读取信息 - # 控制最大并发指令数量 - self.max_thread_num = 2 - self.task_id = 0 #任务id -- - self.workth_list = [] #线程句柄list - # self.long_instr_num = 0 #耗时指令数量 - # self.long_time_instr = ['nikto'] #耗时操作不计入批量执行的数量，不加不减 - self.node_queue = queue.Queue() - - self.lock = threading.Lock() #线程锁 - self.node_num = 0 #在处理Node线程的处理 - self.brun = True - self.cookie = "" #cookie参数 - - def res_in_quere(self,bres,instr,reslut,start_time,end_time,th_DBM,source_result,ext_params,work_node): - ''' - 执行结果入队列 - :param bres: - :param instr: - :param reslut: - :return: - ''' - #入数据库 -- bres True和False 都入数据库2025-3-10---加node_path(2025-3-18)#？ - if th_DBM.ok: - work_node.do_sn += 1 # 指令的执行序列是一个任务共用,要线程锁，错误问题也不大 - th_DBM.insetr_result(self.task_id, instr, reslut, work_node.do_sn, start_time, end_time, source_result, - ext_params, work_node.path) - else: - self.logger.error("数据库连接失败！！") - - #结果入队列---2025-3-18所有的指令均需返回给LLM便于节点状态的更新，所以bres作用要调整。 - res = {'执行指令':instr,'结果':reslut} - str_res = json.dumps(res,ensure_ascii=False) #直接字符串组合也可以-待验证 - work_node.llm_type = 1 - work_node.add_res(str_res) #入节点结果队列 - - def do_worker_th(self): - #线程的dbm需要一个线程一个 - th_DBM = DBManager() - th_DBM.connect() - while self.brun: - try: - work_node = self.node_queue.get(block=False) - # 测试时使用 - with self.lock: - self.node_num += 1 - # 开始执行指令 - while work_node.instr_queue: - #for instruction in work_node.instr_queue: #这里要重新调整#? - instruction = work_node.instr_queue.pop(0) - start_time = get_local_timestr() # 指令执行开始时间 - bres, instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction) - end_time = get_local_timestr() # 指令执行结束时间 - self.res_in_quere(bres, instr, reslut, start_time, end_time, th_DBM, source_result, - ext_params, work_node) # 执行结果入队列 - - # #针对一个节点的指令执行完成后，提交LLM规划下一步操作 - #self.CCM.llm_quere.put(work_node) - - # 保存记录--测试时使用--后期增加人为停止测试时可以使用 - with self.lock: - self.node_num -= 1 - if self.node_num == 0 and self.node_queue.empty(): # - self.logger.debug("此批次指令执行完成！") - with open("attack_tree", 'wb') as f: - pickle.dump(self.CCM.attack_tree, f) - - except queue.Empty: - time.sleep(20) - - def start_task(self,target_name,target_in): - ''' - - :param target_name: 任务目标名字 - :param target_in: 任务目标访问地址 - :return: - ''' - #判断目标合法性 - bok,target,type = self.TargetM.validate_and_extract(target_in) - if bok: - self.target = target - self.type = type #1-IP，2-domain - self.task_id = self.DBM.start_task(target_name,target_in) #数据库新增任务记录 - #获取基本信息: 读取数据库或预生成指令，获取基本的已知信息 - know_info = "无" #? - #启动--初始化指令 - self.CCM.start_do(target,self.task_id) - - #创建工作线程----2025-3-18调整为一个节点一个线程， - for i in range(self.max_thread_num): - w_th = threading.Thread(target=self.do_worker_th) - w_th.start() - self.workth_list.append(w_th) - - #等待线程结束--执行生成报告 - for t in self.workth_list: - t.join() - #生成报告 - pass - else: - return False,"{target}检测目标不合规，请检查！" - - def stop_task(self): - self.brun = False - self.CCM.stop_do() #清空一些全局变量 - self.InstrM.init_data() - #结束任务需要收尾处理#？ - - -if __name__ == "__main__": - import json - TM = TaskManager() - FM = FileManager() - current_path = os.path.dirname(os.path.realpath(__file__)) - strMsg = FM.read_file("test",1) - - test_type = 2 - instr_index = 19 - iput_index = -1 # 0是根节点 - indexs = [] - if test_type == 0: #新目标测试 - # 启动--初始化指令 - node_list = TM.CCM.start_do("58.216.217.70", 1) - #异步处理，需要等待线程结束了 - for th in TM.CCM.llmth_list: - th.join() - elif test_type == 1: - #测试执行指令 - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - # 遍历node，查看有instr的ndoe - nodes = TM.CCM.attack_tree.traverse_dfs() - if indexs: - for index in indexs: - node = nodes[index] - if node.instr_queue: # list - TM.node_queue.put(node) - else: - for node in nodes: - if node.instr_queue: - TM.node_queue.put(node) - - #创建线程执行指令 - for i in range(TM.max_thread_num): - w_th = threading.Thread(target=TM.do_worker_th) - w_th.start() - TM.workth_list.append(w_th) - # 等待线程结束--执行生成报告 - for t in TM.workth_list: - t.join() - elif test_type ==2: - #测试LLM返回下一步指令 - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - #遍历node，查看有res的数据 - iput_max_num = 0 - iput_num = 0 - nodes = TM.CCM.attack_tree.traverse_dfs() - if indexs: - for index in indexs: - node = nodes[index] - if node.res_quere: - TM.CCM.llm_quere.put(node) - else: - for node in nodes: - if node.res_quere: - TM.CCM.llm_quere.put(node) - - #创建llm工作线程 - TM.CCM.brun = True - for i in range(TM.CCM.max_thread_num): - l_th = threading.Thread(target=TM.CCM.th_llm_worker()) - l_th.start() - TM.CCM.llmth_list.append(l_th) - # 等待线程结束 - for t in TM.CCM.llmth_list: - t.join() - elif test_type ==3: #执行指定指令 - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - # 遍历node，查看有instr的ndoe - nodes = TM.CCM.attack_tree.traverse_dfs() - instrlist = nodes[instr_index].instr_queue - instrlist = [''' - mkdir -p /tmp/nfs_test && mount -t nfs -o nolock 192.168.204.137:/ /tmp/nfs_test && ls /tmp/nfs_test - '''] - for instr in instrlist: - start_time = get_local_timestr() # 指令执行开始时间 - bres, instr, reslut, source_result, ext_params = TM.InstrM.execute_instruction(instr) - end_time = get_local_timestr() # 指令执行结束时间 - - res = {'执行结果': reslut} - str_res = json.dumps(res,ensure_ascii=False) # 直接字符串组合也可以-待验证 - print(str_res) - elif test_type == 4: #修改Message - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - #创建一个新的节点 - from mycode.AttackMap import TreeNode - testnode = TreeNode("test",0,0) - TM.CCM.LLM.build_initial_prompt(testnode)#新的Message - systems = testnode.messages[0]["content"] - #print(systems) - # 遍历node，查看有instr的ndoe - nodes = TM.CCM.attack_tree.traverse_bfs() - for node in nodes: - node.messages[0]["content"] = systems - with open("attack_tree", 'wb') as f: - pickle.dump(TM.CCM.attack_tree, f) - elif test_type ==5: #显示指令和结果list - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - nodes = TM.CCM.attack_tree.traverse_dfs() - if iput_index == -1: - for node in nodes: - print(f"----{node.path}-{node.status}----\n****instr_quere") - print(f"{','.join(node.instr_queue)}\n****res_quere") - try: - print(f"{','.join(node.res_quere)}") - except: - print(f"{json.dumps(node.res_quere)}") - elif iput_index == -2:#只输出有instruction的数据 - index = 0 - for node in nodes: - if node.instr_queue: - print(f"----{index}--{node.path}--{node.status}----") - print(f"{','.join(node.instr_queue)}") - index += 1 - else: - print(f"********\n{','.join(nodes[iput_index].instr_queue)}\n********") - print(f"&&&&&&&&\n{','.join(nodes[iput_index].res_quere)}\n&&&&&&&&") - elif test_type == 6: #给指定节点添加测试指令 - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - nodes = TM.CCM.attack_tree.traverse_dfs() - str_instr = "nmap -sV -p- 192.168.204.137 -T4 -oN nmap_full_scan.txt" - index = 9 - nodes[index].instr_queue.append(str_instr) - nodes[index].res_quere = [] - with open("attack_tree", 'wb') as f: - pickle.dump(TM.CCM.attack_tree, f) - elif test_type == 7: #给指定节点修改指令的执行结果 - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - nodes = TM.CCM.attack_tree.traverse_dfs() - str_instr = "psql -h 192.168.204.137 -U postgres -c '\l'" - start_time = get_local_timestr() # 指令执行开始时间 - bres, instr, reslut, source_result, ext_params = TM.InstrM.execute_instruction(str_instr) - end_time = get_local_timestr() # 指令执行结束时间 - # 入数据库 -- bres True和False 都入数据库2025-3-10---加node_path(2025-3-18)#？ - if TM.DBM.ok: - TM.DBM.insetr_result(0, instr, reslut, 0, start_time, end_time, source_result, - ext_params, "独立命令执行") - index = 9 - nodes[index].res_quere.clear() - nodes[index].res_quere.append(reslut) - - with open("attack_tree", 'wb') as f: - pickle.dump(TM.CCM.attack_tree, f) - elif test_type ==8: #显示有漏洞信息的数据 - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - nodes = TM.CCM.attack_tree.traverse_dfs() - if nodes: - for node in nodes: - if node.vul_type != "未发现": - print(f"{node.path}----{node.vul_type}") - elif test_type == 9: #处理自定义llm回复内容 - with open("attack_tree", "rb") as f: - TM.CCM.attack_tree = pickle.load(f) - nodes = TM.CCM.attack_tree.traverse_dfs() - node = nodes[5] - strconent = ''' - {'role': 'assistant', 'content': '{"action":"update_status", "node": "25端口", "status": "已完成", "vulnerability": {"name":"SMTP用户枚举漏洞","risk":"中危","info":"VRFY命令可验证有效用户"}}\n\n```bash-[目标系统->192.168.204.137->25端口]\nsmtp-user-enum -M VRFY -U /usr/share/wordlists/metasploit/unix_users.txt -t 192.168.204.137\n```\n\n```bash-[目标系统->192.168.204.137->25端口]\nnc -nv 192.168.204.137 25 << EOF\nEXPN root\nMAIL FROM: attacker@example.com\nRCPT TO: external@example.com\nDATA\nTest open relay\n.\nQUIT\nEOF\n```'} - ''' - strjson = json.loads(strconent) - node_cmds,commands = TM.CCM.LLM.fetch_instruction(strjson["content"]) - TM.CCM.tree_manager(node_cmds) - - # 更新tree - bok, new_commands = TM.CCM.tree_manager(node_cmds, node, commands, TM.DBM) - # 分析指令入对应节点 - if bok: # 节点指令若存在错误，测试指令都不处理，需要LLM重新生成 - node_list = TM.CCM.instr_in_node(new_commands, node) - - #报不保存待定-- - with open("attack_tree", 'wb') as f: - pickle.dump(TM.CCM.attack_tree, f) - - - else: - #完整过程测试---要设定终止条件 - pass diff --git a/config.yaml b/config.yaml index 7df3f1a..7eee4b1 100644 --- a/config.yaml +++ b/config.yaml @@ -1,3 +1,8 @@ +#工作模式 +App_Work_type: 0 #开发模式，只允许单步模式 + +#线程休眠的时间 +sleep_time: 20 #日志记录 file_log_level: INFO #是否记录日志 @@ -14,7 +19,6 @@ mysql: database: zfsafe #LLM-Type -LLM_type: 2 #0-腾讯云，1-DS，2-2233ai,3-GPT LLM_max_chain_count: 10 #为了避免推理链过长，造成推理效果变差，应该控制一个推理链的长度上限 #用户初始密码 @@ -25,3 +29,9 @@ serverIP: 192.168.3.190 serverPort: 18010 DevID: 12345678901234567890123456789012 sockettimeout: 600 #10分钟 + +#node_tree_file_name +TreeFile: tree_data/attack_tree + +#task +Task_max_threads: 5 diff --git a/main.py b/main.py deleted file mode 100644 index 50be11d..0000000 --- a/main.py +++ /dev/null @@ -1,29 +0,0 @@ -import asyncio -import uvicorn -import TaskManager -import os -from web import create_app -from hypercorn.asyncio import serve -from hypercorn.config import Config - -async def run_quart_app(): - app = create_app() - config = Config() - config.bind = ["0.0.0.0:5001"] - await serve(app, config) - -#对某个目标进行测试 -def startWork(targets): - TaskM = TaskManager() - tlist = targets.split(',') - for target in tlist: - TaskM.start_task(target) - - -# Press the green button in the gutter to run the script. -if __name__ == '__main__': - print(f"Current working directory (run.py): {os.getcwd()}") - #启动web项目 - asyncio.run(run_quart_app()) - uvicorn.run("run:app", host="0.0.0.0", port=5001, workers=4, reload=True) - diff --git a/mycode/AttackMap.py b/mycode/AttackMap.py index f7233ba..d80d4bc 100644 --- a/mycode/AttackMap.py +++ b/mycode/AttackMap.py @@ -1,6 +1,7 @@ import queue import copy import re +import threading #渗透测试树结构维护类 class AttackTree: @@ -9,6 +10,7 @@ class AttackTree: self.root = root_node self.root.path = f"目标系统->{root_node.name}" + def set_root(self,root_node): self.root = root_node @@ -47,6 +49,24 @@ class AttackTree: self.traverse_dfs(child, result) return result + #生成节点树字典数据 + def node_to_dict(self,node): + return { + "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, + "node_workstatus":node.get_work_status(), + "children":[self.node_to_dict(child) for child in node.children] if node.children else [] + } + + #树简化列表，用户传输到前端 + def get_node_dict(self): + node_dict = self.node_to_dict(self.root) #递归生成 + return node_dict + def find_node_by_name(self, name): """根据名称查找节点（广度优先）""" nodes = self.traverse_bfs() @@ -96,6 +116,17 @@ class AttackTree: return None return current_node + #更新节点的bwork状态 + def update_node_bwork(self,node_path): + node = self.find_node_by_nodepath(node_path) + if not node: + return False,False + if node.bwork: + node.bwork = False + else: + node.bwork = True + return True,node.bwork + def find_nodes_by_status(self, status): """根据状态查找所有匹配节点""" return [node for node in self.traverse_bfs() if node.status == status] @@ -135,22 +166,28 @@ class TreeNode: def __init__(self, name,task_id,status="未完成", vul_type="未发现"): self.task_id = task_id #任务id self.name = name # 节点名称 - self.status = status # 节点状态 - self.vul_type = vul_type # 漏洞类型 + #self.node_lock = threading.Lock() #线程锁 + self.bwork = True # 当前节点是否工作，默认True --停止/启动 + self.status = status # 节点测试状态 -- 由llm返回指令触发更新 + #work_status需要跟两个list统一管理：初始0，入instr_queue为1，入instr_node_mq为2，入res_queue为3，入llm_node_mq为4，llm处理完0或1 + self._work_status = 0 #0-无任务，1-待执行测试指令，2-执行指令中，3-待提交Llm,4-提交llm中， 2025-4-6新增，用来动态显示节点的工作细节。 + self.vul_type = vul_type # 漏洞类型--目前赋值时没拆json self.vul_name = "" self.vul_grade = "" self.vul_info = "" self.children = [] # 子节点列表 self.parent = None # 父节点引用 self.path = "" #当前节点的路径 - self.bwork = True #当前节点是否工作，默认True - self.messages = [] # 针对当前节点积累的messages -- 针对不同节点提交不同的messages + self.llm_type = 0 #llm提交类型 0--初始状态无任务状态，1--指令结果反馈，2--llm错误反馈 self.llm_sn = 0 #针对该节点llm提交次数 + self._llm_quere = [] #待提交llm的数据 + self.do_sn = 0 #针对该节点instr执行次数 - self.instr_queue = [] # queue.Queue() #针对当前节点的执行指令----重要约束：一个节点只能有一个线程在执行指令 - self.res_quere = [] # queue.Queue() #指令执行的结果，一批一批 + self._instr_queue = [] #针对当前节点的待执行指令----重要约束：一个节点只能有一个线程在执行指令 + + self.his_instr = [] #保留执行指令的记录｛“instr”:***,"result":***｝ #用户补充信息 self.cookie = "" self.ext_info = "" @@ -196,7 +233,7 @@ class TreeNode: else: print("非法的信息体类型！") - + #添加子节点 def add_child(self, child_node): child_node.parent = self child_node.path = self.path + f"->{child_node.name}" #子节点的路径赋值 @@ -205,17 +242,70 @@ class TreeNode: self.children.append(child_node) + #修改节点的执行状态--return bchange + def update_work_status(self,work_status): + if self._work_status == work_status: + return False + self._work_status = work_status + return True + + def get_work_status(self): + #加锁有没有意义--待定 + return self._work_status + + #-------后期扩充逻辑，目前wokr_status的修改交给上层类对象------- def add_instr(self,instr): - self.instr_queue.append(instr) + self._instr_queue.append(instr) def get_instr(self): - return self.instr_queue.pop(0) if self.instr_queue else None + return self._instr_queue.pop(0) if self._instr_queue else None + + def get_instr_user(self): + return self._instr_queue + + def del_instr(self,instr): + if instr in self._instr_queue: + self._instr_queue.remove(instr) + #指令删除后要判断是否清空指令了 + if not self._instr_queue: + self._work_status = 0 #状态调整为没有带执行指令 + return True,"" + else: + return False,"该指令不在队列中！" + + def add_res(self,str_res): #结构化结果字串 - self.res_quere.append(str_res) + self._llm_quere.append(str_res) def get_res(self): - return self.res_queue.pop(0) if self.res_queue else None + return self._llm_quere.pop(0) if self._llm_quere else None + + def get_res_user(self): + return self._llm_quere + + def get_work_status(self): + return self._work_status + + def updatemsg(self,newtype,newcontent,index): + newmsg = {"llm_type":int(newtype),"result":newcontent} + if self._llm_quere: + self._llm_quere[0] = newmsg + else:#新增消息 + self._llm_quere.append(newmsg) + #更新节点状态 + self._work_status = 3 #待提交 + return True,"" + + def is_instr_empty(self): + if self._instr_queue: + return False + return True + + def is_llm_empty(self): + if self._llm_quere: + return False + return True def __repr__(self): return f"TreeNode({self.name}, {self.status}, {self.vul_type})" diff --git a/mycode/ClientSocket.py b/mycode/ClientSocket.py new file mode 100644 index 0000000..7fab729 --- /dev/null +++ b/mycode/ClientSocket.py @@ -0,0 +1,4 @@ +class ClientSocket: + def __init__(self): + self.user_id = -1 + self.th_read = None \ No newline at end of file diff --git a/mycode/ControlCenter.py b/mycode/ControlCenter.py index 63a29ae..0397190 100644 --- a/mycode/ControlCenter.py +++ b/mycode/ControlCenter.py @@ -5,35 +5,20 @@ import re import queue import time import threading -import pickle from mycode.AttackMap import AttackTree from mycode.AttackMap import TreeNode -from mycode.LLMManager import LLMManager -from myutils.ConfigManager import myCongif #单一实例 from myutils.MyLogger_logger import LogHandler from mycode.DBManager import DBManager class ControlCenter: - def __init__(self,DBM,TM): + def __init__(self): self.logger = LogHandler().get_logger("ControlCenter") - self.task_id = None - self.target = None - self.attack_tree = None - self.DBM = DBM - self.TM = TM - #LLM对象 - self.LLM = LLMManager(myCongif.get_data("LLM_type")) - self.llm_quere = queue.Queue() #提交LLM指令的队列---- - self.max_thread_num = 1 # 控制最大并发指令数量 - self.llmth_list = [] #llm线程list - self.brun = False def __del__(self): self.brun =False self.task_id = None self.target = None self.attack_tree = None - self.DBM = None def init_cc_data(self): #一次任务一次数据 @@ -48,152 +33,7 @@ class ControlCenter: # ?包括是否对目标进行初始化的信息收集 return {"已知信息":"无"} - def start_do(self,target,task_id): - '''一个新任务的开始''' - self.task_id = task_id - self.target = target - #创建/初始化测试树 - if self.attack_tree: - self.attack_tree = None #释放 - root_node = TreeNode(target,task_id) - self.attack_tree = AttackTree(root_node)#创建测试树，同时更新根节点相关内容 - #初始化启动提示信息 - know_info = self.get_user_init_info() - self.LLM.build_initial_prompt(root_node) - #提交到待处理队列 - self.put_one_llm_work(know_info,root_node,1) - - # 启动LLM请求提交线程 - self.brun = True - #启动线程 - ineed_create_num = self.max_thread_num - len(self.llmth_list) #正常应该就是0或者是max_num - for i in range(ineed_create_num): - l_th = threading.Thread(target=self.th_llm_worker) - l_th.start() - self.llmth_list.append(l_th) - - def put_one_llm_work(self,str_res,node,llm_type): - '''提交任务到llm_quere''' - if llm_type == 0: - self.logger.debug("llm_type不能设置为0") - return - node.llm_type = llm_type #目前处理逻辑中一个node应该只会在一个queue中，且只会有一次记录。所以llm_type复用 - node.res_quere.append(str_res) - #提交到待处理队列 - self.llm_quere.put(node) - - def get_one_llm_work(self,node): - '''获取该节点的llm提交数据，会清空type和res_quere''' - llm_type = node.llm_type - node.llm_type = 0 #回归0 - res_list = node.res_quere[:] #浅拷贝，复制第一层 - node.res_quere.clear() #清空待处理数据，相当于把原应用关系接触 - return llm_type,res_list - - def restore_one_llm_work(self,node,llm_type,res_list): - node.llm_type = llm_type - node.res_quere = res_list - - #llm请求提交线程 - def th_llm_worker(self):#LLM没有修改的全局变量，应该可以共用一个client - ''' - 几个规则--TM的work线程同 - 1.线程获取一个节点后，其他线程不能再获取这个节点（遇到被执行的节点，直接放弃执行）--- 加了没办法保存中间结果进行测试 - 2.llm返回的指令，只可能是该节点，或者是子节点的，不符合这条规则的都不处理，避免llm处理混乱。 - :return: - ''' - # 线程的dbm需要一个线程一个 - th_DBM = DBManager() - th_DBM.connect() - while self.brun: - try: - #节点锁 - node = self.llm_quere.get(block=False) - self.get_llm_instruction(node,th_DBM) - #释放锁 - - # 暂存状态--测试时使用--限制条件llm工作线程只能未1个 - with open("attack_tree", 'wb') as f: - pickle.dump(self.attack_tree, f) - - except queue.Empty: - self.logger.debug("llm队列中暂时无新的提交任务！") - time.sleep(30) - - #约束1：一个节点只能同时提交一次，未测试的节点不要重复 - def get_llm_instruction(self,node,DBM): - user_Prompt = "" - ext_Prompt = "" - llm_type, res_list = self.get_one_llm_work(node) - try: - res_str = json.dumps(res_list, ensure_ascii=False) - except TypeError as e: - self.logger.error(f"{res_list}序列化失败：{e},需要自查程序添加代码！") - return # 直接返回 - if llm_type == 0: - self.logger.error("这里type不应该为0，请自查程序逻辑！") - return - #2025-3-20增加了llm_type - user_Prompt = f''' -当前分支路径：{node.path} -当前节点信息： -- 节点名称：{node.name} -- 节点状态：{node.status} -- 漏洞类型：{node.vul_type} -''' - if llm_type == 1: #提交指令执行结果 --- 正常提交 - # 构造本次提交的prompt - ext_Prompt = f''' -上一步结果：{res_str} -任务：生成下一步渗透测试指令或判断是否完成该节点测试。 -''' - elif llm_type ==2: #llm返回的指令存在问题，需要再次请求返回 - ext_Prompt = f''' -反馈类型：节点指令格式错误 -错误信息：{res_str} -任务：请按格式要求重新生成该节点上一次返回中生成的所有指令。 -''' -# ''' -# elif llm_type ==4: #未生成节点列表 -# ext_Prompt = f''' -# 反馈类型：需要继续补充信息 -# 缺失信息：{res_str} -# 任务： -# 1.请生成这些节点的新增节点指令，并生成对应的测试指令； -# 2.这些节点的父节点为当前节点，请正确生成这些节点的节点路径； -# 3.若节点未能全部新增，必须返回未新增的节点列表 -# 4.若有未生成指令的节点，必须返回未生成指令的节点列表。 -# ''' - elif llm_type ==5: - ext_Prompt = f''' -反馈类型：测试指令格式错误 -错误信息：{res_str} -任务：请根据格式要求，重新生成该测试指令。 -''' - else: - self.logger.debug("意外的类型参数") - return - if not ext_Prompt: - self.logger.error("未成功获取本次提交的user_prompt") - return - #提交LLM - user_Prompt = user_Prompt + ext_Prompt - node_cmds, commands = self.LLM.get_llm_instruction(user_Prompt,DBM, node) # message要更新 - ''' - 对于LLM返回的错误处理机制 - 1.验证节点是否都有测试指令返回 - 2.LLM的回复开始反复时（有点难判断） - ''' - # 更新tree - bok,new_commands = self.tree_manager(node_cmds, node,commands,DBM) - # 分析指令入对应节点 - if bok: #节点指令若存在错误，测试指令都不处理，需要LLM重新生成 - node_list = self.instr_in_node(new_commands, node) - # 插入TM的node_queue中,交TM线程处理---除了LLM在不同的请求返回针对同一节点的测试指令，正常业务不会产生两次进队列 - for node in node_list: - self.TM.node_queue.put(node) - - def verify_node_cmds(self,node_cmds,node): + def verify_node_cmds(self,node_cmds): ''' 验证节点指令的合规性，持续维护 :param node_cmds: @@ -223,184 +63,12 @@ class ControlCenter: strerror = {"节点指令错误": f"{node_json}不可识别的action值！"} break if not strerror: - return True - #提交一个错误反馈任务 - self.put_one_llm_work(strerror,node,2) - return False - - def tree_manager(self,node_cmds,node,commands,DBM): - '''更新渗透测试树 - node_cmds是json-list - 2025-03-22添加commands参数，用于处理LLM对同一个节点返回了测试指令，但还返回了no_instruction节点指令 - ''' - if not node_cmds: # or len(node_cmds)==0: 正常not判断就可以有没有节点指令 - return True,commands - #对节点指令进行校验 - if not self.verify_node_cmds(node_cmds,node): - return False,commands #节点指令存在问题，终止执行 - - #执行节点操作---先执行add_node，怕返回顺序不一直 - residue_node_cmds = [] - for node_json in node_cmds: - action = node_json["action"] - if action == "add_node": # 新增节点 - parent_node_name = node_json["parent"] - status = node_json["status"] - node_names = node_json["nodes"].split(',') - # 新增节点原则上应该都是当前节点增加子节点 - if node.name == parent_node_name or parent_node_name.endswith(node.name): - 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的传递待验证 - 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) - else: - self.logger.error(f"添加子节点时，遇到父节点名称不一致的，需要介入！！{node_json}") # 丢弃该节点 - else:#其他指令添加到list - residue_node_cmds.append(node_json) + return True,strerror + return False,strerror - #执行剩余的节点指令--不分先后 - for node_json in residue_node_cmds: - action = node_json["action"] - if action == "update_status": - 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.name = 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 - 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任务 - continue - #验证对应节点是否已经创建---本节点或子节点，其他节点不处理（更狠一点就是本节点都不行） - if node_name == node.name: - nodes.append(node_name) - # str_add = "请生成测试指令" - # self.put_one_llm_work(str_add,node,1) - else: - for child_node in node.children: - if child_node.name == node_name: - nodes.append(node_name) - # str_add = "无" - # self.put_one_llm_work(str_add, child_node, 1) - break - 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: - str_add = {"未新增的节点": nodes} - self.logger.debug(str_add) - # 提交一个继续反馈任务--继续后续工作 2025-3-25不自动处理 - # self.put_one_llm_work(str_add, node, 4) - # self.logger.debug(f"未新增的节点有：{nodes}") - else: - self.logger.error("****不应该执行到这！程序逻辑存在问题！") - return True,commands - - #阻塞轮询补充指令 - def get_other_instruction(self,nodes,DBM,cur_node): - res_str = ','.join(nodes) - new_commands = [] - while res_str: - self.logger.debug(f"开始针对f{res_str}这些节点请求测试指令") - user_Prompt = f''' - 当前分支路径：{cur_node.path} - 当前节点信息： - - 节点名称：{cur_node.name} - - 节点状态：{cur_node.status} - - 漏洞类型：{cur_node.vul_type} - 反馈类型：需要补充信息 - 缺失信息：针对{res_str}的测试指令 - 任务： - 1.请生成这些节点的测试指令； - 2.这些节点的父节点为当前节点，请正确生成这些节点的节点路径； - 3.若还有节点未能生成测试指令，必须返回未生成指令的节点列表。 - ''' - res_str = "" - node_cmds, commands = self.LLM.get_llm_instruction(user_Prompt, DBM, cur_node) # message要更新 - #把返回的测试指令进行追加 - new_commands.extend(commands) - #判断是否还有未添加指令的节点 - for node_json in node_cmds: #正常应该只有一条no_instruction - if "no_instruction" in node_json and "nodes" in node_json: - tmp_nodes = [] - node_names = node_json["nodes"].split(',') - for node_name in node_names: - if node_name in nodes: - tmp_nodes.append(node_name) - res_str = ','.join(tmp_nodes) - break - self.logger.debug("为添加指令的节点，都已完成指令的添加！") - return new_commands - - def instr_in_node(self,commands,node): - node_list = [] #一次返回的测试指令 - for command in commands: - # 使用正则匹配方括号中的node_path（非贪婪模式） - match = re.search(r'\[(.*?)\]', command) - if match: - node_path = match.group(1) - #'''强制约束，不是本节点或者是子节点的指令不处理''' - find_node = self.attack_tree.find_node_by_nodepath_parent(node_path,node) - if find_node: - instruction = re.sub(r'\[.*?\]', "", command,count=1,flags=re.DOTALL) - find_node.instr_queue.append(instruction) - #入输出队列 - if find_node not in node_list: - node_list.append(find_node) - else: - self.logger.error(f"基于节点路径没有找到对应的节点{node_path},父节点都没找到！")#丢弃该指令 - else: - self.logger.error(f"得到的指令格式不符合规范：{command}")#丢弃该指令--- - #这里对于丢弃指令，有几种方案： - # 1.直接丢弃不处理，但需要考虑会不会产生节点缺失指令的问题，需要有机制验证节点；------ 需要有个独立线程，节点要加锁--首选待改进方案 - # 2.入当前节点的res_queue,但有可能当前节点没有其他指令，不会触发提交，另外就算提交，会不会产生预设范围外的返回，不确定； - # 3.独立队列处理 - return node_list + def restore_one_llm_work(self,node,llm_type,res_list): + node.llm_type = llm_type + node.res_quere = res_list #需要用户确认的信息--待完善 def need_user_know(self,strinfo,node): diff --git a/mycode/DBManager.py b/mycode/DBManager.py index 6a83954..49b167e 100644 --- a/mycode/DBManager.py +++ b/mycode/DBManager.py @@ -6,6 +6,7 @@ import json from myutils.ConfigManager import myCongif from myutils.MyLogger_logger import LogHandler from myutils.MyTime import get_local_timestr +from datetime import timedelta class DBManager: #实例化数据库管理对象，并连接数据库 @@ -118,6 +119,22 @@ class DBManager: self.lock.release() return bok + def safe_do_select(self,strsql,params,itype=0): + results = [] + if self.Retest_conn(): + cursor = self.connection.cursor() #每次查询使用新的游标 ---待验证（用于处理出现一次查询不到数据的问题） + try: + 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() + return results + def is_json(self,s:str) -> bool: if not isinstance(s, str): return False @@ -129,9 +146,23 @@ class DBManager: except Exception: return False # 处理其他意外异常（如输入 None） + def timedelta_to_str(delta: timedelta) -> str: + hours, remainder = divmod(delta.total_seconds(), 3600) + minutes, seconds = divmod(remainder, 60) + return f"{int(hours):02}:{int(minutes):02}:{int(seconds):02}" + #---------------------特定数据库操作函数--------------------- + def get_system_info(self): + strsql = "select * from zf_system;" + data = self.do_select(strsql,1) + return data - def start_task(self,task_name,task_target) -> int: + def get_run_tasks(self): + strsql = "select ID,task_target,task_status,work_type,cookie_info,llm_type from task where task_status <> 2;" + datas = self.do_select(strsql) + return datas + + def start_task(self,test_target,cookie_info,work_type,llm_type) -> int: ''' 数据库添加检测任务 :param task_name: @@ -140,8 +171,9 @@ class DBManager: ''' task_id =0 start_time = get_local_timestr() - sql = "INSERT INTO task (task_name,task_target,start_time) VALUES (%s,%s,%s)" - params = (task_name,task_target,start_time) + 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 return task_id @@ -183,7 +215,7 @@ class DBManager: return self.safe_do_sql(sql,params) #llm数据入库 - def insert_llm(self,task_id,prompt,reasoning_content,content,post_time,node): + def insert_llm(self,task_id,prompt,reasoning_content,content,post_time,llm_sn,path): str_reasoning = "" str_content = "" try: @@ -215,9 +247,28 @@ class DBManager: str_reasoning = str_reasoning.encode('utf-8').decode('unicode_escape') str_content = str_content.encode('utf-8').decode('unicode_escape') - params = (task_id,node.llm_sn,prompt,str_reasoning,str_content,post_time,node.path) + params = (task_id,llm_sn,prompt,str_reasoning,str_content,post_time,path) return self.safe_do_sql(sql,params) + #获取任务的测试指令执行情况 + def get_task_instrs(self,task_id,nodename): + instrs = [] + return instrs + + #获取任务的漏洞检测情况 + def get_task_vul(self,task_id,nodename,vultype,vullevel): + vuls =[] + return vuls + + #获取该任务该节点的所有 已经执行的任务 + 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; + ''' + params = (task_id,nodepath) + datas = self.safe_do_select(strsql,params) + return datas + def test(self): # 建立数据库连接 conn = pymysql.connect( diff --git a/mycode/InstructionManager.py b/mycode/InstructionManager.py index ad8420c..9b58f2c 100644 --- a/mycode/InstructionManager.py +++ b/mycode/InstructionManager.py @@ -13,8 +13,7 @@ from tools.ToolBase import ToolBase from myutils.ReturnParams import ReturnParams class InstructionManager: - def __init__(self,TM): - self.TM = TM #任务类对象穿透 + def __init__(self): self.tool_registry = {} # 安全工具list self.load_tools() # 加载工具类 @@ -36,7 +35,7 @@ class InstructionManager: cls = getattr(module,module_name) if(issubclass(cls, ToolBase) and cls != ToolBase): tool_name = module_name.lower()[:-4] - self.tool_registry[tool_name] = cls(self.TM) #类对象穿透 + self.tool_registry[tool_name] = cls() except ImportError as e: print(f"加载工具 {module_name} 失败：{str(e)}") @@ -61,14 +60,15 @@ class InstructionManager: else: bres = False instr = instruction #保障后续代码的一致性 - source_result = result = f"未知工具：{tool_name}" + 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}") - return bres,instr,result,source_result,ext_params + #return bres,instr,result,source_result,ext_params + return instr, result, source_result, ext_params #取消bres的返回，所有指令执行结果多需要返回到Llm,用于控制节点的状态 #过来指令：合规、判重、待执行等 def _instruction_filter(self,instruction): @@ -78,6 +78,8 @@ class InstructionManager: def _fetch_prompt(self): pass +g_instrM = InstructionManager() #全局唯一 + if __name__ == "__main__": instrM = InstructionManager() instrS = ['whois haitutech.cn', diff --git a/mycode/LLMBase.py b/mycode/LLMBase.py deleted file mode 100644 index 1a73954..0000000 --- a/mycode/LLMBase.py +++ /dev/null @@ -1,8 +0,0 @@ -import openai -from openai import OpenAI -#LLM的基类 - -class LLMBase: - def __init__(self): - pass - diff --git a/mycode/LLMManager.py b/mycode/LLMManager.py index 9b8824c..480eb89 100644 --- a/mycode/LLMManager.py +++ b/mycode/LLMManager.py @@ -5,22 +5,19 @@ pip install openai ''' import openai import json -import threading import re import os from openai import OpenAI -from mycode.DBManager import DBManager +from myutils.ConfigManager import myCongif from myutils.MyTime import get_local_timestr from myutils.MyLogger_logger import LogHandler class LLMManager: - def __init__(self,illm_type=3): + def __init__(self,illm_type): self.logger = LogHandler().get_logger("LLMManager") self.api_key = None self.api_url = None - self.task_id =0 #一个任务一个id - # self.llm_sn = 0 # llm执行序列号，--一任务一序列 - # self.llm_sn_lock = threading.Lock() # + #temperature设置 #DS------代码生成/数学解题：0.0 -- 数据抽取/分析：1.0 -- 通用对话：1.3 -- 翻译：1.3 -- 创意类写作：1.5 if illm_type == 0: #腾讯云 @@ -51,15 +48,6 @@ class LLMManager: #self.client = openai self.client = OpenAI(api_key=self.api_key,base_url=self.api_url) - #******测试使用，设置slef.message的值 - def test_old_message(self,strMessage): - try: - self.messages = json.loads(strMessage) - except json.JSONDecodeError as e: - print(f"JSON解析错误: {str(e)}") - except Exception as e: - print(f"错误: {str(e)}") - ''' **决策原则** - 根据节点类型和状态，优先执行基础测试（如端口扫描、服务扫描）。 @@ -67,7 +55,7 @@ class LLMManager: - 确保每个新增节点匹配测试指令。 ''' # 初始化messages - def build_initial_prompt(self,node,str_ip=""): + def build_initial_prompt(self,node): if not node: return #根节点初始化message @@ -85,12 +73,11 @@ class LLMManager: 2. 若端口扫描发现开放端口，对可能存在中高危以上风险的端口新增节点并提供测试指令； 3. 若当前节点是端口且未进行服务扫描，则执行服务扫描； 4. 若服务扫描发现服务版本或漏洞，则新增漏洞测试节点并提供测试指令； -5. 若漏洞验证成功，则根据结果决定是否需要进一步测试，若需要进一步测试，则为测试内容新增子节点并提供测试指令； -6. 当当前节点执行完成所有可能的测试指令，更新状态为“已完成”。 +5. 若漏洞验证成功，则根据结果决定是否需要进一步测试，若需要进一步测试，则必须将进一步测试的内容作为子节点并提供测试指令； +6. 当当前节点没有新的测试指令时，更新状态为“已完成”。 **测试指令生成准则** 1.明确每个测试指令的测试目标，并优先尝试最简单、最直接的办法，不要在同一个请求生成测试效果覆盖的指令； 2.使用递进逻辑组织指令：先尝试基础测试方法，根据执行结果决定是否进行更深入的测试; -3.本地的IP地址为：192.168.204.135。 **节点指令格式** - 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1，节点2\", \"status\": \"未完成\"}； - 未生成指令节点列表：{\"action\": \"no_instruction\", \"nodes\": \"节点1,节点2\"}； @@ -111,14 +98,8 @@ mysql -u root -p 192.168.1.100 ``` '''}] # 一个messages - def init_data(self,task_id=0): - #初始化LLM数据 - self.llm_sn = 0 - self.task_id = task_id - self.messages = [] - # 调用LLM生成指令 - def get_llm_instruction(self,prompt,th_DBM,node): + def get_llm_instruction(self,prompt,node): ''' 1.由于大模型API不记录用户请求的上下文，一个任务的LLM不能并发！ :param prompt:用户本次输入的内容 @@ -126,16 +107,22 @@ mysql -u root -p 192.168.1.100 ''' #添加本次输入入该节点的message队列 message = {"role":"user","content":prompt} - node.messages.append(message) + node.messages.append(message) #更新节点message #提交LLM post_time = get_local_timestr() + if self.model == "o3-mini-2025-01-31": + response = self.client.chat.completions.create( + model=self.model, + reasoning_effort="high", + messages = node.messages + ) + else: + response = self.client.chat.completions.create( + model=self.model, + messages=node.messages + ) - response = self.client.chat.completions.create( - model=self.model, - reasoning_effort="high", - messages = node.messages - ) #LLM返回结果处理 reasoning_content = "" content = "" @@ -160,17 +147,10 @@ mysql -u root -p 192.168.1.100 node.messages.append({'role': 'assistant', 'content': content}) else: self.logger.error("处理到未预设的模型！") - return None - - #LLM记录存数据库 - node.llm_sn += 1 - bres = th_DBM.insert_llm(self.task_id,prompt,reasoning_content,content,post_time,node) - if not bres: - self.logger.error(f"{node.name}-llm入库失败！") - + return "","","","","" #按格式规定对指令进行提取 node_cmds,commands = self.fetch_instruction(content) - return node_cmds,commands + return node_cmds,commands,reasoning_content, content, post_time def fetch_instruction(self,response_text): ''' @@ -241,9 +221,64 @@ mysql -u root -p 192.168.1.100 ) print(response) - if __name__ == "__main__": llm = LLMManager(3) - llm.test_llm() + 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/TargetManager.py b/mycode/TargetManager.py index 0ccc20a..5e97c4a 100644 --- a/mycode/TargetManager.py +++ b/mycode/TargetManager.py @@ -42,6 +42,8 @@ class TargetManager: else: return False, None,type +g_TM = TargetManager() + if __name__ == "__main__": tm = TargetManager() # 示例测试 diff --git a/mycode/TaskManager.py b/mycode/TaskManager.py new file mode 100644 index 0000000..31a6dee --- /dev/null +++ b/mycode/TaskManager.py @@ -0,0 +1,202 @@ +from myutils.ConfigManager import myCongif +from mycode.TaskObject import TaskObject +from mycode.DBManager import app_DBM +from myutils.PickleManager import g_PKM + +import threading + +class TaskManager: + def __init__(self): + self.tasks = {} # 执行中的任务，test_id为key + self.num_threads = myCongif.get_data("Task_max_threads") + #获取系统信息 -- 用户可修改的都放在DB中，不修改的放config + data = app_DBM.get_system_info() + self.local_ip = data[0] + self.version = data[1] + self.tasks_lock = threading.Lock() #加个线程锁?不使用1，quart要使用的异步锁，2.限制只允许一个用户登录，3.遍历到删除的问题不大 + self.web_cur_task = 0 #web端当前显示的 + + + #判断目标是不是在当前执行任务中，---没加锁，最多跟删除有冲突，问题应该不大 + def is_target_in_tasks(self,task_target): + for task in self.tasks.values(): + if task_target == task.target: + return True + return False + + #程序启动后，加载未完成的测试任务 + def load_tasks(self): + '''程序启动时，加载未执行完成的任务''' + datas = app_DBM.get_run_tasks() + for data in datas: + task_id = data[0] + task_target = data[1] + task_status = data[2] + work_type = data[3] + cookie_info = data[4] + llm_type = data[5] + # 创建任务对象 + task = TaskObject(task_target, cookie_info, work_type, llm_type, self.num_threads, self.local_ip,self) + #读取attact_tree + attack_tree = g_PKM.ReadData(str(task_id)) + #开始任务 ---会根据task_status来判断是否需要启动工作线程 + task.start_task(task_id,task_status,attack_tree) + # 保留task对象 + self.tasks[task_id] = task + + #新建测试任务 + def create_task(self, test_target,cookie_info,llm_type,work_type): + """创建新任务--create和load复用？-- + 1.创建和初始化task_object; + 2.创建task_id + 3.启动工作线程 + return T/F + """ + if self.is_target_in_tasks(test_target): + raise ValueError(f"Task {test_target} already exists") + #创建任务对象 + task = TaskObject(test_target,cookie_info,work_type,llm_type,self.num_threads,self.local_ip,self) + #获取task_id -- test_target,cookie_info,work_type,llm_type 入数据库 + task_id = app_DBM.start_task(test_target,cookie_info,work_type,llm_type) + if task_id >0: + #创建后启动工作--同时赋值task_id + task.start_task(task_id) + #保留task对象 + self.tasks[task_id] = task + return True + else: + return False + + #控制task启停----线程不停 + def control_taks(self,task_id): + task = self.tasks[task_id] + if task: + if task.task_status == 0: # 0-暂停，1-执行中，2-已完成 + task.task_status = 1 + elif task.task_status == 1: + task.task_status = 0 + else: + return False,"当前任务状态不允许修改，请联系管理员！",task.task_status + else: + return False,"没有找到对应的任务",None + return True,"",task.task_status + + # 获取任务list + def get_task_list(self): + tasks = [] + for task in self.tasks.values(): + one_json = {"taskID": task.task_id, "testTarget": task.target, "taskStatus": task.task_status, "safeRank": task.safe_rank, + "workType": task.work_type} + tasks.append(one_json) + rejson = {"tasks": tasks} + return rejson + + #获取节点树 + def get_node_tree(self,task_id): + task = self.tasks[task_id] + if task: + self.web_cur_task = task_id + tree_dict = task.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] + if task: + if task.task_status == 0: + task.work_type = new_work_type + return True + return False + + #控制节点的工作状态 + def node_bwork_control(self,task_id,node_path): + task = self.tasks[task_id] + if task: + bsuccess,new_bwork = task.attack_tree.update_node_bwork(node_path) + if bsuccess: + pass #是否要更新IO数据？----待验证是否有只修改部分数据的方案 + return bsuccess,new_bwork + return False,False + + #节点单步--只允許web端调用 + async def node_one_step(self,task_id,node_path): + task = self.tasks[task_id] + node = task.attack_tree.find_node_by_nodepath(node_path) + #web端触发的任务，需要判断当前的执行状态 + bsuccess,error = await task.put_one_node(node) + return bsuccess,error + + #任务单点--只允许web端调用 + async def task_one_step(self,task_id): + task = self.tasks[task_id] + if task: + bsuccess,error = await task.put_one_task() + return bsuccess,error + else: + return False,"task_id值存在问题！" + + #获取节点待执行任务 + def get_task_node_todo_instr(self,task_id,nodepath): + todoinstr = [] + task = self.tasks[task_id] + if task: + node = task.attack_tree.find_node_by_nodepath(nodepath) + if node: + todoinstr = node.get_instr_user() + return todoinstr + + #获取节点的MSG信息 + def get_task_node_MSG(self,task_id,nodepath): + task = self.tasks[task_id] + if task: + node = task.attack_tree.find_node_by_nodepath(nodepath) + if node: + tmpMsg = node.get_res_user() + if tmpMsg: + return node.messages,tmpMsg[0] #待提交消息正常应该只有一条 + else: + return node.messages,{} + return [],{} + + def update_node_MSG(self,task_id,nodepath,newtype,newcontent): + task = self.tasks[task_id] + if task: + node = task.attack_tree.find_node_by_nodepath(nodepath) + if node: + work_status = node.get_work_status() + if work_status == 0 or work_status == 3: + bsuccess,error = node.updatemsg(newtype,newcontent,0) #取的第一条，也就修改第一条 + return bsuccess,error + else: + return False,"当前节点的工作状态不允许修改MSG！" + return False,"找不到对应节点！" + + def del_node_instr(self,task_id,nodepath,instr): + task = self.tasks[task_id] + if task: + node = task.attack_tree.find_node_by_nodepath(nodepath) + if node: + return node.del_instr(instr) + return False,"找不到对应节点！" + + + #------------以下函数还未验证处理----------- + + def start_task(self, task_id): + """启动指定任务""" + task = self.tasks.get(task_id) + if task: + task.start(self.num_threads) + else: + print(f"Task {task_id} not found") + + def stop_task(self, task_id): + """停止指定任务""" + task = self.tasks.get(task_id) + if task: + task.stop() + else: + print(f"Task {task_id} not found") + +g_TaskM = TaskManager() #单一实例 \ No newline at end of file diff --git a/mycode/TaskObject.py b/mycode/TaskObject.py new file mode 100644 index 0000000..a735f27 --- /dev/null +++ b/mycode/TaskObject.py @@ -0,0 +1,578 @@ +''' +渗透测试任务管理类 一次任务的闭合性要检查2025-3-10 一次任务后要清理LLM和InstrM的数据 +''' +from mycode.TargetManager import TargetManager # 从模块导入类 +#from LLMManager import LLMManager # 同理修正其他导入 +from mycode.ControlCenter import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。 +from mycode.InstructionManager import g_instrM +from mycode.DBManager import DBManager,app_DBM +from mycode.LLMManager import LLMManager +from mycode.AttackMap import AttackTree,TreeNode +from myutils.MyTime import get_local_timestr +from myutils.MyLogger_logger import LogHandler +from myutils.PickleManager import g_PKM +from myutils.ConfigManager import myCongif +from mycode.WebSocketManager import g_WSM +import asyncio +import queue +import time +import os +import re +import threading +import json + +class TaskObject: + + def __init__(self,test_target,cookie_info,work_type,llm_type,num_threads,local_ip,taskM): + #功能类相关 + self.taskM = taskM + self.logger = LogHandler().get_logger("TaskObject") + self.InstrM = g_instrM # 类对象渗透，要约束只读取信息，且不允许有类全局对象--持续检查 + self.CCM = ControlCenter() #一个任务一个CCM + self.LLM = LLMManager(llm_type) # LLM对象调整为一个任务一个对象，这样可以为不同的任务选择不同的LLM + #全局变量 + self.app_work_type = myCongif.get_data("App_Work_type") #app工作为0时，只允许单步模式工作，是附加规则，不影响正常逻辑处理 + self.brun = False #任务的停止可以用该变量来控制 + self.sleep_time = myCongif.get_data("sleep_time") + self.target = test_target + self.cookie = cookie_info + self.work_type = work_type #工作模式 0-人工，1-自动 + self.task_id = None + self.task_status = 0 #0-暂停，1-执行中，2-已完成 + 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.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_node_queue = queue.Queue() #待提交LLM的节点队列 + #自检线程-------- + self.check_th = None #自检线程句柄 + + #---------------三个线程------------ + #测试指令执行线程 + def do_worker_th(self): + #线程的dbm需要一个线程一个 + th_DBM = DBManager() + th_DBM.connect() + while self.brun: + if self.task_status == 1: + try: + work_node = self.instr_node_queue.get(block=False)#正常一个队列中一个节点只会出现一次，进而也就只有一个线程在处理 + # 开始执行指令 + results = [] + while True: + instruction = work_node.get_instr() + if not instruction: + break + start_time = get_local_timestr() # 指令执行开始时间 + instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction) + 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 + th_DBM.insetr_result(self.task_id, instr, reslut, work_node.do_sn, start_time, end_time, + source_result, + ext_params, work_node.path) + else: + self.logger.error("数据库连接失败！！") + #暂存结果 + oneres = {'执行指令': instr, '结果': reslut} + results.append(oneres) + #指令都执行结束后，入节点待提交队列 + str_res = json.dumps(results, ensure_ascii=False) + # 提交llm待处理任务 --更新节点work_status + self.put_node_reslist(work_node, str_res, 1) + # 保存记录 + g_PKM.WriteData(self.attack_tree,str(self.task_id)) + except queue.Empty: + time.sleep(self.sleep_time) + else:#不是工作状态则休眠 + time.sleep(self.sleep_time) + + #llm请求提交线程 + def th_llm_worker(self): + ''' + 几个规则--TM的work线程同 + 1.线程获取一个节点后，其他线程不能再获取这个节点（遇到被执行的节点，直接放弃执行）--- 加了没办法保存中间结果进行测试 + 2.llm返回的指令，只可能是该节点，或者是子节点的，不符合这条规则的都不处理，避免llm处理混乱。 + :return: + ''' + # 线程的dbm需要一个线程一个 + th_DBM = DBManager() + th_DBM.connect() + while self.brun: + if self.task_status == 1: + try: + llm_node = self.llm_node_queue.get(block=False) #获取一个待处理节点 + #开始处理 + tmp_commands = [] + # {llm_node.status} --- 暂时固化为未完成 + user_Prompt = f''' +当前分支路径：{llm_node.path} +当前节点信息： +- 节点名称：{llm_node.name} +- 节点状态：未完成 +- 漏洞类型：{llm_node.vul_type} + ''' + while True: + llm_data = llm_node.get_res() + if llm_data is None: + break + llm_type = llm_data["llm_type"] + str_res = llm_data["result"] + #获取提示词 + prompt = self.get_llm_prompt(llm_type,str_res,user_Prompt) + # 提交llm请求返回数据--并对返回数据进行处理，节点指令直接执行，测试指令根据工作模式执行 + node_cmds, commands,reasoning_content, content, post_time = self.LLM.get_llm_instruction(prompt,llm_node) # message要更新 + # LLM记录存数据库 + if th_DBM.ok: + llm_node.llm_sn += 1 + bres = th_DBM.insert_llm(self.task_id, prompt, reasoning_content, content, post_time, llm_node.llm_sn,llm_node.path) + if not bres: + self.logger.error(f"{llm_node.name}-llm入库失败！") + else: + self.logger.error("数据库连接失败！") + ''' + 对于LLM返回的错误处理机制 + 1.验证节点是否都有测试指令返回 + 2.LLM的回复开始反复时（有点难判断） + ''' + # 更新tree + bok, new_commands = 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) + #一个节点完成，节点树持久化---待验证是否有局部更新持久化的方案 + g_PKM.WriteData(self.attack_tree,str(self.task_id)) + except queue.Empty: + self.logger.debug("llm队列中暂时无新的提交任务！") + time.sleep(self.sleep_time) + else: + time.sleep(self.sleep_time) + + #自检线程 + def th_check(self): + print("自检线程待实现中！") + + #------------入两个nodeMQ-禁止直接调用入队列----------- + def put_instr_mq(self,node): + #这里不做状态的判断，调用前处理 + self.instr_node_queue.put(node) + self.update_node_work_status(node,2) #在执行--1.work_status不影响整个任务的执行，错了问题不大，2--attack_tree持久化需要出去lock信息。 + + def put_llm_mq(self,node): + #同instr_mq + self.llm_node_queue.put(node) + self.update_node_work_status(node,4) #提交中 + + async def put_instr_mq_async(self,node): + #这里不做状态的判断，调用前处理 + self.instr_node_queue.put(node) + await self.update_node_work_status_async(node,2) #在执行--1.work_status不影响整个任务的执行，错了问题不大，2--attack_tree持久化需要出去lock信息。 + + async def put_llm_mq_async(self,node): + #同instr_mq + self.llm_node_queue.put(node) + await self.update_node_work_status_async(node,4) #提交中 + + async def update_node_work_status_async(self,node,work_status): + #更新状态 + bchange = node.update_work_status(work_status) + #基于websocket推送到前端 + if bchange: + #判断是否是web端最新获取数据的task + if self.taskM.web_cur_task == self.task_id: + idatatype = 1 + strdata = {"node_path":node.path,"node_workstatus":work_status} + await g_WSM.send_data(idatatype,strdata) + + #------------入Node的两个list--禁止直接调用入list------- + def put_node_reslist(self, node, str_res, llm_type): + # 推送llm提交任务到节点的待处理任务中，并根据工作模式判断是否如llm_node_quere + one_llm = {'llm_type': llm_type, 'result': str_res} + node.add_res(one_llm) # 入节点结果队列 + self.update_node_work_status(node,3) #待提交llm + # 如果是自动执行的模式则入队列交给llm线程处理 + if self.app_work_type == 1: + if self.work_type == 1 and node.bwork: + self.put_llm_mq(node) #变4 + + def put_node_instrlist(self, commands, node): #如果当前节点没有进一般指令返回，需要修改节点执行状态 + node_list = [] #一次返回的测试指令 + for command in commands: + # 使用正则匹配方括号中的node_path（非贪婪模式） + match = re.search(r'\[(.*?)\]', command) + if match: + node_path = match.group(1) + #'''强制约束，不是本节点或者是子节点的指令不处理''' + find_node = self.attack_tree.find_node_by_nodepath_parent(node_path,node) + 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: + self.logger.error(f"基于节点路径没有找到对应的节点{node_path},父节点都没找到！当前节点{node.path}")#丢弃该指令 + else: + self.logger.error(f"得到的指令格式不符合规范：{command}")#丢弃该指令--- + #这里对于丢弃指令，有几种方案： + # 1.直接丢弃不处理，但需要考虑会不会产生节点缺失指令的问题，需要有机制验证节点；------ 需要有个独立线程，节点要加锁--首选待改进方案 + # 2.入当前节点的res_queue,但有可能当前节点没有其他指令，不会触发提交，另外就算提交，会不会产生预设范围外的返回，不确定； + # 3.独立队列处理 + #判断当前节点是否有指令 + if node not in node_list: + #修改该节点状态为0--无待执行任务 + self.update_node_work_status(node,0) + #入instr队列 + if self.app_work_type == 1: + if self.work_type == 1: #是自动执行模式 + for node in node_list: + if node.bwork: + self.put_instr_mq(node) #2-执行中 + + def put_work_node(self): + '''遍历节点需要处理的任务，提交mq''' + nodes = self.attack_tree.traverse_bfs() + for node in nodes: + if not node.is_instr_empty(): #待执行指令有值 + if not node.is_llm_empty(): + self.logger.error(f"{node.path}即存在待执行指令，还存在待提交的llm，需要人工介入！！") + else: + if node.bwork: + self.put_instr_mq(node) #提交执行 + elif not node.is_llm_empty(): #待提交llm有值 + if not node.is_instr_empty(): + self.logger.error(f"{node.path}即存在待执行指令，还存在待提交的llm，需要人工介入！！") + else: + if node.bwork: + self.put_llm_mq(node) #提交执行 + + #web端提交单步任务--节点单步 + async def put_one_node(self,node): + #提交某个节点的代表任务 + if self.task_status ==1 and self.work_type==0 and node.bwork: + if not node.is_instr_empty(): #待执行指令有值 + if not node.is_llm_empty(): + self.logger.error(f"{node.path}即存在待执行指令，还存在待提交的llm，需要人工介入！！") + return False,"该节点的待执行任务数据不正确，请联系管理员！" + else: + if node.bwork: + await self.put_instr_mq_async(node) #提交执行 + elif not node.is_llm_empty(): #待提交llm有值 + if not node.is_instr_empty(): + self.logger.error(f"{node.path}即存在待执行指令，还存在待提交的llm，需要人工介入！！") + return False, "该节点的待执行任务数据不正确，请联系管理员！" + else: + if node.bwork: + await self.put_llm_mq_async(node) #提交执行 + else: + await self.update_node_work_status_async(node,0) #只是修补措施，保障状态的一致性 + return False,"当前节点没有待执行任务！" + return True,"已提交单步任务" + else: + return False,"当前的任务或节点状态不允许执行单步，请检查！" + + #web端提交任务单步--任务单步 + async def put_one_task(self): + if self.task_status == 1 and self.work_type == 0: + nodes = self.attack_tree.traverse_bfs() + for node in nodes: + await self.put_one_node(node) + return True,"已提交单步任务" + else: + return False,"当前的任务状态不允许执行单步，请检查！" + + #修改节点的执行状态，并需要基于websocket推送到前端显示 同步线程调用 + def update_node_work_status(self,node,work_status): + #更新状态 + bchange = node.update_work_status(work_status) + #基于websocket推送到前端 + if bchange: + #判断是否是web端最新获取数据的task + if self.taskM.web_cur_task == self.task_id: + idatatype = 1 + strdata = {"node_path":node.path,"node_workstatus":work_status} + asyncio.run(g_WSM.send_data(idatatype,strdata)) + + #获取本次的提交提示词 + def get_llm_prompt(self,llm_type,str_res,user_Prompt): + if llm_type == 0: + ext_Prompt = f''' +初始信息：{str_res} +任务：请开始对该目标的渗透测试工作。 +''' + elif llm_type == 1: # 提交指令执行结果 --- 正常提交 + # 构造本次提交的prompt + ext_Prompt = f''' +上一步结果：{str_res} +任务：生成下一步渗透测试指令或判断是否完成该节点测试。 +''' + elif llm_type == 2: # llm返回的指令存在问题，需要再次请求返回 + ext_Prompt = f''' +反馈类型：节点指令格式错误 +错误信息：{str_res} +任务：请按格式要求重新生成该节点上一次返回中生成的所有指令。 +''' + elif llm_type ==3: #对节点没有指令的，请求指令 + ext_Prompt = f''' +任务：{str_res}。 +''' + elif llm_type == 5: + ext_Prompt = f''' +反馈类型：测试指令格式错误 +错误信息：{str_res} +任务：请根据格式要求，重新生成该测试指令。 +''' + else: + self.logger.debug("意外的类型参数") + return "" + + user_Prompt = user_Prompt + ext_Prompt + return user_Prompt + + #处理节点指令 + def tree_manager(self,node_cmds,node,commands,DBM): + '''更新渗透测试树 + node_cmds是json-list + 2025-03-22添加commands参数，用于处理LLM对同一个节点返回了测试指令，但还返回了no_instruction节点指令 + ''' + if not node_cmds: # or len(node_cmds)==0: 正常not判断就可以有没有节点指令 + return True,commands + + #对节点指令进行校验 + bok,strerror = self.CCM.verify_node_cmds(node_cmds) + if not bok: #节点指令存在问题，则不进行后续处理，提交一个错误反馈任务 + # 提交llm待处理任务 + self.put_node_reslist(node, strerror, 2) + return False,commands + + #先执行add_node操作 + residue_node_cmds = [] + badd_node = False + for node_json in node_cmds: + action = node_json["action"] + if action == "add_node": # 新增节点 + badd_node = True + parent_node_name = node_json["parent"] + status = node_json["status"] + 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的传递待验证 + 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) + else: + self.logger.error(f"添加子节点时，遇到父节点名称不一致的，需要介入！！{node_json}") # 丢弃该节点 + else:#其他指令添加到list + residue_node_cmds.append(node_json) + + if badd_node and self.taskM.web_cur_task == self.task_id: #如果新增了节点，且该节点树是当前查看的数据，需要通知前端更新数据 + idatatype = 2 + strdata = "update accack_tree!" + asyncio.run(g_WSM.send_data(idatatype, strdata)) + #先取消当前task,已经通知前端重新获取,这样可以避免后端不必要的数据推送 + self.taskM.web_cur_task = 0 + + #执行剩余的节点指令--不分先后 + for node_json in residue_node_cmds: + action = node_json["action"] + if action == "update_status": + 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 + 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: + str_add = {"未新增的节点": nodes} + self.logger.debug(str_add) + # 提交一个继续反馈任务--继续后续工作 2025-3-25不自动处理 + # self.put_one_llm_work(str_add, node, 4) + # self.logger.debug(f"未新增的节点有：{nodes}") + else: + self.logger.error("****不应该执行到这！程序逻辑存在问题！") + return True,commands + + #阻塞轮询补充指令 + def get_other_instruction(self,nodes,DBM,cur_node): + res_str = ','.join(nodes) + new_commands = [] + while res_str: + self.logger.debug(f"开始针对f{res_str}这些节点请求测试指令") + user_Prompt = f''' +当前分支路径：{cur_node.path} +当前节点信息： +- 节点名称：{cur_node.name} +- 节点状态：{cur_node.status} +- 漏洞类型：{cur_node.vul_type} +反馈类型：需要补充以下子节点的测试指令:{res_str} +任务： +1.请生成这些子节点的测试指令,注意不要生成重复的测试指令； +2.这些节点的父节点为当前节点，请正确生成这些节点的节点路径； +3.只有当还有节点未能生成测试指令或不完整时，才返回未生成指令的节点列表。 + ''' + res_str = "" + #node_cmds, commands = self.LLM.get_llm_instruction(user_Prompt, DBM, cur_node) # message要更新 + node_cmds, commands, reasoning_content, content, post_time = self.LLM.get_llm_instruction(user_Prompt, + cur_node) # message要更新 + # LLM记录存数据库 + cur_node.llm_sn += 1 + bres = DBM.insert_llm(self.task_id, user_Prompt, reasoning_content, content, post_time, cur_node.llm_sn,cur_node.path) + if not bres: + self.logger.error(f"{cur_node.name}-llm入库失败！") + #把返回的测试指令进行追加 + new_commands.extend(commands) + #判断是否还有未添加指令的节点 + for node_json in node_cmds: #正常应该只有一条no_instruction --暂时只处理 + if "no_instruction" in node_json and "nodes" in node_json: + tmp_nodes = [] + node_names = node_json["nodes"].split(',') + for node_name in node_names: + if node_name in nodes: + tmp_nodes.append(node_name) + res_str = ','.join(tmp_nodes) + break + self.logger.debug("未添加指令的节点，都已完成指令的添加！") + return new_commands + + def start_task(self,task_id,task_status=1,attack_tree=None): + self.task_id = task_id + ''' + 启动该测试任务 + ''' + #判断目标合法性 + # bok,target,type = self.TargetM.validate_and_extract(self.target) #是否还需要判断待定? + # if not bok: + # return False, "{target}检测目标不合规，请检查！" + #初始化节点树 + if attack_tree:#有值的情况是load + self.attack_tree =attack_tree + #加载未完成的任务 + if self.work_type ==1:#自动模式 + #提交到mq,待线程执行 + self.put_work_node() + else: #无值的情况是new_create + root_node = TreeNode(self.target, self.task_id) #根节点 + self.attack_tree = AttackTree(root_node) #创建测试树，同时更新根节点相关内容 + self.LLM.build_initial_prompt(root_node) #对根节点初始化system-msg + #插入一个user消息 + # 提交第一个llm任务，开始工作 + know_info = f"本测试主机的IP地址为:{self.local_ip}" + if self.cookie: + know_info = know_info + f",本站点的cookie值为{self.cookie}" + self.put_node_reslist(root_node,know_info,0) #入待提交list + #初始保存个attack_tree文件 + g_PKM.WriteData(self.attack_tree,str(self.task_id)) + #启动工作线程 + self.task_status = task_status + self.brun = True #线程正常启动 + #启动指令工作线程 + for i in range(self.max_thread_num): + w_th = threading.Thread(target=self.do_worker_th) + w_th.start() + self.workth_list.append(w_th) + #启动llm提交线程--llm暂时单线程，多线程处理时attack_tree需要加锁 + l_th = threading.Thread(target=self.th_llm_worker) + l_th.start() + self.llmth_list.append(l_th) + #启动自检线程 + self.check_th = threading.Thread(target=self.th_check) + self.check_th.start() + + def stop_task(self): #还未处理 + self.brun = False + 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/mycode/WebSocketManager.py b/mycode/WebSocketManager.py new file mode 100644 index 0000000..75d8cab --- /dev/null +++ b/mycode/WebSocketManager.py @@ -0,0 +1,49 @@ +import json +import struct +from quart import websocket + +class WebSocketManager: + def __init__(self): + self.ws_clients={} + + async def register(self, user_id, ws_proxy): + ws = ws_proxy._get_current_object() # 获取代理背后的真实对象 + self.ws_clients[user_id] = ws + + async def unregister(self, user_id=0): + if user_id in self.ws_clients: + del self.ws_clients[user_id] + print(f"{user_id}-socket退出") + + def create_binary_error_message(self,error_text, idatatype=0): + # 将错误信息包装成 JSON + body = json.dumps({"error": error_text}).encode('utf-8') + idata_len = len(body) + # 构造数据头：固定字符串 "TFTF"，后面 8 字节分别为 idatatype 和 idata_len (使用大端字节序) + header = b"TFTF" + struct.pack("!II", idatatype, idata_len) + return header + body + + async def send_data(self,idatatype,data,user_id=0): + """ + 发送数据格式： + - 固定 4 字节："TFTF"（ASCII） + - 接下来 8 字节为数据头，包含两个无符号整数（每个 4 字节），分别为 idatatype 与 idata_len + - 最后 idata_len 字节为数据体（JSON字符串），数据体中包含 node_path 和 node_workstatus + """ + if user_id not in self.ws_clients: + return + ws = self.ws_clients[user_id] + + # 将 data 转换为 JSON 字符串并转换为字节 + body = json.dumps(data).encode('utf-8') + idata_len = len(body) + # 使用 struct.pack 构造数据头（大端字节序） + header = b"TFTF" + struct.pack("!II", idatatype, idata_len) + message = header + body + try: + await ws.send(message) + except Exception as e: + print(f"发送失败: {e}") + await self.unregister(user_id) # 异常时自动注销连接 + +g_WSM = WebSocketManager() \ No newline at end of file diff --git a/myutils/MyTime.py b/myutils/MyTime.py index 49926af..40ad33b 100644 --- a/myutils/MyTime.py +++ b/myutils/MyTime.py @@ -9,4 +9,14 @@ def get_local_timestr() -> str: # 转换为结构化时间对象，再格式化为字符串 struct_time = time.localtime(timestamp) # 本地时区时间‌:ml-citation{ref="3,5" data="citationList"} formatted_time = time.strftime("%Y-%m-%d %H:%M:%S",struct_time) - return formatted_time \ No newline at end of file + return formatted_time + + + +if __name__ =="__main__": + list_text = [] + if not list_text: + list_text.append("111") + else: + list_text[0]= "111" + print(list_text) \ No newline at end of file diff --git a/myutils/PickleManager.py b/myutils/PickleManager.py new file mode 100644 index 0000000..6548d39 --- /dev/null +++ b/myutils/PickleManager.py @@ -0,0 +1,32 @@ +import pickle +import threading +from myutils.ConfigManager import myCongif + +class PickleManager: + def __init__(self): + self.lock = threading.Lock() # 线程锁 + self.tree_file = myCongif.get_data("TreeFile") + + def WriteData(self,attack_tree,filename=""): + if filename: + filepath = "tree_data/"+filename + else: + filepath = self.tree_file + + 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 + + with self.lock: + with open(filepath, "rb") as f: + attack_tree = pickle.load(f) + return attack_tree + +g_PKM = PickleManager() \ No newline at end of file diff --git a/pipfile b/pipfile index 5260e24..14b9cee 100644 --- a/pipfile +++ b/pipfile @@ -30,6 +30,12 @@ pip install quart-cors -i https://pypi.tuna.tsinghua.edu.cn/simple pip install pymemcache -i https://pypi.tuna.tsinghua.edu.cn/simple pip install quart-sqlalchemy -i https://pypi.tuna.tsinghua.edu.cn/simple pip install pillow -i https://pypi.tuna.tsinghua.edu.cn/simple +#redis--session使用redis缓存--kali +pip install redis aioredis -i https://pypi.tuna.tsinghua.edu.cn/simple +pip install quart-session -i https://pypi.tuna.tsinghua.edu.cn/simple +systemctl start redis-server +systemctl enable redis-server + ---arial.ttf字体--- diff --git a/run.py b/run.py new file mode 100644 index 0000000..ab91a23 --- /dev/null +++ b/run.py @@ -0,0 +1,54 @@ +import asyncio +import os +from mycode.TaskManager import g_TaskM +from web import create_app +from hypercorn.asyncio import serve +from hypercorn.config import Config + +async def run_quart_app(): + app = create_app() + config = Config() + config.bind = ["0.0.0.0:5001"] + config.use_reloader = True # 启用热重载 + config.reload_include_patterns = ["*.py", "templates/*", "static/*"] # 监控模板和静态文件 + 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) + diff --git a/tools/MsfconsoleTool.py b/tools/MsfconsoleTool.py index da8ccbe..d0b3cd7 100644 --- a/tools/MsfconsoleTool.py +++ b/tools/MsfconsoleTool.py @@ -9,8 +9,8 @@ class MsfconsoleTool(ToolBase): ''' Metasploit 这样的交互工具，保持一个会话，析构时结束会话 ''' - def __init__(self,TM): - super().__init__(TM) + def __init__(self): + super().__init__() self.bOK = False self.msf_lock = threading.Lock() # 线程锁 diff --git a/tools/RpcinfoTool.py b/tools/RpcinfoTool.py new file mode 100644 index 0000000..4469875 --- /dev/null +++ b/tools/RpcinfoTool.py @@ -0,0 +1,11 @@ +from tools.ToolBase import ToolBase + +class RpcinfoTool(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/ToolBase.py b/tools/ToolBase.py index 284bba6..18f0e61 100644 --- a/tools/ToolBase.py +++ b/tools/ToolBase.py @@ -14,10 +14,11 @@ import sys from myutils.ReturnParams import ReturnParams class ToolBase(abc.ABC): - def __init__(self,TM): + def __init__(self): #self.res_type = 0 #补充一个结果类型 0-初始状态,1-有安全问题, #由于工具类会被多个线程调用，全局变量不能修改，只能读取 - self.TM = TM + pass + def create_extparams(self): diff --git a/web/API/__init__.py b/web/API/__init__.py index 92cdafb..2954561 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 +from . import user,task,wsm diff --git a/web/API/task.py b/web/API/task.py new file mode 100644 index 0000000..52434dc --- /dev/null +++ b/web/API/task.py @@ -0,0 +1,194 @@ +from . import api +from quart import Quart, render_template, redirect, url_for, request,jsonify +from mycode.TargetManager import g_TM +from mycode.DBManager import app_DBM +from mycode.TaskManager import g_TaskM + + +def is_valid_target(test_target: str) -> bool: + """ + 验证 test_target 的逻辑（这里用简单示例代替已有逻辑） + 例如：测试目标不能为空且长度大于3 + """ + if test_target: + return True + return False + +@api.route('/task/start',methods=['POST']) +async def start_task(): #开始任务 + data = await request.get_json() + test_target = data.get("testTarget") + cookie_info = data.get("cookieInfo") + llm_type = data.get("curmodel") # //0-腾讯云，1-DS，2-2233.ai,3-GPT 目前只有1-2，2025-4-4 + work_type = data.get("workType") #0-人工，1-自动 + #新增任务处理 + bok,_,_ = g_TM.validate_and_extract(test_target) + if not bok: + # 返回错误信息，状态码 400 表示请求错误 + return jsonify({"error": "测试目标验证失败，请检查输入内容！"}), 400 + #开始任务 + try: + b_success = g_TaskM.create_task(test_target,cookie_info,llm_type,work_type) + #再启动 + if not b_success: + return jsonify({"error": "检测任务创建失败，请联系管理员！"}), 500 + except: + return jsonify({"error": "该目标已经在测试中，请检查！"}), 400 + #跳转到任务管理页面 + return redirect(url_for('main.get_html', html='task_manager.html')) + +@api.route('/task/getlist',methods=['GET']) +async def get_task_list(): + #task_list = app_DBM.get_task_list() #从内存取--2025-4-6 + task_list = g_TaskM.get_task_list() + if task_list: + return jsonify(task_list) + else: + return jsonify({"error":"查询任务数据出错！"}),500 + +@api.route('/task/getinstr',methods=['POST']) +async def get_instr(): + data = await request.get_json() + task_id = data.get("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) + +@api.route('/task/getvul',methods=['POST']) +async def get_vul(): + data = await request.get_json() + task_id = data.get("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) + +@api.route('/task/gettree',methods=['POST']) +async def get_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_node_tree(task_id) + return jsonify({"tree":tree_dict}) + +@api.route('/task/taskcontrol',methods=['POST']) +async def task_status_control(): + '''控制任务状态 + 1.对于执行时间长的指令，如何处理？强制停止的话，要有个执行中指令的缓存，强制停止该指令返回到待执行，执行完成，该指令到执行完成； + ''' + data = await request.get_json() + task_id = data.get("cur_task_id") + if not task_id: + return jsonify({'error': 'Missing task_id'}), 400 + #只做暂停和继续间的切换，以服务器端的状态为准 + bsuccess,strerror,new_task_status = g_TaskM.control_taks(task_id) + if bsuccess: + return jsonify({'newstatus':new_task_status}) + return jsonify({'error': strerror}), 400 + +@api.route('/task/taskstep',methods=['POST']) +async def task_one_step(): + '''单步推进任务--也就是待处理node 返回bsuccess,error + 1.执行单步的前提条件是，工作线程都要在工作； + 2.遍历节点把需要处理的节点进入待处理queue，instr和llm只能一个有数据（强制约束） + ''' + 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 = await g_TaskM.task_one_step(task_id) + return jsonify({"bsuccess":bsuccess,"error":error}) + +@api.route('/task/nodestep',methods=['POST']) +async def node_one_step(): + data = await request.get_json() + task_id = data.get("task_id") + node_path = data.get("node_path") + if not task_id: + return jsonify({'error': 'Missing task_id'}), 400 + bsuccess,error = await g_TaskM.node_one_step(task_id,node_path) + return jsonify({"bsuccess":bsuccess,"error":error}) + +@api.route('/task/taskworktype',methods=['POST']) +async def task_work_type_control(): + data = await request.get_json() + task_id = data.get("cur_task_id") + newwork_type = data.get("mode") + if not task_id: + return jsonify({'error': 'Missing task_id or newwork_type'}), 400 + bsuccess = g_TaskM.update_task_work_type(task_id,newwork_type) + return jsonify({"bsuccess": bsuccess}) + +@api.route('/task/nodecontrol',methods=['POST']) +async def node_work_status_control(): + data = await request.get_json() + task_id = data.get("task_id") + nodepath = data.get("node_path") + if not task_id or not nodepath: + return jsonify({'error': 'Missing task_id or node_path'}), 400 + #修改节点的工作状态 + bsuccess,newbwork = g_TaskM.node_bwork_control(task_id,nodepath) + if not bsuccess: + return jsonify({'error': 'node_path not bfind'}), 400 + return jsonify({"newbwork":newbwork}) + +@api.route('/task/nodegetinstr',methods=['POST']) +async def 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,"todoInstrs":todoInstrs}) + +@api.route('/task/nodegetmsg',methods=['POST']) +async def node_get_msg(): + 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 + submitted,pending = g_TaskM.get_task_node_MSG(task_id,nodepath) + return jsonify({"submitted": submitted, "pending": pending}) + +@api.route('/task/nodeupdatemsg',methods=['POST']) +async def node_update_msg(): + data = await request.get_json() + task_id = data.get("task_id") + nodepath = data.get("node_path") + newllm_type = data.get("llmtype") + newcontent = data.get("content") + if not task_id: + return jsonify({'error': 'Missing task_id'}), 400 + bsuccess,error =g_TaskM.update_node_MSG(task_id,nodepath,newllm_type,newcontent) + return jsonify({"bsuccess":bsuccess,"error":error}) + +@api.route('/task/delnodeinstr',methods=['POST']) +async def node_del_instr(): + data = await request.get_json() + task_id = data.get("task_id") + nodepath = data.get("node_path") + instr = data.get("item") + if not task_id: + return jsonify({'error': 'Missing task_id'}), 400 + bsuccess,error = g_TaskM.del_node_instr(task_id,nodepath,instr) + return jsonify({"bsuccess": bsuccess, "error": error}) + + + + + + + + + + diff --git a/web/API/user.py b/web/API/user.py index 1c6f683..fabaadb 100644 --- a/web/API/user.py +++ b/web/API/user.py @@ -1,5 +1,6 @@ import os import hashlib +from mycode.DBManager import app_DBM from quart import Quart, render_template, request, session, redirect, url_for,jsonify,send_file,flash from quart_sqlalchemy import SQLAlchemy from quart_session import Session @@ -38,13 +39,13 @@ async def user_login(): #用户登录 #return 'captcha error!', 400 #比对用户名和密码 strsql = f"select password from user where username = '{username}'" - db_password = mDBM.do_select(strsql,1) + db_password = app_DBM.do_select(strsql,1) passwd_md5 = get_md5(password) if db_password: if db_password[0] == passwd_md5: #后续需要对密码进行MD5加默 print("登录成功") session['user'] = username - return redirect(url_for('main.get_html', html='view_main.html')) + return redirect(url_for('main.get_html', html='index.html')) await flash('用户名或密码错误', 'error') return redirect(url_for('main.login')) @@ -52,7 +53,7 @@ async def user_login(): #用户登录 @login_required async def user_info(): #获取用户列表 strsql = "select username,status,people,tellnum from user;"; - data = mDBM.do_select(strsql) + data = app_DBM.do_select(strsql) if data: user_list = [{"username": user[0], "status": user[1], "people":user[2],"tellnum":user[3]} for user in data] @@ -69,14 +70,14 @@ async def user_adduser(): #新增用户 tellnum = (await request.form)['tellnum'] strsql = f"select username from user where username = '{username}';" password = myCongif.get_data('pw') - data = mDBM.do_select(strsql) + data = app_DBM.do_select(strsql) if data: reStatus = 0 reMsg = '用户名重复，请重新输入！' else: strsql = (f"INSERT INTO user (username ,password ,status,people,tellnum ) VALUES " f"('{username}','{password}',1,'{people}','{tellnum}');") - ret = mDBM.do_sql(strsql) + ret = app_DBM.do_sql(strsql) if ret == True: reStatus = 1 reMsg = '添加用户成功' @@ -94,12 +95,12 @@ async def user_change_passwd(): #修改密码 old_md5= get_md5(oldpasswd) print(old_md5) strsql = f"select id from user where password='{old_md5}';" - data = mDBM.do_select(strsql,1) + data = app_DBM.do_select(strsql,1) reStatus = 0 if data: new_md5 = get_md5(newpasswd) strsql = f"update user set password = '{new_md5}' where password = '{old_md5}';" - ret = mDBM.do_sql(strsql) + ret = app_DBM.do_sql(strsql) if ret: reStatus = 1 reMsg = '修改密码成功' @@ -117,7 +118,7 @@ async def user_change_user_info(): #修改用户信息 people = (await request.form)['people'] tellnum = (await request.form)['tellnum'] strsql = f"update user set people='{people}',tellnum='{tellnum}' where username='{username}';" - ret = mDBM.do_sql(strsql) + ret = app_DBM.do_sql(strsql) if ret == True: reStatus = 1 reMsg = '修改用户信息成功' diff --git a/web/API/wsm.py b/web/API/wsm.py new file mode 100644 index 0000000..f824c3b --- /dev/null +++ b/web/API/wsm.py @@ -0,0 +1,43 @@ +import json +from . import api +from quart import Quart, websocket, jsonify +from mycode.WebSocketManager import g_WSM + +# WebSocket 路由，端口默认与 HTTP 同端口，例如 5000（开发时） +@api.websocket("/ws") +async def ws(): + """ + WebSocket 连接入口： + 1. 客户端连接成功后，首先应发送登录数据包，例如 {"user_id": 1} + 2. 后端解析登录数据包，将 user_id 与 websocket 绑定（注册） + 3. 后续进入消息接收循环，根据数据协议（TFTF+头+体格式）处理数据 + """ + # 接收登录数据包（假设为纯 JSON 包，非二进制格式） + login_msg = await websocket.receive() + try: + login_data = json.loads(login_msg) + user_id = login_data.get("user_id") + if user_id is None: + error_msg = g_WSM.create_binary_error_message("没有提供user_id") + await websocket.send(error_msg) + return + # 注册当前连接 + await g_WSM.register(user_id, websocket) + except Exception as e: + errof_msg = g_WSM.create_binary_error_message(str(e)) + await websocket.send(errof_msg) + return + + try: + # 主循环：接收后续消息并处理；此处可以根据协议解析消息 ---暂时无web发送数据的业务 + while True: + message = await websocket.receive() + # 这里可以添加对接收到的消息解析处理，根据协议分离 TFTF 等 + print("收到消息:", message) + # 例如，直接回显消息；实际项目中应解析消息格式并调用相应的业务逻辑 + # await websocket.send(message) + except Exception as e: + print("WebSocket 出现异常:", e) + finally: + # 注销当前连接 + await g_WSM.unregister(user_id) \ No newline at end of file diff --git a/web/__init__.py b/web/__init__.py index 4c5e75d..fe7bdf0 100644 --- a/web/__init__.py +++ b/web/__init__.py @@ -1,3 +1,4 @@ +import aioredis from quart import Quart,session,redirect, url_for from quart_session import Session from quart_cors import cors @@ -33,19 +34,13 @@ class MemcachedSessionInterface: #只是能用，不明所以 def create_app(): app = Quart(__name__) app.config['SECRET_KEY'] = 'zfxxkj_2024_!@#' - - if myCongif.get_data("model_platform") == "acl": - app.config['SESSION_TYPE'] = 'memcached' # session类型 - elif myCongif.get_data("model_platform") =="cpu": - app.config['SESSION_TYPE'] = 'redis' # session类型 - + app.config["TEMPLATES_AUTO_RELOAD"] = True #动态加载模板文件 #app.config['SESSION_FILE_DIR'] = './sessions' # session保存路径 #app.config['SESSION_MEMCACHED'] = base.Client(('localhost', 11211)) app.config['SESSION_PERMANENT'] = True # 如果设置为True，则关闭浏览器session就失效。 app.config['SESSION_USE_SIGNER'] = False # 是否对发送到浏览器上session的cookie值进行加密 - - memcached_client = base.Client(('localhost', 11211)) - app.session_interface = MemcachedSessionInterface(memcached_client) + app.config['SESSION_TYPE'] = 'redis' # session类型 + app.config['SESSION_REDIS'] = aioredis.from_url('redis://localhost:6379') Session(app) # 注册main diff --git a/web/main/routes.py b/web/main/routes.py index 82a9b2b..f85a39c 100644 --- a/web/main/routes.py +++ b/web/main/routes.py @@ -14,7 +14,7 @@ def login_required(f): @wraps(f) async def decorated_function(*args, **kwargs): if 'user' not in session: - return redirect(url_for('main.index',error='未登录，请重新登录')) + return redirect(url_for('main.login',error='未登录，请重新登录')) return await f(*args, **kwargs) return decorated_function diff --git a/web/main/static/resources/css/node_tree.css b/web/main/static/resources/css/node_tree.css new file mode 100644 index 0000000..8a3481e --- /dev/null +++ b/web/main/static/resources/css/node_tree.css @@ -0,0 +1,134 @@ +/* 左侧树容器：固定高度，允许滚动，根节点居中 */ + .node-tree-area { + width: 100%; + height: 100%; /* 示例高度，根据需求调整 */ + border: 1px solid #ddd; + overflow: hidden; /* 超出时出现滚动条 */ + background-color: #f8f9fa; + text-align: center; /* 内部 inline-block 居中 */ + position: relative; + } + +/* 树节点内容区域，不包含刷新按钮 */ +.tree-content { + height: 100%; + overflow: auto; + padding-top: 5px; /* 留出顶部刷新按钮位置 */ +} + +/* 顶部刷新按钮 */ + .refresh-container { + position: absolute; + top: 5px; + left: 5px; + z-index: 100; + } + +.tree-refresh { + position: absolute; + top: 5px; + left: 5px; + z-index: 100; + cursor: pointer; + border: none; + background-color: #007bff; /* 蓝色背景 */ + color: #fff; /* 白色文字 */ + font-size: 20px; + padding: 5px 10px; + border-radius: 4px; +} +.tree-refresh:hover { + background-color: #0056b3; /* 悬停时深蓝 */ +} + + /* 让树的 ul 使用 inline-block 显示，便于根节点居中 */ + .tree-root-ul { + display: inline-block; + text-align: left; /* 子节点左对齐 */ + margin: 20px; + padding:0; + } + .tree-root-ul ul { + margin-left: 20px; /* 子节点缩进 */ + padding-left: 20px; + border-left: 1px dashed #ccc; /* 增加分隔线效果 */ + } + /* 节点容器：增加立体效果 */ + .node-container { + display: inline-flex; + align-items: center; + } + /* 切换图标 */ + .toggle-icon { + margin-right: 5px; + font-weight: bold; + cursor: pointer; + } + /* 每个节点样式 */ + .tree-node { + display: inline-block; + padding: 6px 10px; + margin: 3px 0; + border: 1px solid transparent; + border-radius: 4px; + background-color: #fff; + cursor: pointer; + transition: all 0.2s ease; + box-shadow: 2px 2px 4px rgba(0, 0, 0, 0.15); + } + .tree-node:hover { + background-color: #e6f7ff; + transform: translateY(-1px); + } + .tree-node.selected { + border-color: #1890ff; + background-color: #bae7ff; + } + /* 针对漏洞级别的样式 */ + .tree-node.vul-low { + border-color: #87d068; + background-color: #f6ffed; + } + .tree-node.vul-medium { + border-color: #faad14; + background-color: #fff7e6; + } + .tree-node.vul-high { + border-color: #ff4d4f; + background-color: #fff1f0; + } +.tree-node.no-work { + border-color: #151515; + background-color: #c4c4c4; + } + /* 收缩/展开时隐藏子节点 ul */ + .collapsed { + display: none; + } + /* 右侧信息区域 */ + .node-info-area { + border: 1px solid #ddd; + padding: 10px; + min-height: 200px; + } + /* 按钮区域 */ + .node-actions button { + margin-right: 10px; + margin-bottom: 10px; + } + + /* 限制模态框对话框的最大高度（比如距离屏幕上下各留 20px） */ +.modal-dialog { + max-height: calc(100vh - 60px); + /* 如果模态框尺寸需要自适应宽度，可以考虑设置 width: auto; */ +} + +/* 限制模态框内容区域的最大高度并启用垂直滚动 */ +.modal-content { + max-height: calc(100vh - 60px); +} + +/* 模态框主体区域启用滚动 */ +.modal-body { + overflow-y: auto; +} diff --git a/web/main/static/resources/scripts/aiortc-client-new.js b/web/main/static/resources/scripts/aiortc-client-new.js deleted file mode 100644 index ad17162..0000000 --- a/web/main/static/resources/scripts/aiortc-client-new.js +++ /dev/null @@ -1,563 +0,0 @@ -let video_list = {}; //element_id -- socket -let run_list = {}; //element_id -- runtag -let berror_state_list = {}; //element_id -- 错误信息显示 -let m_count = 0; -let connection_version = {}; // 保存每个 element_id 的版本号 -let channel_list = null; -const fourViewButton = document.getElementById('fourView'); -const nineViewButton = document.getElementById('nineView'); - -//页面即将卸载时执行 -window.addEventListener('beforeunload', function (event) { - // 关闭所有 WebSocket 连接或执行其他清理操作 - for(let key in video_list){ - const videoFrame = document.getElementById(`video-${key}`); - const event = new Event('closeVideo'); - videoFrame.dispatchEvent(event); - - delete video_list[key]; - berror_state_list[key] = false; - } -}); - -//页面加载时执行 -document.addEventListener('DOMContentLoaded', async function() { - console.log('DOM fully loaded and parsed'); - // 发送请求获取额外数据 --- 这个接口功能有点大了---暂时只是更新通道树2024-7-29 - try { - let response = await fetch('/api/channel/list'); - if (!response.ok) { - throw new Error('Network response was not ok'); - } - channel_list = await response.json(); - // 遍历输出每个元素的信息 - let area_name = "" - let html = ''; - html += ''; - } - area_name = `${channel.area_name}`; - html += `
  • ${area_name}`; - html += ''; - html += '
    • '; - } - html += ''; - const treeView = document.getElementById('treeView'); - treeView.innerHTML = html - generateVideoNodes(4); - } catch (error) { - console.error('Failed to fetch data:', error); - } - }); - -document.addEventListener('click', function() { - console.log("第一次页面点击，开始显示视频--已注释",m_count); - // if(m_count != 0){ - // count = m_count - // //获取视频接口 - // const url = `/api/viewlist?count=${count}`; - // fetch(url) - // .then(response => response.json()) - // .then(data => { - // console.log('Success:', data); - // clist = data.clist; - // elist = data.elist; - // nlist = data.nlist; - // for(let i=0;i { - // console.error('Error:', error); - // }); - // } -}, { once: true }); - -//视频窗口 -document.getElementById('fourView').addEventListener('click', function() { - if (fourViewButton.classList.contains('btn-primary')) { - return; // 如果按钮已经是选中状态，直接返回 - } - const videoGrid = document.getElementById('videoGrid'); - videoGrid.classList.remove('nine'); - videoGrid.classList.add('four'); - generateVideoNodes(4); - //更新按钮点击状态 - fourViewButton.classList.remove('btn-secondary'); - fourViewButton.classList.add('btn-primary'); - nineViewButton.classList.remove('btn-primary'); - nineViewButton.classList.add('btn-secondary'); -}); - -document.getElementById('nineView').addEventListener('click', function() { - if (nineViewButton.classList.contains('btn-primary')) { - return; // 如果按钮已经是选中状态，直接返回 - } - const videoGrid = document.getElementById('videoGrid'); - videoGrid.classList.remove('four'); - videoGrid.classList.add('nine'); - generateVideoNodes(9); - //更新按钮点击状态 - fourViewButton.classList.remove('btn-primary'); - fourViewButton.classList.add('btn-secondary'); - nineViewButton.classList.remove('btn-secondary'); - nineViewButton.classList.add('btn-primary'); -}); - -function generateVideoNodes(count) { //在这里显示视频-初始化 ---这里使用重置逻辑 - //结束在播放的socket - for(let key in video_list){ - //flv使用 - // const videoFrame = document.getElementById(`video-${key}`); - // const event = new Event('closeVideo'); - // videoFrame.dispatchEvent(event); - - //通用关闭 - run_list[key] = false; - video_list[key].close(); - berror_state_list[key] = false; - delete video_list[key]; - } - //切换窗口布局 - const videoGrid = document.getElementById('videoGrid'); - let html = ''; - for (let i = 0; i < count; i++) { - let frameWidth = count === 4 ? 'calc(50% - 10px)' : 'calc(33.33% - 10px)'; - html += ` -
    -
    -
    Video Stream ${i+1}
    -
    - - -
    -
    -
    -
    `; - } - videoGrid.innerHTML = html; - - //开始还原视频，获取视频接口 - // if(m_count != 0){ - const url = `/api/viewlist?count=${count}`; - fetch(url) - .then(response => response.json()) - .then(data => { - console.log('Success:', data); - clist = data.clist; - elist = data.elist; - nlist = data.nlist; - for(let i=0;i { - console.error('Error:', error); - }); - // } - //m_count = count -} - -function toggleFullScreen(id) { - console.log('toggleFullScreen'); - const videoFrame = document.querySelector(`[data-frame-id="${id}"]`); - if (!document.fullscreenElement) { - videoFrame.requestFullscreen().catch(err => { - alert(`Error attempting to enable full-screen mode: ${err.message} (${err.name})`); - }); - } else { - document.exitFullscreen(); - }; -} - -function allowDrop(event) { - event.preventDefault(); -} - -function drag(event) { - event.dataTransfer.setData("text", event.target.dataset.nodeId); - event.dataTransfer.setData("name", event.target.dataset.nodeName); -} - -function drop(event) { - event.preventDefault(); - const nodeId = event.dataTransfer.getData("text"); - const nodeName = event.dataTransfer.getData("name"); - const frameId = event.currentTarget.dataset.frameId; - - //需要判断下当前窗口是否已经在播放视频 - const imgElement = document.getElementById(`video-${frameId}`); - const titleElement = document.querySelector(`[data-frame-id="${frameId}"] .video-title`); - - if (titleElement.textContent !== `Video Stream ${Number(frameId)+1}`) { - showModal('请先关闭当前窗口视频，然后再播放新的视频。'); - return; - }; - //发送视频链接接口 - const url = '/api/start_stream'; - const data = {"channel_id":nodeId,"element_id":frameId}; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - const istatus = data.status; - if(istatus === 0){ - showModal(data.msg); // 使用 Modal 显示消息 - return; - } - else{ - //获取视频流 - console.log("drop触发") - connectToStream(frameId,nodeId,nodeName); - //startFLVStream(frameId,nodeId,nodeName); #基于FLV的开发程度：后端直接用RTSP流转发是有画面的，但CPU占用太高，用不了。2024-8-30 - } - }) - .catch((error) => { - showModal(`Error: ${error.message}`); // 使用 Modal 显示错误信息 - return; - }); - //console.log('retrun 只是把fetch结束，这里的代码还是会执行'); -} - -function connect(channel_id,element_id,imgcanvas,ctx,offscreenCtx,offscreenCanvas,streamUrl) { - //判断是否有重复socket，进行删除 - if(element_id in video_list) { - run_list[element_id] = false; - video_list[element_id].close(); - delete video_list[element_id]; - console.log("有历史数据未删干净！！---",element_id) //要不要等待待定 - } - // 每次连接时增加版本号 - const current_version = (connection_version[element_id] || 0) + 1; - connection_version[element_id] = current_version; - const socket = new WebSocket(streamUrl); - socket.binaryType = 'arraybuffer'; // 设置为二进制数据接收 - socket.customData = { channel_id: channel_id, element_id: element_id, - imgcanvas:imgcanvas,ctx:ctx,offscreenCtx:offscreenCtx,offscreenCanvas:offscreenCanvas, - version_id: current_version,streamUrl:streamUrl}; // 自定义属性 -- JS异步事件只能等到当前同步任务（代码块）完成之后才有可能被触发。 - //新的连接 - video_list[element_id] = socket; - run_list[element_id] = true; - berror_state_list[element_id] = false; - imgcanvas.style.display = 'block'; - - // 处理连接打开事件 - socket.onopen = function(){ - console.log('WebSocket connection established--',socket.customData.channel_id); - }; - - socket.onmessage = function(event) { - let el_id = socket.customData.element_id - let cl_id = socket.customData.channel_id - let imgcanvas = socket.customData.imgcanvas - let ctx = socket.customData.ctx - let offctx = socket.customData.offscreenCtx - let offscreenCanvas = socket.customData.offscreenCanvas - - // 转换为字符串来检查前缀 - let message = new TextDecoder().decode(event.data.slice(0, 6)); // 取前6个字节 - if (message.startsWith('frame:')){ - //如有错误信息显示 -- 清除错误信息 - if(berror_state_list[el_id]){ - console.log("清除错误信息！"); - removeErrorMessage(imgcanvas); - berror_state_list[el_id] = false; - } - // 接收到 JPG 图像数据，转换为 Blob - let img = new Image(); - let blob = new Blob([event.data.slice(6)], { type: 'image/jpeg' }); - // 将 Blob 转换为可用的图像 URL - img.src = URL.createObjectURL(blob); - //定义图片加载函数 - img.onload = function() { - imgcanvas.width = offscreenCanvas.width = img.width; - imgcanvas.height = offscreenCanvas.height = img.height; - - // 在 OffscreenCanvas 上绘制 - offctx.clearRect(0, 0, imgcanvas.width, imgcanvas.height); - offctx.drawImage(img, 0, 0, imgcanvas.width, imgcanvas.height); - // 将 OffscreenCanvas 的内容复制到主 canvas - ctx.drawImage(offscreenCanvas, 0, 0); - - // 用完就释放 - URL.revokeObjectURL(img.src); - // blob = null - // img = null - // message = null - // event.data = null - // event = null - }; - }else if(message.startsWith('error:')){ - const errorText = new TextDecoder().decode(event.data.slice(6)); // 截掉前缀 'error:' - //目前只处理一个错误信息，暂不区分 - displayErrorMessage(imgcanvas, "该视频源未获取到画面，请检查后刷新重试，默认两分钟后重连"); - berror_state_list[el_id] = true; - } - }; - - socket.onclose = function() { - let el_id = socket.customData.element_id; - let cl_id = socket.customData.channel_id; - //判断是不是要重连 - if(socket.customData.version_id === connection_version[el_id]){ - if(run_list[el_id]){ - console.log(`尝试重新连接... Channel ID: ${cl_id}`); - setTimeout(() => connect(cl_id, el_id, socket.customData.imgcanvas, - socket.customData.ctx,socket.customData.streamUrl), 1000*10); // 尝试在10秒后重新连接 - } - delete video_list[el_id]; - delete connection_version[el_id]; - delete run_list[el_id]; - } - }; - - socket.onerror = function() { - console.log(`WebSocket错误，Channel ID: ${socket.customData.channel_id}`); - socket.close(1000, "Normal Closure"); - }; -} - - -function connectToStream(element_id,channel_id,channel_name) { - console.log("开始连接视频",element_id,channel_id); - //更新控件状态--设置视频区域的标题 - const titleElement = document.querySelector(`[data-frame-id="${element_id}"] .video-title`); - titleElement.textContent = channel_name; - //视频控件 - //const imgElement = document.getElementById(`video-${element_id}`); - //imgElement.alt = `Stream ${channel_name}`; - const imgcanvas = document.getElementById(`video-${element_id}`); - const ctx = imgcanvas.getContext('2d') - // 创建 OffscreenCanvas - const offscreenCanvas = new OffscreenCanvas(imgcanvas.width, imgcanvas.height); - const offscreenCtx = offscreenCanvas.getContext('2d'); - - const streamUrl = `ws://${window.location.host}/api/ws/video_feed/${channel_id}`; - - //创建websocket连接，并接收和显示图片 - connect(channel_id,element_id,imgcanvas,ctx,offscreenCtx,offscreenCanvas,streamUrl); //执行websocket连接 -- 异步的应该会直接返回 -} - -function closeVideo(id) { - if(id in video_list) { - const imgcanvas = document.getElementById(`video-${id}`); - const titleElement = document.querySelector(`[data-frame-id="${id}"] .video-title`); - //断socket - run_list[id] = false; - video_list[id].close(); - delete video_list[id]; - //清空控件状态 - imgcanvas.style.display = 'none'; // 停止播放时隐藏元素 - titleElement.textContent = `Video Stream ${id+1}`; - removeErrorMessage(imgcanvas); - berror_state_list[id] = false; - //删记录 - const url = '/api/close_stream'; - const data = {"element_id":id}; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - console.log('Success:', data); - const istatus = data.status; - if(istatus == 0){ - showModal(data.msg); // 使用 Modal 显示消息 - return; - } - }) - .catch((error) => { - showModal(`Error: ${error.message}`); // 使用 Modal 显示错误信息 - return; - }); - } - else{ - showModal('当前视频窗口未播放视频。'); - return; - } -} - -function startFLVStream(element_id,channel_id,channel_name) { - - // 设置视频区域的标题 - const titleElement = document.querySelector(`[data-frame-id="${element_id}"] .video-title`); - titleElement.textContent = channel_name; - //获取视频 - // const imgElement = document.getElementById(`video-${element_id}`); - // imgElement.alt = `Stream ${channel_name}`; - const videoElement = document.getElementById(`video-${element_id}`); - let reconnectAttempts = 0; - const maxReconnectAttempts = 3; - const flvUrl = `ws://${window.location.host}/api/ws/video_feed/${channel_id}`; - - function initFLVPlayer() { - if (flvjs.isSupported()) { - //要避免重复播放 - if(element_id in video_list) { - closeFLVStream(element_id) - }else{ - video_list[element_id] = element_id; - berror_state_list[element_id] = true; - } - - flvPlayer = flvjs.createPlayer({ - type: 'flv', - url: flvUrl, - }); - - flvPlayer.attachMediaElement(videoElement); - flvPlayer.load(); - flvPlayer.play(); - - // 设定超时时间，例如10秒 - timeoutId = setTimeout(() => { - console.error('No video data received. Closing connection.'); - flvPlayer.destroy(); // 停止视频 - // 显示错误信息或提示 - displayErrorMessage(videoElement, "该视频源获取画面超时，请检查后刷新重试，默认两分钟后重连"); - berror_state_list[element_id] = true; - }, 130000); // 130秒 - - // 错误处理 - flvPlayer.on(flvjs.Events.ERROR, (errorType, errorDetail) => { - console.error(`FLV Error: ${errorType} - ${errorDetail}`); - if (reconnectAttempts < maxReconnectAttempts) { - reconnectAttempts++; - console.log("开始重连") - setTimeout(initFLVPlayer, 30000); // 尝试重连 - } else { - displayErrorMessage(videoElement, "重连超时，请检查后重试，可联系技术支持!"); - berror_state_list[element_id] = true; - } - }); - - // 监听播放事件，如果播放成功则清除超时计时器 - flvPlayer.on(flvjs.Events.STATISTICS_INFO, () => { - clearTimeout(timeoutId); - timeoutId = null; - removeErrorMessage(videoElement); - berror_state_list[element_id] = false; - }); - - // 关闭视频流时销毁播放器 - videoElement.addEventListener('closeVideo', () => { - if(flvPlayer){ - flvPlayer.destroy(); - videoElement.removeEventListener('closeVideo', onCloseVideo); - flvPlayer.off(flvjs.Events.ERROR); - flvPlayer.off(flvjs.Events.STATISTICS_INFO); - delete flvPlayer - } - }); - - } else { - console.error('FLV is not supported in this browser.'); - } - } - - initFLVPlayer(); -} - -// 主动关闭视频的函数 -function closeFLVStream(id) { - const titleElement = document.querySelector(`[data-frame-id="${id}"] .video-title`); - if (titleElement.textContent === `Video Stream ${Number(id)+1}`) { - showModal('当前视频窗口未播放视频。'); - return; - }; - - //发送视频链接接口 - const url = '/api/close_stream'; - const data = {"element_id":id}; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - console.log('Success:', data); - const istatus = data.status; - if(istatus == 0){ - showModal(data.msg); // 使用 Modal 显示消息 - return; - } - else{ - const videoFrame = document.getElementById(`video-${element_id}`); - const event = new Event('closeVideo'); - videoFrame.dispatchEvent(event); - - //videoFrame.style.display = 'none'; // 停止播放时隐藏 img 元素 - titleElement.textContent = `Video Stream ${id+1}`; - removeErrorMessage(videoFrame); - berror_state_list[key] = false; - delete video_list[id]; - } - }) - .catch((error) => { - showModal(`Error: ${error.message}`); // 使用 Modal 显示错误信息 - return; - }); - -} - - -function displayErrorMessage(imgElement, message) { - removeErrorMessage(imgElement) - imgElement.style.display = 'none'; // 隐藏图片 - const errorElement = document.createElement('div'); - errorElement.textContent = message; - errorElement.classList.add('error-message'); - imgElement.parentNode.appendChild(errorElement); -} - -function removeErrorMessage(imgElement) { - const errorElement = imgElement.parentNode.querySelector('.error-message'); - if (errorElement) { - imgElement.parentNode.removeChild(errorElement); - imgElement.style.display = 'block'; - } -} - diff --git a/web/main/static/resources/scripts/base.js b/web/main/static/resources/scripts/base.js index c625c0f..0031633 100644 --- a/web/main/static/resources/scripts/base.js +++ b/web/main/static/resources/scripts/base.js @@ -35,7 +35,7 @@ function set_select_data(select_ele_id,datas){ }); } -//设定选项选中状态 +//设定选项选中状态---下拉框 function set_select_selct(select_ele_id,option_str){ let bfind = false; const select_Ele = document.getElementById(select_ele_id); @@ -49,6 +49,29 @@ function set_select_selct(select_ele_id,option_str){ return bfind; } +//设定单选按钮选中状态 ----单选按钮 +function set_radio_selection(groupName, targetValue) { + // 获取所有同名的 radio 按钮 + const radios = document.getElementsByName(groupName); + + // 遍历 radio 按钮 + for (const radio of radios) { + if (radio.value === targetValue) { + radio.checked = true; // 选中匹配项 + return true; // 返回成功 + } + } + + console.error(`未找到值为 ${targetValue} 的 radio 按钮`); + return false; // 未找到匹配项 +} + +// 获取选中值--返回的value值 ----单选按钮 +function get_radio_value(groupName) { + return document.querySelector(`input[name="${groupName}"]:checked`)?.value; +} + + //将输入框内容转换为单精度型，若转换失败则返回0 -- 若需要整型，parseInt function getInputValueAsFloat(id) { var value = document.getElementById(id).value; diff --git a/web/main/static/resources/scripts/channel_manager.js b/web/main/static/resources/scripts/channel_manager.js deleted file mode 100644 index 4fdc00f..0000000 --- a/web/main/static/resources/scripts/channel_manager.js +++ /dev/null @@ -1,735 +0,0 @@ -const apiEndpoint = '/api/channel/list'; -const rowsPerPage = 10; -//算法配置窗口部分控件 -const searchEndpoint = '/api/channel/select'; -const canvas = document.getElementById('myCanvas'); -const ctx = canvas.getContext('2d'); -const backgroundCanvas = document.getElementById('backgroundCanvas'); -const backgroundCtx = backgroundCanvas.getContext('2d'); -const img = new Image(); -const tbody = document.getElementById('schedule-body');//布防计划 - -let currentPage = 1; -let channelData = []; -let channelData_bak = []; -let areaData = ["请选择"]; -let currentEditingRow = null; -let cid_schedule = "-1"; -let m_polygon = ""; -let check_area = 0; -let draw_status = false; //是否是绘制状态，处于绘制状态才能开始绘制 -let b_img = false; //有没有加载图片成功，如果没有初始化的时候就不绘制线条了。 -let points = []; //检测区域的点坐标数组 -//布防计划 - - -document.addEventListener('DOMContentLoaded', function () { - fetchChannelData(); //初始化通道管理页面元素数据 - - document.getElementById('searchButton').addEventListener('click', function () { - performSearch(); - }); - //新增通道模块--保存按钮 - document.getElementById('saveButton').addEventListener('click', function () { - addChannel(1); - }); - //修改通道模块--保存按钮 - document.getElementById('saveButton_cc').addEventListener('click', function () { - addChannel(2); - }); - //算法配置模块--取消按钮 - document.getElementById('cancelButton_mx').addEventListener('click', function () { - close_mx_model(); - }); - //保存算法配置--保存按钮 - document.getElementById('saveButton_mx').addEventListener('click', function () { - save_mx_model(); - }); - //开始绘制区域按钮 - document.getElementById('but_hzqy').addEventListener('click', function () { - startDraw(); - }); -}); - -//添加和修改通道 1--新增，2--修改 -function addChannel(itype) { - let area; - let cName; - let Rtsp; - let cid; - const spinnerOverlay = document.getElementById("spinnerOverlay"); - let saveButton = null; - let CNameInput = null; - let RTSPInput = null; - - if(itype ==1){ - saveButton = document.getElementById('saveButton'); - CNameInput = document.getElementById('CNameInput'); - RTSPInput = document.getElementById('RTSPInput'); - area = document.getElementById('areaSelect_M').value; - cid = -1 - } - else if(itype ==2){ - saveButton = document.getElementById('saveButton_cc'); - CNameInput = document.getElementById('CNameInput_cc'); - RTSPInput = document.getElementById('RTSPInput_cc'); - area = document.getElementById('areaSelect_CC').value; - cid = currentEditingRow.cells[0].innerText; - } - console.log("点击了保存按钮"); - cName = CNameInput.value.trim(); - Rtsp = RTSPInput.value.trim(); - if(area === "请选择"){ - alert('请选择所属区域'); - } - else{ - if (cName && Rtsp) { - saveButton.disabled = true; - //发送视频链接接口 - const url = '/api/channel/add'; - const data = {"area":area,"cName":cName,"Rtsp":Rtsp,"cid":cid}; - // 显示 Spinners - spinnerOverlay.style.display = "flex"; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - const istatus = data.status; - saveButton.disabled = false; - alert(data.msg); // 使用 Modal 显示消息 - if(istatus == 1){ - //刷新列表 - fetchChannelData(); - if(itype ==1){ - //添加通道成功 - $('#channelModal').modal('hide'); - } - else if(itype==2){ - //修改通道成功 - currentEditingRow = null; - $('#ChangeC').modal('hide'); - } - } - }) - .catch((error) => { - alert(`Error: ${error.message}`); // 使用 Modal 显示错误信息 - // 启用保存按钮 - saveButton.disabled = false; - return; - }) - .finally(()=>{ - // 隐藏 Spinners - spinnerOverlay.style.display = "none"; - }); - } else { - alert('通道名称和RTSP地址不能为空'); - } - } -} - -//初始化通道管理页面元素数据 -async function fetchChannelData() { //获取通道相关信息（/api/channel/list），刷新通道表格控件数据 - try { - const response = await fetch(apiEndpoint); - channelData = await response.json(); - channelData_bak = channelData; - - url = "/api/channel/area/list" - area_response = await fetch(url); - areaDatas = await area_response.json(); - areaData = ["请选择"]; //清空下 - areaDatas.forEach((area) => { - areaData.push(area.area_name) - }); - - renderTable(); //刷新表格 - renderPagination(); //刷新分页元素 - renderAreaOptions(); //所属区域下来框 - } catch (error) { - console.error('Error fetching channel data:', error); - } -} - -//刷新表单页面数据 -function renderTable() { - const tableBody = document.getElementById('table-body'); - tableBody.innerHTML = ''; - - const start = (currentPage - 1) * rowsPerPage; - const end = start + rowsPerPage; - const pageData = channelData.slice(start, end); - const surplus_count = rowsPerPage - pageData.length; - - pageData.forEach((channel) => { -// if(area_name!==channel.area_name){ //这里要求区域名称一样的要在一起 -// area_name = channel.area_name; -// areaData.push(area_name); -// } - const row = document.createElement('tr'); - row.innerHTML = ` - ${channel.ID} - ${channel.area_name} - ${channel.channel_name} - ${channel.ulr} - ${channel.model_name} - - - - - - `; - tableBody.appendChild(row); - row.querySelector('.modify-btn').addEventListener('click', () => modifyChannel(row)); - row.querySelector('.algorithm-btn').addEventListener('click', () => configureAlgorithm(row)); - row.querySelector('.delete-btn').addEventListener('click', () => deleteChannel(row)); - }); -} - -//关键字查询数据 -async function performSearch() { - try { - const area = document.getElementById('areaSelect').value; - const channelName = document.getElementById('channelNameInput').value; - if(area === "请选择" && channelName===""){ - channelData = channelData_bak; - } - else if(area === "请选择"){ - channelData = []; - channelData_bak.forEach((channel) => { - if(channelName === channel.channel_name){ - channelData.push(channel); - } - }); - } - else if(channelName === ""){ - channelData = []; - channelData_bak.forEach((channel) => { - if(area === channel.area_name){ - channelData.push(channel); - } - }); - } - else{ - channelData = []; - channelData_bak.forEach((channel) => { - if(area === channel.area_name && channelName === channel.channel_name){ - channelData.push(channel); - } - }); - } - // 渲染表格和分页控件 - currentPage = 1; // 重置当前页为第一页 - renderTable(); - renderPagination(); - } catch (error) { - console.error('Error performing search:', error); - } -} - -//点击修改按钮，显示修改通道信息模块 --只是显示 -function modifyChannel(row) { -// const cid = row.cells[0].innerText; - const areaName = row.cells[1].innerText; - const channelName = row.cells[2].innerText; - const url = row.cells[3].innerText; - - const area = document.getElementById('areaSelect_CC'); - const CName = document.getElementById('CNameInput_cc'); - const RTSP = document.getElementById('RTSPInput_cc'); - - for(let i=0;i< area.options.length;i++){ - if(area.options[i].value === areaName){ - area.options[i].selected = true; - break; - } - } - CName.value = channelName; - RTSP.value = url; - - currentEditingRow = row; - $('#ChangeC').modal('show'); -} - -//点击算法按钮，显示算法配置模块 --只是显示 -function configureAlgorithm(row) { - //获取当前行信息 - currentEditingRow = row; - const cid = row.cells[0].innerText; - //清除数据，若需要的话 - ctx.clearRect(0, 0, canvas.width, canvas.height); //清除左侧绘画和画线信息 - tbody.innerHTML = ''; //清空布防控件数据 - points = []; //清空绘制检测区域 - draw_status = false; - b_img = false; - document.getElementById('but_hzqy').textContent = "绘制区域"; - //开始初始化算法管理模块 - show_channel_model_schedule(cid); //获取并显示结构化数据 - show_channel_img(cid); //获取并显示一帧图片 -- 获取不到图片就是黑画面 --并要绘制检测区域 - - //显示窗口 - $('#MX_M').modal('show'); -} - -//获取一帧图片 -function show_channel_img(cid){ - const data = {"cid":cid}; - fetch('/api/channel/img', { - method: 'POST', - headers: { - 'Content-Type': 'application/json' - }, - body: JSON.stringify(data) - }) - .then(response => response.json()) - .then(data => { - if (data.image) { - b_img = true; - img.src = 'data:image/jpeg;base64,' + data.image; - } else { - console.error('Error:', data.error); - } - }) - .catch(error => console.error('Error:', error)); -} - -//图片加载事项 -img.onload = () => { //清除、画图和画线应该分开 - // 设置画布宽高 - backgroundCanvas.width = canvas.width = img.width; - backgroundCanvas.height = canvas.height = img.height; - // 将图片绘制到背景画布上 - backgroundCtx.drawImage(img, 0, 0, img.width, img.height); - drawLines(); - // 将背景画布的内容复制到前台画布上 - //ctx.drawImage(backgroundCanvas, 0, 0, canvas.width, canvas.height); //绘制画面 -}; - -//开始和重新绘制 -function startDraw(){ - if(!document.getElementById('zdjc').checked){ - alert("请先选择指定区域！"); - return; - } - let but = document.getElementById('but_hzqy'); - if(!draw_status){//开始绘制 - if(points.length >0){ - if (confirm('开始绘制将清除未提交保存的绘制数据，是否继续？')) { - draw_status = true; - points = []; - //按钮文字调整为结束绘制 - but.textContent = '结 束 绘 制'; - // 清除前台画布 - ctx.clearRect(0, 0, canvas.width, canvas.height); - // 将背景画布的内容复制到前台画布上 - ctx.drawImage(backgroundCanvas, 0, 0, canvas.width, canvas.height); - } - } - else{ - draw_status = true; - but.textContent = '结束绘制'; - } - } - else{//结束绘制 - draw_status = false; - but.textContent = '绘制区域'; - } -} - -//单选按钮点击事件处理 -function handleRadioClick(event) { - const selectedRadio = event.target; - console.log('Selected Radio:', selectedRadio.id); - // 根据选中的单选按钮执行相应操作 - if (selectedRadio.id === 'qjjc') { - if(draw_status){ - alert("请先结束绘制后，再切换检测方案！"); - document.getElementById('zdjc').checked = true; - } - else{ - // 处理全画面生效的逻辑 - if(points.length>0){ - if (!confirm('切换到全画面生效，将清除已绘制的区域信息，是否切换？')) { - document.getElementById('zdjc').checked = true; - }else{ - points = []; - } - } - } - //console.log('全画面生效'); - } else if (selectedRadio.id === 'zdjc') { - // 处理指定区域的逻辑 - console.log('指定区域'); - } -} - -// 鼠标点击事件处理--动态绘图 -canvas.addEventListener('click', (event) => { - if(draw_status){ - const rect = canvas.getBoundingClientRect(); - const scaleX = canvas.width / rect.width; - const scaleY = canvas.height / rect.height; - - // 获取鼠标相对于canvas的位置 - const x = (event.clientX - rect.left) * scaleX; - const y = (event.clientY - rect.top) * scaleY; - points.push({ x, y }); - console.log(points); - //绘制线条 - drawLines(); - } -}); - -// 绘制区域，各点连接 -function drawLines() { - if(b_img){ - // 清除前台画布 - ctx.clearRect(0, 0, canvas.width, canvas.height); - // 将背景画布的内容复制到前台画布上 - ctx.drawImage(backgroundCanvas, 0, 0, canvas.width, canvas.height); - - // 绘制点和线 - ctx.strokeStyle = 'red'; - ctx.lineWidth = 2; - - if (points.length > 0) { - ctx.beginPath(); - ctx.moveTo(points[0].x, points[0].y); - - for (let i = 1; i < points.length; i++) { - ctx.lineTo(points[i].x, points[i].y); - } - - // 连接最后一个点到起点 - ctx.lineTo(points[0].x, points[0].y); - ctx.stroke(); - } - - points.forEach(point => { - ctx.beginPath(); - ctx.arc(point.x, point.y, 5, 0, Math.PI * 2); - ctx.fillStyle = 'red'; - ctx.fill(); - }); - } -} - -//获取并显示该通道相关算法的结构化数据 --- 这里用GET会更加贴切一些 -function show_channel_model_schedule(cid){ - //发送视频链接接口 - const url = '/api/channel/C2M'; - const data = {"cid":cid}; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - const m_datas = data.m_datas; //算法清单 - const c2m_data = data.c2m_data; //该通道管理算法的相关数据，会有空的情况 - const schedule = data.schedule; //布防计划 - //console.log("m_datas--",m_datas); - //console.log("c2m_data--",c2m_data); - //console.log("schedule--",schedule); - //配置算法下拉清单 - select_datas = ["请选择"]; - m_datas.forEach(option => { - select_datas.push(option.model_name); - }); - set_select_data("model_select",select_datas); - select_str = currentEditingRow.cells[4].innerText; - //检测区域 - if(c2m_data.length >0){ - model_id = c2m_data[0].model_id; - model_name = currentEditingRow.cells[4].innerText; - set_select_selct("model_select",model_name); - check_area = c2m_data[0].check_area - if( check_area == 0){ //全画面生效 - document.getElementById('qjjc').checked = true; - m_polygon = ""; - } - else{//指定区域 - document.getElementById('zdjc').checked = true; - m_polygon = c2m_data[0].polygon; - console.log("m_polygon--",m_polygon); - if(m_polygon !== ""){ //指定区域了，一般是会有数据的。 - const coords = parseCoordStr(m_polygon); - points = coords; - } - } - //阈值 - document.getElementById('zxyz').value = c2m_data[0].conf_thres - document.getElementById('iouyz').value = c2m_data[0].iou_thres - } - - //布防计划 - const days = ['一', '二', '三', '四', '五', '六','日']; - const num_days=['0','1','2','3','4','5','6'] - days.forEach((day, dayIndex) => { - const row = document.createElement('tr'); - const dayCell = document.createElement('th'); - dayCell.textContent = day; - row.appendChild(dayCell); - num_day = num_days[dayIndex] - for (let hour = 0; hour < 24; hour++) { - const cell = document.createElement('td'); - if(schedule.length >0){ - const status = schedule.find(item => item.day === num_day && item.hour === hour); - if (status && status.status === 1) { - cell.classList.add('blocked'); - } else { - cell.classList.add('allowed'); - } - } - else{ - cell.classList.add('blocked'); - } - row.appendChild(cell); - - cell.addEventListener('click', () => { - if (cell.classList.contains('blocked')) { - cell.classList.remove('blocked'); - cell.classList.add('allowed'); - // Update status in the database - //updateStatus(day, hour, 0); - } else { - cell.classList.remove('allowed'); - cell.classList.add('blocked'); - // Update status in the database - //updateStatus(day, hour, 1); - } - }); - } - tbody.appendChild(row); - }); - }) - .catch((error) => { - alert(`Error: ${error.message}`); // 使用 Modal 显示错误信息 - return; - }); -} - -// 将字符串转换为数组 -function parseCoordStr(str) { - return str.match(/\(([^)]+)\)/g).map(pair => { - const [x, y] = pair.replace(/[()]/g, '').split(',').map(Number); - return { x, y }; - }); -} - -//关闭算法配置窗口 -function close_mx_model(){ - if (confirm('确定退出窗口吗？未保存的修改将丢失！')) { - $('#MX_M').modal('hide'); - } -} - -//保存算法配置窗口数据 -function save_mx_model(){ - let model_name; //算法名称 - let check_area; //检测区域标识 0-全局，1-指定范围 - let polygon_str; //具体的检测区域 - let conf_thres; //置信阈值 - let iou_thres; //iou阈值 - let schedule; //布防计划 - const saveButton = document.getElementById('saveButton_mx'); - saveButton.disabled = true; //不可点击状态 - //配置算法 - model_name = document.getElementById("model_select").value; - //检测区域 - if(document.getElementById('zdjc').checked){ - check_area = 1; - console.log("points--",points); - if (points.length > 0){ - const formattedArray = points.map(point => `(${point.x},${point.y})`); - polygon_str = `[${formattedArray.join(',')}]`; - }else{ - polygon_str = ""; - } - - }else{ - check_area = 0; - polygon_str = ""; - } - //置信阈值和IOU阈值 - conf_thres = getInputValueAsFloat('zxyz'); - iou_thres = getInputValueAsFloat('iouyz'); - //验证数据 - if(model_name !== "请选择"){ - console.log(model_name); - if(conf_thres <= 0 || conf_thres>=1 || iou_thres <= 0 || iou_thres>=1){ - alert("阈值的有效范围是大于0，小于1；请输入正确的阈值(默认可0.5)。"); - saveButton.disabled = false; //不可点击状态 - return; - } - } - //布防计划 - // 定义一个对象来存储数据 - const scheduleData = { - '0': Array(24).fill(0), - '1': Array(24).fill(0), - '2': Array(24).fill(0), - '3': Array(24).fill(0), - '4': Array(24).fill(0), - '5': Array(24).fill(0), - '6': Array(24).fill(0) - }; - // 遍历 tbody 的每一行 - [...tbody.children].forEach((row, dayIndex) => { - // 获取当前行的所有单元格 - const cells = row.getElementsByTagName('td'); - - // 遍历每一个单元格 - for (let hour = 0; hour < cells.length; hour++) { - // 检查单元格的 class 是否包含 'blocked' - if (cells[hour].classList.contains('blocked')) { - // 将对应的 scheduleData 位置设置为 1 - scheduleData[dayIndex][hour] = 1; - } else { - // 将对应的 scheduleData 位置设置为 0 - scheduleData[dayIndex][hour] = 0; - } - } - }); - // 将 scheduleData 对象转换为 JSON 字符串 - const scheduleData_json = JSON.stringify(scheduleData); - //提交到服务器 -// console.log("model_name--",model_name); -// console.log("check_area--",check_area); -// console.log("polygon_str--",polygon_str); -// console.log("iou_thres--",iou_thres); -// console.log("conf_thres--",conf_thres); -// console.log("schedule-- ",scheduleData_json); - cid = currentEditingRow.cells[0].innerText; - const url = '/api/channel/chanegeC2M'; - const data = {"model_name":model_name,"check_area":check_area,"polygon_str":polygon_str,"iou_thres":iou_thres, - "conf_thres":conf_thres,"schedule":scheduleData_json,"cid":cid}; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - const istatus = data.status; - if(istatus === 0){ - alert(data.msg); // 使用 Modal 显示消息 - // 启用保存按钮 - saveButton.disabled = false; - return; - } - else{ - // 启用保存按钮 - saveButton.disabled = false; - //刷新列表 - fetchChannelData(); - //$('#MX_M').modal('hide'); - alert("修改成功！"); - } - }) - .catch((error) => { - alert(`Error: ${error.message}`); - // 启用保存按钮 - saveButton.disabled = false; - return; - }); - -} - -//删除通道 -function deleteChannel(row) { - if (confirm('确定删除此通道吗？')) { - cid = row.cells[0].innerText; - //发送视频链接接口 - const url = '/api/channel/del'; - const data = {"cid":cid}; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - const istatus = data.status; - if(istatus === 0){ - alert(data.msg); // 使用 Modal 显示消息 - return; - } - else{ - //刷新列表 - row.remove(); - alert("删除通道成功！"); - } - }) - .catch((error) => { - alert(`Error: ${error.message}`); // 使用 Modal 显示错误信息 - return; - }); - - } -} - -//刷新分页标签 -function renderPagination() { - const pagination = document.getElementById('pagination'); - pagination.innerHTML = ''; - - const totalPages = Math.ceil(channelData.length / rowsPerPage); - for (let i = 1; i <= totalPages; i++) { - const pageItem = document.createElement('li'); - pageItem.className = 'page-item' + (i === currentPage ? ' active' : ''); - pageItem.innerHTML = `${i}`; - pageItem.addEventListener('click', (event) => { - event.preventDefault(); - currentPage = i; - renderTable(); - renderPagination(); - }); - pagination.appendChild(pageItem); - } -} - -//刷新区域下拉控件 -function renderAreaOptions() { - const areaSelect = document.getElementById('areaSelect'); - const areaSelect_M = document.getElementById('areaSelect_M') - const areaSelect_CC = document.getElementById('areaSelect_CC') - //先清空 - areaSelect.innerHTML = ''; - areaSelect_M.innerHTML = ''; - areaSelect_CC.innerHTML = ''; - //再添加 - areaData.forEach(option => { - const optionElement = document.createElement('option'); - optionElement.textContent = option; - areaSelect.appendChild(optionElement); - - const optionElement_m = document.createElement('option'); - optionElement_m.textContent = option; - areaSelect_M.appendChild(optionElement_m); - - const optionElement_cc = document.createElement('option'); - optionElement_cc.textContent = option; - areaSelect_CC.appendChild(optionElement_cc); - }); -} - - - - - diff --git a/web/main/static/resources/scripts/model_manager.js b/web/main/static/resources/scripts/model_manager.js deleted file mode 100644 index ec0b9fa..0000000 --- a/web/main/static/resources/scripts/model_manager.js +++ /dev/null @@ -1,303 +0,0 @@ -let currentPage = 1; -const rowsPerPage = 10; -let modelData = []; -let modelData_bak = []; //用于关键字检索 -let currentEditingRow = null; - - - -//页面加载初始化 -document.addEventListener('DOMContentLoaded', function () { - fetchModelData(); //初始化通道管理页面元素数据 - //新增算法模态框---保存按钮 - document.getElementById('saveButton_model').addEventListener('click',function () { - addModel(); - }); - - //配置算法模态框--保存按钮 - document.getElementById('saveButton_config_model').addEventListener('click',function () { - post_configureModel(); - }); - - //升级算法模态框--保存按钮 - document.getElementById('saveButton_upmodel').addEventListener('click',function () { - post_modifyModel(); - }); - - //查询按钮 - document.getElementById('searchMButton').addEventListener('click',function () { - searchModel(); - }); -}); - -//获取算法列表数据，并更新页面 -async function fetchModelData() { - try{ - let response = await fetch('/api/model/list'); - if (!response.ok) { - throw new Error('Network response was not ok'); - } - modelData = await response.json(); - modelData_bak = modelData; - currentPage = 1; // 重置当前页为第一页 - renderTable(); //刷新表格 - renderPagination(); - }catch (error) { - console.error('Error fetching model data:', error); - } -} - -//刷新表单页面数据 -function renderTable() { - const tableBody = document.getElementById('table-body-model'); - tableBody.innerHTML = ''; //清空 - - const start = (currentPage - 1) * rowsPerPage; - const end = start + rowsPerPage; - const pageData = modelData.slice(start, end); - const surplus_count = rowsPerPage - pageData.length; - - - pageData.forEach((model) => { - const row = document.createElement('tr'); - row.innerHTML = ` - ${model.ID} - ${model.name} - ${model.version} - ${model.duration_time} - ${model.proportion} - - - - - - `; - tableBody.appendChild(row); - row.querySelector('.modify-btn').addEventListener('click', () => modifyModel(row)); - row.querySelector('.algorithm-btn').addEventListener('click', () => configureModel(row)); - row.querySelector('.delete-btn').addEventListener('click', () => deleteModel(row)); - }); -} - -//刷新分页标签 -function renderPagination() { - const pagination = document.getElementById('pagination-model'); - pagination.innerHTML = ''; - - const totalPages = Math.ceil(modelData.length / rowsPerPage); - for (let i = 1; i <= totalPages; i++) { - const pageItem = document.createElement('li'); - pageItem.className = 'page-item' + (i === currentPage ? ' active' : ''); - pageItem.innerHTML = `${i}`; - pageItem.addEventListener('click', (event) => { - event.preventDefault(); - currentPage = i; - renderTable(); - renderPagination(); - }); - pagination.appendChild(pageItem); - } -} - -//显示升级算法模态框 -function modifyModel(row){ - currentEditingRow = row; - model_name = row.cells[1].innerText; - version_name = row.cells[2].innerText; - - document.getElementById('update_mname_label').innerText = `算法名称: ${model_name}`; - document.getElementById('update_mversion_label').innerText = `当前版本: ${version_name}`; - $('#updateMM').modal('show'); -} - -//升级算法模态框--点击保存按钮 -function post_modifyModel(){ - mid = currentEditingRow.cells[0].innerText; - const btn = document.getElementById('saveButton_upmodel'); - const fileInput = document.getElementById('updateModelFile'); - const file = fileInput.files[0]; - if(file){ - btn.disabled = true; //不可点击 - const formData = new FormData(); - formData.append('file', file); - formData.append('mid', mid); - - fetch('/api/model/upgrade', { - method: 'POST', - body: formData, - }) - .then(response => response.json()) - .then(data => { - const istatus = data.status; - alert(data.msg); - btn.disabled = false; - if(istatus == 1 ){ - fetchModelData(); - $('#updateMM').modal('hide'); - } - }) - .catch(error => { - console.error('Error:', error); - alert('升级失败，请重试。'); - btn.disabled = false; - }); - } - else{ - alert('请选择升级包进行上传。'); - btn.disabled = false; - } -} - -//显示配置算法模态框 -function configureModel(row){ - currentEditingRow = row; - model_name = row.cells[1].innerText; - duration_time = row.cells[3].innerText; - proportion = row.cells[4].innerText; - //设置模态框控件遍历 - document.getElementById('config_mname_label').innerText = `算法名称: ${model_name}`; - document.getElementById('duration_timeInput').value = duration_time; - document.getElementById('proportionInput').value = proportion; - $('#configMM').modal('show'); -} - -//配置算法模态框--点击保存按钮 -function post_configureModel(){ - mid = currentEditingRow.cells[0].innerText; - duration_time = parseInt(document.getElementById('duration_timeInput').value); - proportion = parseFloat(document.getElementById('proportionInput').value); - if(isNaN(duration_time) || isNaN(proportion) ){ - alert("请输入数字！"); - return; - } - if(proportion<=0 || proportion>=1){ - alert("占比阈值需要大于0，且小于1"); - return; - } - //提交数据 - const url = '/api/model/changecnf'; - const data = {"mid":mid,"duration_time":duration_time,"proportion":proportion}; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - const istatus = data.status; - if(istatus === 0){ - alert(data.msg); - return; - } - else{ - //刷新列表 - fetchModelData(); - alert(data.msg); - $('#configMM').modal('hide'); - } - }) - .catch((error) => { - alert(`Error: ${error.message}`); - return; - }); -} - -//删除算法记录 -function deleteModel(row){ - if (confirm('确定删除此算法吗？')) { - mid = row.cells[0].innerText; - const url = '/api/model/del'; - const data = {"mid":mid}; - // 发送 POST 请求 - fetch(url, { - method: 'POST', // 指定请求方法为 POST - headers: { - 'Content-Type': 'application/json' // 设置请求头，告诉服务器请求体的数据类型为 JSON - }, - body: JSON.stringify(data) // 将 JavaScript 对象转换为 JSON 字符串 - }) - .then(response => response.json()) // 将响应解析为 JSON - .then(data => { - const istatus = data.status; - if(istatus === 0){ - alert(data.msg); - return; - } - else{ - //刷新列表 - row.remove(); - alert("删除算法成功！"); - } - }) - .catch((error) => { - alert(`Error: ${error.message}`); - return; - }); - - } -} - -//新增算法--保存按钮 -function addModel(){ - const btn = document.getElementById('saveButton_model'); - const fileInput = document.getElementById('uploadModelFile'); - const file = fileInput.files[0]; - const mName = document.getElementById('MNameInput').value; - - if (file && mName) { - btn.disabled = true; //不可点击 - const formData = new FormData(); - formData.append('file', file); - formData.append('mName', mName); - - fetch('/api/model/add', { - method: 'POST', - body: formData, - }) - .then(response => response.json()) - .then(data => { - const istatus = data.status; - alert(data.msg); - btn.disabled = false; - if(istatus == 1 ){ - fetchModelData(); - $('#addMM').modal('hide'); - } - }) - .catch(error => { - console.error('Error:', error); - alert('上传失败，请重试。'); - btn.disabled = false; - }); - } else { - alert('请填写算法名称并选择一个升级包进行上传。'); - btn.disabled = false; - } -} - -//关键字检索 -function searchModel(){ - try { - const modelName = document.getElementById('modelNameInput').value; - if(modelName===""){ - modelData = modelData_bak; - } - else{ - modelData = []; - modelData_bak.forEach((model) => { - if(model.name.includes(modelName)){ - modelData.push(model); - } - }); - } - // 渲染表格和分页控件 - currentPage = 1; // 重置当前页为第一页 - renderTable(); - renderPagination(); - } catch (error) { - console.error('Error performing search:', error); - } -} \ No newline at end of file diff --git a/web/main/static/resources/scripts/my_web_socket.js b/web/main/static/resources/scripts/my_web_socket.js new file mode 100644 index 0000000..efa1efc --- /dev/null +++ b/web/main/static/resources/scripts/my_web_socket.js @@ -0,0 +1,82 @@ +function getWsUrl() { + const protocol = window.location.protocol === "https:" ? "wss" : "ws"; + const host = window.location.host; // 包含主机和端口（如 localhost:5000） + const path = "/api/ws"; // 你的 WebSocket 路由 + return `${protocol}://${host}${path}`; +} + +function initWebSocket() { + // 替换成你实际的 WebSocket 地址，例如 ws://localhost:5000/ws + wsUrl = getWsUrl(); + const ws = new WebSocket(wsUrl); + // 设置接收二进制数据为 ArrayBuffer + ws.binaryType = "arraybuffer"; + + ws.onopen = function() { + console.log("WebSocket 已连接"); + // 发送登录数据包，假设格式为 JSON（后端可以根据需要解析此包） --暂时默认0只有一个 + ws.send(JSON.stringify({ user_id: 0 })); + }; + + ws.onmessage = function(event) { + // 解析收到的二进制数据 + const buffer = event.data; + const dataView = new DataView(buffer); + // 读取前四个字节，验证魔数 "TFTF" + let magic = ""; + for (let i = 0; i < 4; i++) { + magic += String.fromCharCode(dataView.getUint8(i)); + } + if (magic !== "TFTF") { + console.error("收到非法格式数据:", magic); + return; + } + // 读取数据头：接下来的 8 字节分别为 idata_type 和 idata_len（均为 32 位无符号整数，假设大端） + const idata_type = dataView.getUint32(4, false); + const idata_len = dataView.getUint32(8, false); + + // 数据体从偏移量 12 开始，长度为 idata_len + const bodyBytes = new Uint8Array(buffer, 12, idata_len); + const decoder = new TextDecoder("utf-8"); + const bodyText = decoder.decode(bodyBytes); + // 假设数据体为 JSON 字符串 + let bodyData; + try { + bodyData = JSON.parse(bodyText); + } catch (error) { + console.error("解析数据体出错:", error); + return; + } + + // 针对 idata_type 为 1 的处理：更新节点 + if(idata_type === 0){ + console.log(bodyData); + } + else if (idata_type === 1) {//type-1为更新节点的node_workstatus + // bodyData 应该包含 node_path 和 node_workstatus 两个字段 + updateTreeNode(bodyData.node_path, bodyData.node_workstatus);//node_tree.js + } + else if(idata_type === 2){ + if(cur_task_id === 0){return;} + // 清空选中状态 + selectedNodeData = null; + //刷新界面内容 + update_select_node_data_show("-","-","-","-","-",false) + // 重新加载节点树数据 + loadNodeTree(cur_task_id); + } + else { + console.error("未知的数据类型"); + } + }; + + ws.onerror = function(error) { + console.error("WebSocket 错误:", error); + }; + + ws.onclose = function() { + console.log("WebSocket 连接关闭"); + }; + + return ws; +} diff --git a/web/main/static/resources/scripts/node_tree.js b/web/main/static/resources/scripts/node_tree.js new file mode 100644 index 0000000..2145890 --- /dev/null +++ b/web/main/static/resources/scripts/node_tree.js @@ -0,0 +1,929 @@ + // 全局变量，用于保存当前选中的节点数据 + 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/gettree", { + 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 = "执行中"; + document.getElementById("btnToggleStatus").textContent = "暂停"; + }else { + document.getElementById("node_bwork").textContent = "暂停中"; + document.getElementById("btnToggleStatus").textContent = "继续"; + } + setNodeBtnStatus(); + } + + //节点按钮的状态控制 + function setNodeBtnStatus(){ + const btn_TS = document.getElementById("btnToggleStatus"); + const btn_NodeStep = document.getElementById("btnNodeStep"); + const btn_VI = document.getElementById("btnViewInstr"); + const btn_VM = document.getElementById("btnViewMsg"); + const btn_AI = document.getElementById("btnAddInfo"); + const btn_AC = document.getElementById("btnAddChild"); + if(!selectedNodeData){ + //没有选择node，按钮全部置不可用 + btn_TS.disabled = true; + btn_TS.classList.add("disabled-btn"); + btn_NodeStep.disabled = true; + btn_NodeStep.classList.add("disabled-btn"); + btn_VI.disabled = true; + btn_VI.classList.add("disabled-btn"); + btn_VM.disabled = true; + btn_VM.classList.add("disabled-btn"); + btn_AI.disabled = true; + btn_AI.classList.add("disabled-btn"); + btn_AC.disabled = true; + btn_AC.classList.add("disabled-btn"); + } + else{ + //5个可用 + btn_TS.disabled = false; + btn_TS.classList.remove("disabled-btn"); + btn_VI.disabled = false; + btn_VI.classList.remove("disabled-btn"); + btn_VM.disabled = false; + btn_VM.classList.remove("disabled-btn"); + btn_AI.disabled = false; + btn_AI.classList.remove("disabled-btn"); + btn_AC.disabled = false; + btn_AC.classList.remove("disabled-btn"); + if(cur_task.taskStatus === 1 && cur_task.workType === 0 && selectedNodeData.node_bwork==="true"){ + btn_NodeStep.disabled = false; + btn_NodeStep.classList.remove("disabled-btn"); + } + else{ + btn_NodeStep.disabled = true; + btn_NodeStep.classList.add("disabled-btn"); + } + } + } + + // 刷新按钮事件绑定 + document.getElementById("btnRefresh").addEventListener("click", () => { + // 重新加载节点树数据 + loadNodeTree(cur_task_id); + }); + + // 按钮事件：当未选中节点时提示 + function checkSelectedNode() { + if (!selectedNodeData) { + alert("请先选择节点"); + return false; + } + return true; + } + + //节点-暂停/继续，bwork控制 + document.getElementById("btnToggleStatus").addEventListener("click", () => { + if (!checkSelectedNode()) return; + // bwork的状态已后端状态为准，只做切换 + update_node_bwork(cur_task_id,selectedNodeData.node_path); + }); + async function update_node_bwork(task_id,node_path){ + try { + const res = await fetch("/api/task/nodecontrol", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ task_id,node_path }), //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 newbwork = data.newbwork; + //更新数据 + const selectedEl = document.querySelector(".tree-node.selected"); + if (selectedEl) { + selectedEl.setAttribute("data-node_bwork", newbwork); + selectedNodeData.node_bwork = newbwork; + } + //刷新界面 + const btn_NodeStep = document.getElementById("btnNodeStep"); + if(newbwork){ + document.getElementById("node_bwork").textContent ="执行中"; + document.getElementById("btnToggleStatus").textContent = "暂停"; + if(cur_task.taskStatus === 1 && cur_task.workType === 0){ + btn_NodeStep.disabled = false; + btn_NodeStep.classList.remove("disabled-btn"); + } + }else { + document.getElementById("node_bwork").textContent = "暂停中"; + document.getElementById("btnToggleStatus").textContent = "继续"; + btn_NodeStep.disabled = true; + btn_NodeStep.classList.add("disabled-btn"); + } + }catch (error) { + alert("修改节点的bwork失败：", error); + } + } + + //节点-单步工作 + document.getElementById("btnNodeStep").addEventListener("click", () => { + if (!checkSelectedNode()) return; + node_one_step(cur_task_id,selectedNodeData.node_path); + }); + async function node_one_step(task_id,node_path){ + try { + const res = await fetch("/api/task/nodestep", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ task_id,node_path }), //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 bsuccess = data.bsuccess; + if(bsuccess){ + alert("该节点任务已提交，请稍候查看执行结果！") + } + else{ + error = data.erroe; + alert("该节点单步失败！",error) + } + + }catch (error) { + alert("该节点单步失败，请联系管理员！", error); + } + } + + //----------------------查看指令modal---------------------------- + let doneInstrs = []; // 已执行指令的所有数据 + let todoInstrs = []; // 待执行指令的所有数据 + let donePage = 1; // 已执行指令当前页 + let todoPage = 1; // 待执行指令当前页 + const pageSize = 10; // 每页固定显示 10 行 + + document.getElementById("msgModal").addEventListener("hidden.bs.modal", () => { + document.activeElement.blur(); // 清除当前焦点 + }); + + document.getElementById("btnViewInstr").addEventListener("click", () => { + if (!checkSelectedNode()) return; + openInstrModal() + }); + // 打开对话框函数 + function openInstrModal() { + const modalEl = document.getElementById("instrModal"); + // 假设用 Bootstrap 5 的 Modal 组件 + const instrModal = new bootstrap.Modal(modalEl, {keyboard: false}); + // 在打开 modal 时，先更新提示内容，将 loadingMsg 显示“请稍后，数据获取中…” + const loadingMsg = document.getElementById("loadingMsg"); + if (loadingMsg) { + loadingMsg.textContent = "请稍后，数据获取中..."; + } + // 显示对话框 + instrModal.show(); + // 加载指令数据 + loadInstrData(); + } + + // 调用后端接口，获取指令数据 + async function loadInstrData() { + task_id = cur_task_id; + node_path = selectedNodeData.node_path; + try { + const res = await fetch("/api/task/nodegetinstr", { + 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); + } + } + + function confirmCopyTodoInstr(idx) { + // 从全局数组 todoInstrs 中获取指令文本 + const instruction = todoInstrs[idx]; + + if (navigator.clipboard && navigator.clipboard.writeText) { + navigator.clipboard.writeText(instruction) + .then(() => { + alert("已复制: " + instruction); + }) + .catch((err) => { + console.error("使用 Clipboard API 复制失败:", err); + fallbackCopyTextToClipboard(instruction); + }); + } else { + // 如果 clipboard API 不可用，则回退使用 execCommand 方法 + fallbackCopyTextToClipboard(instruction); + } +} + +function fallbackCopyTextToClipboard(text) { + // 创建 textarea 元素 + const textArea = document.createElement("textarea"); + textArea.value = text; + + // 使用 CSS 样式使其不可见，同时保证能够获得焦点 + textArea.style.position = "fixed"; // 避免页面滚动 + textArea.style.top = "0"; + textArea.style.left = "0"; + textArea.style.width = "2em"; + textArea.style.height = "2em"; + textArea.style.padding = "0"; + textArea.style.border = "none"; + textArea.style.outline = "none"; + textArea.style.boxShadow = "none"; + textArea.style.background = "transparent"; + + document.body.appendChild(textArea); + textArea.focus(); + textArea.select(); + + try { + const successful = document.execCommand('copy'); + if (successful) { + alert("已复制: " + text); + } else { + alert("复制失败，请手动复制！"); + } + } catch (err) { + console.error("Fallback: 无法复制", err); + alert("复制失败，请手动复制！"); + } + document.body.removeChild(textArea); +} + + + // 删除待执行指令，先确认 + function confirmDeleteTodoInstr(arrIndex) { + if (!confirm("确认删除该条待执行指令？")) return; + // arrIndex 在当前分页中的索引 + // 先算出全局索引 + const realIndex = (todoPage - 1) * pageSize + arrIndex; + const item = todoInstrs[realIndex]; + // 调用后端删除接口 + deleteTodoInstr(item); + } + + // 调用后端接口删除 + async function deleteTodoInstr(item) { + task_id = cur_task_id; + node_path = selectedNodeData.node_path; + try { + const res = await fetch("/api/task/delnodeinstr", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ task_id,node_path,item }) + }); + if (!res.ok) { + alert(data.error || "删除失败"); + return; + } + const data = await res.json(); + if(data.bsuccess){ + // 删除成功，更新本地数据 + const idx = todoInstrs.findIndex(x => x.id === item.id); + if (idx !== -1) { + todoInstrs.splice(idx, 1); + } + // 重新渲染 + renderTodoInstrTable(todoPage); + //0需要更新work_status为无待执行任务的状态 + if(todoInstrs.length === 0){ + updateTreeNode(node_path, 0); + } + } + else{ + alert("指令删除失败",data.error) + } + } catch (error) { + console.error("删除指令异常:", error); + alert("删除指令异常，请联系管理员！"); + } + } + + // 分页事件 + 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); + } + + //---------------------查看MSGmodal------------------------------ + document.getElementById("msgModal").addEventListener("hidden.bs.modal", () => { + document.activeElement.blur(); // 清除当前焦点 + }); + + document.getElementById("btnViewMsg").addEventListener("click", () => { + if (!checkSelectedNode()) return; + openMsgModal(); + }); + +let submittedMsgs = []; // 存储已提交的 MSG 数据数组 +let pendingMsg = {}; // 存储待提交的 MSG 数据 +let submittedPage = 1; + +function openMsgModal(){ + // 显示 Modal（使用 Bootstrap 5 Modal） + const msgModal = new bootstrap.Modal(document.getElementById("msgModal"), { keyboard: false }); + msgModal.show(); + + // 加载数据：调用后端接口 /api/task/getnodeinstr 或其他接口获取数据 + // 这里仅作示例使用模拟数据 + loadMsgData(); +} + +async function loadMsgData(){ + task_id = cur_task_id; + node_path = selectedNodeData.node_path; + try { + const res = await fetch("/api/task/nodegetmsg", { + 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(); + submittedMsgs = data.submitted || []; + pendingMsg = data.pending || {}; //one_llm = {'llm_type': llm_type, 'result': str_res} + + submittedPage = 1; + renderSubmittedTable(submittedPage); + // 填充待提交区域 --- + document.getElementById("llmtype").value = pendingMsg.llm_type || "0"; + document.getElementById("pendingContent").value = pendingMsg.result || ""; + } catch (error) { + console.error("加载Msg数据异常:", error); + } +} + +function renderSubmittedTable(page){ + const tbody = document.getElementById("submittedTbody"); + const start = (page - 1) * pageSize; + const end = start + pageSize; + const pageData = submittedMsgs.slice(start, end); + tbody.innerHTML = ""; + pageData.forEach((item, index) => { + const tr = document.createElement("tr"); + + // 第一列：序号 + const tdIndex = document.createElement("td"); + tdIndex.textContent = start + index + 1; + tr.appendChild(tdIndex); + + // 第二列：角色 + const tdRole = document.createElement("td"); + tdRole.textContent = item.role; + tr.appendChild(tdRole); + + // 第三列：内容 + const tdContent = document.createElement("td"); + tdContent.textContent = item.content; + tr.appendChild(tdContent); + + tbody.appendChild(tr); + }); + // 不足 10 行时，补空行 + for(let i = pageData.length; i < pageSize; i++){ + const tr = document.createElement("tr"); + tr.innerHTML = `   `; + tbody.appendChild(tr); + } +} + +// 分页按钮事件 +document.getElementById("submittedPrev").addEventListener("click", function(e){ + e.preventDefault(); + if(submittedPage > 1){ + submittedPage--; + renderSubmittedTable(submittedPage); + } +}); +document.getElementById("submittedNext").addEventListener("click", function(e){ + e.preventDefault(); + const maxPage = Math.ceil(submittedMsgs.length / pageSize); + if(submittedPage < maxPage){ + submittedPage++; + renderSubmittedTable(submittedPage); + } +}); + +// 导出功能：导出当前页已提交数据到 CSV +document.getElementById("btnExportSubmitted").addEventListener("click", function(){ + exportSubmittedCurrentPage(); +}); +function exportSubmittedCurrentPage(){ + const start = (submittedPage - 1) * pageSize; + const end = start + pageSize; + const pageData = submittedMsgs.slice(start, end); + let csv = "\uFEFF" + "序号,角色,内容\n"; // 添加 BOM 防乱码 + pageData.forEach((item, i) => { + csv += `${start + i + 1},${item.role},${item.content}\n`; + }); + // 补空行 + for(let i = pageData.length; i < pageSize; i++){ + csv += ",,\n"; + } + const blob = new Blob([csv], { type: "text/csv;charset=utf-8;" }); + const url = URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = "已提交消息.csv"; + a.click(); + URL.revokeObjectURL(url); +} + +//新增请求指令的msg +document.getElementById("btnNeedInstr").addEventListener("click",function(){ + if(selectedNodeData.workstatus === "0"){ + if (!confirm("是否确认要为该节点新增请求指令的信息？")) return; + // 获取用户在待提交区输入的新值 + newType = 3; + newContent = "请针对该节点信息，生成下一步渗透测试指令。"; + //提交到后端更新 + bsuccess = putnewmsg(newType,newContent); + if(bsuccess){ + //更新缓存 + selectedNodeData.workstatus = "3"; + const nodeEl = document.querySelector(`.tree-node[data-node_path="${selectedNodeData.node_path}"]`); + if (nodeEl) { + // 更新 DOM 属性（属性值均为字符串） + nodeEl.setAttribute("data-node_workstatus", "3"); + nodeEl.classList.remove("no-work"); + } + //更新界面 + strnew = getWorkStatus_Str(3); + document.getElementById("node_workstatus").textContent = strnew; + document.getElementById("llmtype").value = newType; + document.getElementById("pendingContent").value = newContent; + } + }else { + alert("只允许在-无待执行任务状态下新增请求指令的msg!") + } +}); +// 保存待提交内容修改 +document.getElementById("btnSavePending").addEventListener("click", function(){ + if(selectedNodeData.workstatus === "3"){ + if (!confirm("是否确认要保存对该节点待提交信息的修改？")) return; + // 获取用户在待提交区输入的新值 + const newType = document.getElementById("llmtype").value; + const newContent = document.getElementById("pendingContent").value; + //提交到后端更新 + putnewmsg(newType,newContent); + }else { + alert("只允许在-待提交llm状态下修改或新增msg!") + } + + +}); +async function putnewmsg(llmtype,content){ + task_id = cur_task_id; + node_path = selectedNodeData.node_path; + try { + const res = await fetch("/api/task/nodeupdatemsg", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({task_id,node_path,llmtype,content}), + }); + if (!res.ok) { + const errorData = await res.json(); + throw new Error(errorData.error || `HTTP错误 ${res.status}`); + } + const data = await res.json(); + bsuccess = data.bsuccess; + if(bsuccess){ + alert("修改成功") + pendingMsg.llmtype = llmtype; + pendingMsg.content = content; + return true; + } + else{ + alert("修改失败：",data.error) + return false; + } + } catch (error) { + console.error("加载Msg数据异常:", error); + return false; + } +} + + + //---------------------添加信息modal------------------------------ + document.getElementById("btnAddInfo").addEventListener("click", () => { + if (!checkSelectedNode()) return; + alert("该功能实现中..."); + }); + document.getElementById("btnAddChild").addEventListener("click", () => { + if (!checkSelectedNode()) return; + alert("该功能实现中..."); + }); + + // 页面加载完成后，加载节点树 + //document.addEventListener("DOMContentLoaded", loadNodeTree); \ No newline at end of file diff --git a/web/main/static/resources/scripts/task_manager.js b/web/main/static/resources/scripts/task_manager.js new file mode 100644 index 0000000..f92c720 --- /dev/null +++ b/web/main/static/resources/scripts/task_manager.js @@ -0,0 +1,447 @@ +let task_list = [] +let cur_task = null //当前选择的task--用于修改缓存时使用 +let cur_task_id = 0 //当前选择的cur_task_id +let ws = null + +// 页面卸载时断开连接 +window.addEventListener("beforeunload", function() { + if (ws) { + ws.close(); + ws =null; + } + task_list = [] + cur_task = null; + cur_task_id = 0; +}); + +// 页面加载完成后调用接口获取任务列表数据 +document.addEventListener("DOMContentLoaded", async () => { + //建立wbsocket + initWebSocket() + //获取左侧任务列表 + getTasklist(); + //当前选中的数据要清空 + cur_task = null; + cur_task_id = 0; + //任务基本信息界面初始化 + document.getElementById("detailTestTarget").textContent = "-"; + document.getElementById("detailTestStatus").textContent = "-"; + document.getElementById("detailSafeStatus").textContent = '-'; + //单选按钮 + set_radio_selection('testMode', 'auto'); + setSetpBtnStatus(); + //节点树界面初始化 + update_select_node_data_show("-","-","-","-","-",false) + //单选按钮点击事件 ------ 是不是更规范的编程方式，应该把控件事件的添加都放页面加载完成后处理 + const autoRadio = document.getElementById("autoMode"); + const manualRadio = document.getElementById("manualMode"); + autoRadio.addEventListener("click", () => updateTestMode(1)); + manualRadio.addEventListener("click", () => updateTestMode(0)); + + //双击任务列表节点读取数据 + searchInstructions(); + searchVulnerabilities(); + // renderTableRows(document.querySelector("#instrTable tbody"), []); + // renderTableRows(document.querySelector("#vulTable tbody"), []); + +}); + +//----------------------左侧任务列表----------------------- +function getstrsafeR(safeRank){ + if(safeRank === 0){ + safeR = "安全"; + } + else{ + safeR = "存在风险"; + } + return safeR; +} +function getstrTaskS(taskStatus){ + if(taskStatus === 0){ + taskS = "暂停中"; + }else if(taskStatus === 1){ + taskS = "执行中"; + }else { + taskS = "已完成"; + } + return taskS; +} + +async function getTasklist(){ + try { + const res = await fetch("/api/task/getlist"); + if (!res.ok) { + const errorData = await res.json(); + throw new Error(errorData.error || `HTTP错误 ${res.status}`); + } + const data = await res.json(); + task_list = data.tasks + const taskList = document.getElementById("taskList"); + taskList.innerHTML = ""; // 清空“加载中”提示 + // 遍历任务数组，生成任务项 + task_list.forEach((task) => { + const taskItem = document.createElement("div"); + taskItem.dataset.taskID =task.taskID //在taskItem添加数据属性，这是关联数据的第二种方法，selectedEl.dataset.taskId; 感觉比cur_task更好 + taskItem.className = "task-item"; + + // 第一行：测试目标 + const targetDiv = document.createElement("div"); + targetDiv.className = "task-target"; + targetDiv.textContent = task.testTarget; + taskItem.appendChild(targetDiv); + + // 第二行：测试状态，带缩进 + const statusDiv = document.createElement("div"); + statusDiv.className = "task-status"; + let safeR = getstrsafeR(task.safeRank); + let taskS = getstrTaskS(task.taskStatus); + statusDiv.textContent = `${taskS}-${safeR}`; + taskItem.appendChild(statusDiv); + + // 可绑定点击事件：点击任务项更新右侧详情 + taskItem.addEventListener("click", () => { + // 取消所有任务项的选中状态 + document.querySelectorAll(".task-item.selected").forEach(item => { + item.classList.remove("selected"); + }); + // 给当前点击的任务项添加选中状态 + taskItem.classList.add("selected"); + //执行业务代码 --参数当前选中task的数据 + cur_task = task; + selected_task_item() + }); + taskList.appendChild(taskItem); + }); + } catch (error) { + console.error("加载任务列表出错：", error); + document.getElementById("taskList").innerHTML = "

        加载任务列表失败！

        "; + } +} + +//选中tasklist--更新界面数据 +function selected_task_item(){ + if(cur_task_id === cur_task.taskID) return; + cur_task_id = cur_task.taskID; + //按钮状态更新 + actionButton = document.getElementById("actionButton"); + if(cur_task.taskStatus === 0){ + actionButton.textContent = "继续"; + }else if(cur_task.taskStatus === 1){ + actionButton.textContent = "暂停"; + }else { + } + //基本信息 + let safeR = getstrsafeR(cur_task.safeRank); + let taskS = getstrTaskS(cur_task.taskStatus); + document.getElementById("detailTestTarget").textContent = cur_task.testTarget; + document.getElementById("detailTestStatus").textContent = taskS; + document.getElementById("detailSafeStatus").textContent = safeR; + //单选按钮 + if(cur_task.workType === 0){ //人工 + set_radio_selection('testMode', 'manual'); + }else { //1-自动 + set_radio_selection('testMode', 'auto'); + } + //更新单步按钮 + setSetpBtnStatus() + //加载任务其他信息--node_tree.js + loadNodeTree(cur_task_id); +} + +//--------------------任务基本信息区域-------------------- +//单选按钮--测试模式修改 +async function updateTestMode(mode){ //0-人工，1-自动 + if(cur_task){ + if(cur_task.workType !== mode){ + if(cur_task.taskStatus === 1){ + alert("执行状态，不允许修改测试模式！"); + if( cur_task.workType === 0){ //人工 + set_radio_selection('testMode', 'manual'); + }else { //1-自动 + set_radio_selection('testMode', 'auto'); + } + } + else {//不是执行状态，可以修改测试模式 + try { + const res = await fetch("/api/task/taskworktype", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ cur_task_id,mode }), //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; + if(bsuccess){ + cur_task.workType = mode; //更新前端缓存 + //更新单步按钮状态 + setSetpBtnStatus() + }else { + alert("修改测试模失败！") + if( cur_task.workType === 0){ //人工 修改失败还原单选按钮点击状态 + set_radio_selection('testMode', 'manual'); + }else { //1-自动 + set_radio_selection('testMode', 'auto'); + } + } + } catch (error) { + console.error("控制任务状态异常:", error); + alert("控制任务状态异常:",error); + if( cur_task.workType === 0){ //人工 修改失败还原单选按钮点击状态 + set_radio_selection('testMode', 'manual'); + }else { //1-自动 + set_radio_selection('testMode', 'auto'); + } + } + } + } + } + + +} + +//点击暂停/继续按钮 +document.getElementById("actionButton").addEventListener("click",() => { + controlTask(); +}); +async function controlTask(){ + if(cur_task_id === 0){ + alert("请先选择一个任务！") + return + } + try { + const res = await fetch("/api/task/taskcontrol", { + 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(); + newstatus = data.newstatus; + //更新页面 + if(newstatus === 0){ + document.getElementById("detailTestStatus").textContent = "暂停中"; + actionButton.textContent = "继续"; + }else if(newstatus === 1){ + document.getElementById("detailTestStatus").textContent = "执行中"; + actionButton.textContent = "暂停"; + }else { + document.getElementById("detailTestStatus").textContent = "已完成"; + actionButton.textContent = "已完成"; + } + cur_task.taskStatus = newstatus; //光有个cur_taks也可以 + setSetpBtnStatus() + //更新task_list的显示 + updateTaskList() + } catch (error) { + console.error("控制任务状态异常:", error); + alert("控制任务状态异常:",error); + } +} + +//修改了涉及到tasklist的展示内容，修改tasklist显示 +function updateTaskList(){ + //更新数据 + const selectedEl = document.querySelector(".task-item.selected"); + if (selectedEl) { + let safeR = getstrsafeR(cur_task.safeRank); + let taskS = getstrTaskS(cur_task.taskStatus); + const statusDiv = selectedEl.querySelector(".task-status"); + statusDiv.textContent = `${taskS}-${safeR}`; + } +} + +//设置单步按钮的可点击状态状态 +function setSetpBtnStatus(){ + if(cur_task){ + task_status = cur_task.taskStatus; + work_type = cur_task.workType; + }else { + task_status = 0; + work_type = 0; + } + btn_TaskStep = document.getElementById("one_step"); + btn_NodeStep = document.getElementById("btnNodeStep"); + if(task_status===1 && work_type===0){ //执行中且是人工模式 + btn_TaskStep.disabled= false; + btn_TaskStep.classList.remove("disabled-btn"); + //node-step + if (selectedNodeData) { + if(selectedNodeData.node_bwork){ //有选中node,且节点为工作状态 + btn_NodeStep.disabled= false; + btn_NodeStep.classList.remove("disabled-btn"); + } + else { + btn_NodeStep.disabled = true; + btn_NodeStep.classList.add("disabled-btn"); //css会去重 + } + } + }else{ //其他情况都是不可用状态 + btn_TaskStep.disabled = true; // 添加 disabled 属性 + btn_TaskStep.classList.add("disabled-btn"); // 添加自定义样式 + if (selectedNodeData) { + btn_NodeStep.disabled = true; + btn_NodeStep.classList.add("disabled-btn"); //css会去重 + } + } +} + +//单步按钮--任务单步 +document.getElementById("one_step").addEventListener("click",() => { + if(cur_task_id===0){ + return + } + one_step_task(); +}); +async function one_step_task(){ + try { + const res = await fetch("/api/task/taskstep", { + 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(); + const bsuccess = data.bsuccess; + if(bsuccess){ + alert("该任务已提交单步工作，请稍候查看执行结果！") + } + else{ + error = data.erroe; + alert("该任务单步失败！",error) + } + }catch (error) { + alert("该节点单步失败，请联系管理员！", error); + } +} + +//------------------测试数据和漏洞数据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 < 10; 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); + } +} + +// 查询测试指令 +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); + } + } + +// 查询漏洞数据 +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(); + renderTableRows(document.querySelector("#vulTable tbody"), data.vuls || []); + document.getElementById("vulPrev").dataset.page = page > 1 ? page - 1 : 1; + document.getElementById("vulNext").dataset.page = page + 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)); +}); + +// 绑定漏洞数据查询按钮事件 +document.getElementById("vulSearchBtn").addEventListener("click", () => { + searchVulnerabilities(); +}); +// 绑定漏洞数据分页点击事件 +document.getElementById("vulPrev").addEventListener("click", (e) => { + e.preventDefault(); + searchVulnerabilities(parseInt(e.target.dataset.page)); +}); +document.getElementById("vulNext").addEventListener("click", (e) => { + e.preventDefault(); + searchVulnerabilities(parseInt(e.target.dataset.page)); +}); diff --git a/web/main/static/resources/scripts/warn_manager.js b/web/main/static/resources/scripts/warn_manager.js deleted file mode 100644 index 84f505a..0000000 --- a/web/main/static/resources/scripts/warn_manager.js +++ /dev/null @@ -1,246 +0,0 @@ -let modelMap = []; //model_name model_id -let channelMap = []; //channel_name channel_id -let warn_data = []; //查询到的报警数据 -let page_data = []; -let currentEditingRow = null; -let currentPage = 1; -const rowsPerPage = 30; -//页面加载初始化 -document.addEventListener('DOMContentLoaded', function () { - perWarnHtml() - - document.getElementById('delPageButton').addEventListener('click', function () { - delpageWarn(); - }); -}); - -//搜索按钮点击 -document.getElementById('searchMButton').addEventListener('click', function() { - shearchWarn(); -}); - -//查询数据 -async function shearchWarn(){ - //查询告警数据 - let modelName = document.getElementById('modelSelect').value; - let channelName = document.getElementById('channelSelect').value; - const startTime = document.getElementById('startTime').value; - const endTime = document.getElementById('endTime').value; - const sCount = 0; // 起始记录数从0开始 - const eCount = 5000; // 每页显示10条记录 - - if(modelName == "请选择"){ - modelName = ""; - } - if(channelName == "请选择"){ - channelName = ""; - } - - // 构造请求体 - const requestData = { - model_name: modelName || "", // 如果为空，则传空字符串 - channel_name: channelName || "", - start_time: startTime || "", - end_time: endTime || "", - s_count: sCount, - e_count: eCount - }; - try{ - // 发送POST请求到后端 - const response = await fetch('/api/warn/search_warn', { - method: 'POST', - headers: { - 'Content-Type': 'application/json' - }, - body: JSON.stringify(requestData) // 将数据转为JSON字符串 - }); - - // 检查响应是否成功 - if (response.ok) { - warn_data = await response.json(); - // 在这里处理查询结果，比如更新表格显示数据 - currentPage = 1; // 重置当前页为第一页 - renderTable(); //刷新表格 - renderPagination(); - } else { - console.error('查询失败:', response.status); - } - } catch (error) { - console.error('请求出错:', error); - } -} - -async function perWarnHtml() { - //获取算法和通道列表，在下拉框显示 - try{ - //算法名称下拉框 - let response = await fetch('/api/model/list'); - if (!response.ok) { - throw new Error('Network response was not ok'); - } - model_datas = await response.json(); - model_select_datas = ["请选择"]; - model_datas.forEach(option => { - model_select_datas.push(option.name); - modelMap[option.name] = option.ID; - }); - set_select_data("modelSelect",model_select_datas); - - //视频通道下拉框 - response = await fetch('/api/channel/tree'); - if (!response.ok) { - throw new Error('Network response was not ok'); - } - channel_datas = await response.json(); - channel_select_datas = ["请选择"]; - channel_datas.forEach(option => { - channel_select_datas.push(option.channel_name); - channelMap[option.channel_name] = option.ID; - }); - set_select_data("channelSelect",channel_select_datas); - //查询数据 - shearchWarn() - }catch (error) { - console.error('Error fetching model data:', error); - } -} - -//刷新表单页面数据 -function renderTable() { - const tableBody = document.getElementById('table-body-warn'); - tableBody.innerHTML = ''; //清空 - - const start = (currentPage - 1) * rowsPerPage; - const end = start + rowsPerPage; - pageData = warn_data.slice(start, end); - const surplus_count = rowsPerPage - pageData.length; - - - pageData.forEach((warn) => { - const row = document.createElement('tr'); - row.innerHTML = ` - ${warn.ID} - ${warn.model_name} - ${warn.channel_name} - ${warn.creat_time} - - - - - - `; - tableBody.appendChild(row); - row.querySelector('.warn-show-btn').addEventListener('click', () => showWarn(row)); - row.querySelector('.warn-video-btn').addEventListener('click', () => videoWarn(row)); - row.querySelector('.warn-delete-btn').addEventListener('click', () => deleteWarn(row)); - }); -} - -//刷新分页标签 -function renderPagination() { - const pagination = document.getElementById('pagination-warn'); - pagination.innerHTML = ''; - - const totalPages = Math.ceil(warn_data.length / rowsPerPage); - for (let i = 1; i <= totalPages; i++) { - const pageItem = document.createElement('li'); - pageItem.className = 'page-item' + (i === currentPage ? ' active' : ''); - pageItem.innerHTML = `${i}`; - pageItem.addEventListener('click', (event) => { - event.preventDefault(); - currentPage = i; - renderTable(); - renderPagination(); - }); - pagination.appendChild(pageItem); - } -} - -//显示报警信息详情 -function showWarn(row){ - currentEditingRow = row; - model_name = row.cells[1].innerText; - channel_name = row.cells[2].innerText; - create_time = row.cells[3].innerText; - img_path = pageData[row.rowIndex-1].img_path; - - document.getElementById('modelName').innerText = `${model_name}`; - document.getElementById('channleName').innerText = `${channel_name}`; - document.getElementById('warnTime').innerText = `${create_time}`; - document.getElementById('warnImg').src = `/api/warn/warn_img?path=${img_path}`; // 设置图片的 src - - $('#showWarn').modal('show'); -} - -//下载视频按钮 -function videoWarn(row){ - video_path = pageData[row.rowIndex-1].video_path; - // const videoPlayer = document.getElementById('videoPlayer'); - // videoPlayer.src = `/api/warn/warn_video?path=${encodeURIComponent(video_path)}`; - // videoPlayer.load(); // 确保重新加载视频 - // $('#showVideo').modal('show'); - // 创建下载链接 - const downloadUrl = `/api/warn/warn_video?path=${encodeURIComponent(video_path)}`; - const a = document.createElement('a'); - a.href = downloadUrl; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); -} - -//删除一条报警数据 -function deleteWarn(row){ - if (confirm('确定删除此报警吗？')) { - let alarmIds=[]; - warn_id = row.cells[0].innerText; - alarmIds.push(warn_id); - // 发送POST请求到后端 - fetch('/api/warn/warn_del', { - method: 'POST', - headers: { - 'Content-Type': 'application/json' - }, - body: JSON.stringify({ alarmIds: alarmIds }) - }) - .then(response => response.json()) - .then(data => { - if (data.success) { - alert('删除成功'); - } else { - alert('删除失败'); - } - }) - .catch(error => { - console.error('Error:', error); - }); - } -} - -//删除当前页的报警数据 -function delpageWarn(){ - if (confirm('确定删除此页报警吗？')) { - let alarmIds=[]; - pageData.forEach((warn) => { - alarmIds.push(warn.ID) - }); - // 发送POST请求到后端 - fetch('/api/warn/warn_del', { - method: 'POST', - headers: { - 'Content-Type': 'application/json' - }, - body: JSON.stringify({ alarmIds: alarmIds }) - }) - .then(response => response.json()) - .then(data => { - if (data.success) { - alert('删除成功'); - } else { - alert('删除失败'); - } - }) - .catch(error => { - console.error('Error:', error); - }); - } -} \ No newline at end of file diff --git a/web/main/templates/assets_manager.html b/web/main/templates/assets_manager.html new file mode 100644 index 0000000..30ea576 --- /dev/null +++ b/web/main/templates/assets_manager.html @@ -0,0 +1,15 @@ +{% extends 'base.html' %} + +{% block title %}ZFSAFE{% endblock %} + + +{% block style %} +{% endblock %} + + +{% block content %} +{% endblock %} + + +{% block script %} +{% endblock %} \ No newline at end of file diff --git a/web/main/templates/base.html b/web/main/templates/base.html index b9bbb53..92f1232 100644 --- a/web/main/templates/base.html +++ b/web/main/templates/base.html @@ -6,6 +6,7 @@ {% block title %}My Website{% endblock %} + {% block style_link %}{% endblock %} - - -
        - -
        -
        -
        -
          -
        • 一区 -
            -
          • 北门通道一
            • -
          • 南门通道二
            • -
          • 通道三
            • -
          -
          • -
        • 二区域 -
            -
          • 通道一
            • -
          -
          • -
        -
        -
        -
        - - -
        -
        -
        Video Stream
        -
        Video Stream
        -
        Video Stream
        -
        Video Stream
        -
        -
        -
        -
        - © 2024 ZFKJ All Rights Reserved -
        - - - - - diff --git a/web/main/templates/system_manager.html b/web/main/templates/system_manager.html index f6f36e2..30ea576 100644 --- a/web/main/templates/system_manager.html +++ b/web/main/templates/system_manager.html @@ -1,286 +1,15 @@ {% extends 'base.html' %} -{% block title %}ZFBOX{% endblock %} +{% block title %}ZFSAFE{% endblock %} + {% block style %} - .btn-blue { - background-color: #007bff; - color: white; - } - .btn-blue:hover { - background-color: #0056b3; - } - - .table-container { - min-height: 300px; /* 设置最小高度，可以根据需要调整 */ - } - /* 缩小表格行高 */ - .table-sm th, - .table-sm td { - padding: 0.2rem; /* 调整这里的值来改变行高 */ - } - .pagination { - justify-content: flex-end; /* 右对齐 */ - } - - .section-title { - position: relative; - margin: 20px 0; - } - .section-title hr { - margin: 0; - border-color: #000; /* You can change the color of the line if needed */ - } {% endblock %} + {% block content %} -
        - - - - - - - - - - -
        -
        -
        -

        1.0.0.1

        -
        -
        -
        -
        -
        -
        -
        -
        -

        1.0.0.1

        -
        -
        -
        -

        - 有线--192.168.3.45 WIFI--192.168.3.45 5G--192.168.3.45

        -
        -
        -
        -
        -
        -
        -
        -
        -
        -
        - -
        -
        -
        -
        - -
        -
        - - - - - - - - - - - -
        ID区域名称操作
        - -
        - -
        {% 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 new file mode 100644 index 0000000..4396a72 --- /dev/null +++ b/web/main/templates/task_manager.html @@ -0,0 +1,425 @@ +{% extends 'base.html' %} + +{% block title %}ZFSAFE{% endblock %} + + +{% block style_link %} + +{% endblock %} + + +{% block style %} + /* 搜索条件区域居中 */ + .search-area { + text-align: center; + margin-bottom: 15px; + } + /* 表格样式 */ + .table thead th { + text-align: center; /* 表头居中 */ + vertical-align: middle; + } + .table tbody td { + vertical-align: middle; + padding: 0.3rem 0.5rem; /* 缩小单元格上下内边距 */ + line-height: 1.5; + } + /* 序号列宽度较小 */ + .table thead th.seq-col, + .table tbody td.seq-col { + width: 50px; + text-align: center; + } + /* 分页控件右对齐 */ + .pagination { + justify-content: end; + } + /* 左侧任务列表整体样式 */ + .task-list { + height: calc(100vh - 210px); + overflow-y: auto; + background-color: #fff; + padding: 5px; + border-right: 1px solid #ddd; + } + /* 每个任务项样式 */ + .task-item { + padding: 5px; + margin-bottom: 5px; + background-color: #f9f9f9; + border: 1px solid #eee; + border-radius: 4px; + cursor: pointer; + transition: background-color 0.3s ease, box-shadow 0.3s ease; + } + /* 任务项选中状态样式 */ + .task-item.selected { + background-color: #e6f7ff; + border: 1px solid #1890ff; + box-shadow: 2px 2px 6px rgba(0, 0, 0, 0.2); + } + .task-item:hover { + background-color: #e6f7ff; + } + .task-item.selected { + border-color: #1890ff; + background-color: #bae7ff; + } + /* 任务目标样式，第一行 */ + .task-target { + font-weight: bold; + } + /* 任务状态，第二行，有缩进 */ + .task-status { + margin-left: 10px; + color: #666; + font-size: 0.8rem; + } + +/* 定义一个全高容器，与左侧任务列表高度保持一致 */ +.full-height { + height: calc(100vh - 210px); +} + +/* 右侧区域采用 flex 布局，垂直排列 */ +.right-container { + display: flex; + flex-direction: column; +} + +/* 基本信息区域，保持固定高度 */ +.basic-info { + /* 根据需要设定高度，或保持内容自适应 */ + flex: 0 0 auto; +} + +/* Tab 内容区域占满剩余空间 */ +.tab-wrapper { + flex: 1 1 auto; + display: flex; + flex-direction: column; + overflow: hidden; /* 防止内部溢出 */ +} + +/* 导航部分根据内容高度 */ +#myTab { + flex: 0 0 auto; +} + +.tab-content{ + flex: 1 1 auto; + overflow-y: auto; +} + +.disabled-btn { + /* 禁用状态样式 */ + background-color: #cccccc; /* 灰色背景 */ + color: #666666; /* 文字颜色变浅 */ + cursor: not-allowed; /* 鼠标显示禁用图标 */ + opacity: 0.7; /* 可选：降低透明度 */ + + /* 禁用点击事件（通过 disabled 属性已实现，此样式仅增强视觉效果） */ + pointer-events: none; /* 可选：彻底阻止鼠标事件 */ +} + +{% endblock %} +{% block modal %} +{% include 'task_manager_modal.html' %} +{% endblock %} + +{% block content %} +
        +
        +
        + +
        + +

        加载中...

        +
        + + +
        + +
        +
        +
        + + 192.168.1.110 +
        +
        +
        + + 执行中 +
        +
        + + 安全 +
        +
        + +
        + + +
        +
        + + +
        +
        +
        +
        + +
        + + + +
        +
        + + +
        + +
        + +
        +
        + +
        +
        + +
        + +
        + +
        +

        加载中...

        +
        +
        +
        + +
        +
        +
        节点信息
        +

        节点名称: -

        +

        测试状态: -

        +

        漏洞类型: -

        +

        漏洞级别: -

        +

        工作状态: -

        +

        执行状态: -

        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        +
        + + +
        + +
        +
        + +
        +
        + +
        +
        + + + + + + + + + + + + + +
        序号节点路径指令序号执行指令执行结果
        + + +
        + + +
        + +
        +
        + +
        +
        + +
        +
        + +
        +
        + +
        +
        + + + + + + + + + + + + + +
        序号节点路径漏洞类型漏洞级别漏洞说明
        + + +
        +
        +
        +
        +
        +
        +
        +{% endblock %} + + +{% block script %} + + + + + +{% endblock %} \ No newline at end of file diff --git a/web/main/templates/task_manager_modal.html b/web/main/templates/task_manager_modal.html new file mode 100644 index 0000000..d6ceb9e --- /dev/null +++ b/web/main/templates/task_manager_modal.html @@ -0,0 +1,213 @@ + + + + + \ No newline at end of file diff --git a/web/main/templates/user_manager.html b/web/main/templates/user_manager.html index f70c0f4..19b2109 100644 --- a/web/main/templates/user_manager.html +++ b/web/main/templates/user_manager.html @@ -1,6 +1,6 @@ {% extends 'base.html' %} -{% block title %}ZFBOX{% endblock %} +{% block title %}ZFSAFE{% endblock %} {% block style %} .note { diff --git a/web/main/templates/view_main.html b/web/main/templates/view_main.html deleted file mode 100644 index 7011713..0000000 --- a/web/main/templates/view_main.html +++ /dev/null @@ -1,143 +0,0 @@ -{% extends 'base.html' %} - -{% block title %}ZFBOX{% endblock %} - -{% block style %} - .nav-scroller { - position: relative; - z-index: 2; - height: 2.75rem; - overflow-y: hidden; - } - - .nav-scroller .nav { - display: flex; - flex-wrap: nowrap; - padding-bottom: 1rem; - margin-top: -1px; - overflow-x: auto; - text-align: center; - white-space: nowrap; - -webkit-overflow-scrolling: touch; - } - - .blue-svg { - color: blue; /* 这将影响所有使用currentColor的fill属性 */ - } - - main { - display: flex; - flex: 1; - overflow: hidden; - } - - .tree-view { - border-right: 1px solid #ddd; - padding: 10px; - overflow-y: auto; - } - - .video-frame { - position: relative; - width: calc(50% - 10px); /* 默认4画面布局，每行2个视频框架 */ - margin: 5px; - float: left; - background-color: #ccc; /* 默认灰色填充 */ - border: 1px solid #000; /* 边框 */ - } - .video-header { - display: flex; - justify-content: space-between; - align-items: center; - background-color: #2c3e50; - color: #fff; - padding: 5px; - } - .video-area { - width: 100%; - padding-bottom: 75%; /* 4:3 比例 */ - background-color: #000; /* 默认灰色填充 */ - position: relative; - border: 1px solid #ddd; /* 视频区域边框 */ - } - - .video-area canvas { - display: none; - position: absolute; - top: 0; - left: 0; - width: 100%; - height: 100%; - object-fit: contain; - } - - .video-buttons { - display: flex; - gap: 10px; - } - - .video-buttons button { - background: none; - border: none; - color: white; - cursor: pointer; - } - - .video-buttons button:hover { - color: #f39c12; - } - - .dropdown-menu { - min-width: 100px; - } - - - #videoGrid.four .video-frame { - width: calc(50% - 10px); /* 每行4个视频框架 */ - } - #videoGrid.eight .video-frame { - width: calc(12.5% - 10px); /* 每行8个视频框架 */ - } - #videoGrid.nine .video-frame { - width: calc(33.33% - 10px); /* 每行3个视频框架，9画面布局 */ - } - .toggle-buttons { - margin-bottom: 5px; - } - .btn-small { - font-size: 0.75rem; - padding: 0.25rem 0.5rem; - } - .error-message { - position: absolute; - top: 50%; - left: 50%; - transform: translate(-50%, -50%); - color: red; - background-color: rgba(255, 255, 255, 0.8); /* 半透明背景 */ - padding: 10px; - border-radius: 5px; - } -{% endblock %} - -{% block content %} -
        -
        - -
        - -
        -
        - -
        -
        - - -
        -
        -
        -{% endblock %} - -{% block script %} - -{% endblock %} diff --git a/web/main/templates/vul_manager.html b/web/main/templates/vul_manager.html new file mode 100644 index 0000000..30ea576 --- /dev/null +++ b/web/main/templates/vul_manager.html @@ -0,0 +1,15 @@ +{% extends 'base.html' %} + +{% block title %}ZFSAFE{% endblock %} + + +{% block style %} +{% endblock %} + + +{% block content %} +{% endblock %} + + +{% block script %} +{% endblock %} \ No newline at end of file