''' 渗透测试任务管理类 一次任务的闭合性要检查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.InstrM = InstructionManager(self) # 类对象渗透，要约束只读取信息 # 控制最大并发指令数量 self.max_thread_num = 2 self.task_id = 0 #任务id -- self.do_sn = 0 #指令执行顺序号--暂时已执行完成顺序记录 self.workth_list = [] #线程句柄list self.batch_num = 0 #一个批次的指令数量 self.long_instr_num = 0 #耗时指令数量 self.long_time_instr = ['nikto'] #耗时操作不计入批量执行的数量，不加不减 self.node_queue = [] # self.lock = threading.Lock() #线程锁 self.node_num = 0 #在处理Node线程的处理 self.do_sn_lock = threading.Lock() #指令执行顺序号锁 self.brun = True 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} work_node.add_res(res) #入节点结果队列 def do_worker_th(self): #线程的dbm需要一个线程一个 th_DBM = DBManager() th_DBM.connect() while self.brun: work_node = None with self.lock: if self.node_queue: # 获取测试节点 work_node = self.node_queue.pop(0) self.node_num += 1 if work_node: #开始执行指令 for instruction in work_node.instr_queue: 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) # 执行结果入队列 #保存记录--测试使用 with self.lock: self.node_num -=1 if self.node_num == 0 and len(self.node_queue) == 0: with open("attack_tree", 'wb') as f: pickle.dump(TM.CCM.attack_tree, f) #针对一个节点的指令执行完成后，提交LLM规划下一步操作 # node_list = self.CCM.get_llm_instruction(work_node) # if not node_list:#该节点测试完成 # continue # for node in node_list: # self.node_queue.put(node) else: time.sleep(10) # try: # instruction = self.instr_queue.get(block=False) #quere线程安全，block=false非阻塞get # self.doing_instr.put(instruction) #入执行队列--忘了入执行队列的作用 # #执行中会对指令进行微调，并有可能不执行直接返回空结果 # 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) #执行结果入队列 # self.done_instr.put(instruction) #执行完成队列 # #执行情况是否需要用户确认 # if ext_params["is_user"]: # pass # #print("该指令执行需要用户确认") # except queue.Empty: # time.sleep(10) #函数结束，局部变量自动释放 def start_task(self,target_name,target_in): #判断目标合法性 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 = "无" #? #启动--初始化指令 node_list = self.CCM.start_do(target,self.task_id) with self.lock: for node in node_list: self.node_queue.append(node) #创建工作线程----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 = 3 if test_type == 1: #测试执行指令 with open("attack_tree", "rb") as f: TM.CCM.attack_tree = pickle.load(f) # 遍历node，查看有res的数据 nodes = TM.CCM.attack_tree.traverse_bfs() for node in nodes: if node.instr_queue: # list TM.node_queue.append(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的数据 nodes = TM.CCM.attack_tree.traverse_bfs() for node in nodes: if node.res_quere: #list node_list = TM.CCM.get_llm_instruction(node) if not node_list:#该节点测试完成 continue with TM.lock: for do_node in node_list: TM.node_queue.append(do_node) # 暂存状态-- with open("attack_tree", 'wb') as f: pickle.dump(TM.CCM.attack_tree, f) elif test_type == 3: #新目标测试 # 启动--初始化指令 node_list = TM.CCM.start_do("192.168.204.137", 0) with TM.lock: for node in node_list: TM.node_queue.append(node) #暂存状态-- with open("attack_tree",'wb') as f: pickle.dump(TM.CCM.attack_tree,f) elif test_type == 4: #读取messages pass else: #完整过程测试---要设定终止条件 pass