''' 渗透测试任务管理类 一次任务的闭合性要检查2025-3-10 一次任务后要清理LLM和InstrM的数据 ''' from TargetManager import TargetManager # 从模块导入类 #from LLMManager import LLMManager # 同理修正其他导入 from mycode.ControlCenter import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。 from myutils.FileManager import FileManager from 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 = 6 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: self.logger.debug("暂无需要执行指令的节点！") 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 = 1 instr_index = 19 iput_index = -1 # 0是根节点 indexs = [] if test_type == 0: #新目标测试 # 启动--初始化指令 node_list = TM.CCM.start_do("192.168.204.137", 0) #异步处理，需要等待线程结束了 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 = [''' # '''] 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 == 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