''' 渗透测试任务管理类 一次任务的闭合性要检查2025-3-10 一次任务后要清理LLM和InstrM的数据 ''' from TargetManager import TargetManager # 从模块导入类 from LLMManager import LLMManager # 同理修正其他导入 from myutils.FileManager import FileManager from InstructionManager import InstructionManager from mycode.DBManager import DBManager from myutils.MyTime import get_local_timestr import queue import time import os import threading class TaskManager: def __init__(self): self.TargetM = TargetManager() # 生成功能对象 self.LLMM = LLMManager(1) self.InstrM = InstructionManager() self.DBM = DBManager() #主进程一个DBM self.DBM.connect() # 控制最大并发指令数量 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.instr_queue = queue.Queue() #线程安全 --待执行指令 self.user_instr = queue.Queue() #需要用确认或手动执行的命令--待执行指令 self.user_done_instr = queue.Queue() #执行完成需要用户确认的指令 self.doing_instr = queue.Queue()#执行中的指令 self.done_instr = queue.Queue() #执行完成的指令 self.res_queue = queue.Queue() #结果队列 self.lock = threading.Lock() #线程锁 self.do_sn_lock = threading.Lock() #指令执行顺序号锁 self.brun = True def is_user_instr(self,instr): ''' 过滤需要人工确认或手动执行的指令 ---- 待完善 :param instr: :return: ''' #if instr.startswith("curl") or instr.startswith("http") or instr.startswith("wget"): if instr.startswith("http") or instr.startswith("wget") or instr.startswith("ssh"): return True def instr_in_quere(self,instr_list): ''' 对于运行需要较长时间的不强求同一批次返回给LLM :param instr_list: :return: ''' for instr in instr_list: if self.is_user_instr(instr): self.user_instr.put(instr) print(f"需要人工确认的指令{instr}") else: matched =False for prefix in self.long_time_instr: if instr.startswith(prefix): matched =True if not matched: with self.lock: self.batch_num += 1 #非耗时指令+1 print(f"&&&&&&当前batch_num:{self.batch_num}") else: with self.lock: self.long_instr_num +=1 #耗时指令数量+1 # 指令入队列 self.instr_queue.put(instr) #入数据库放哪一层要待定-TaskManager , InstructionM,ToolBase-2025-3-10 def res_in_db(self,bres,instr,reslut,start_time,end_time,th_DBM,source_result,ext_params): if th_DBM.ok: with self.do_sn_lock: self.do_sn += 1 #指令的执行序列是一个任务共用,要线程锁，错误问题也不大 th_DBM.insetr_result(self.task_id,instr,reslut,self.do_sn,start_time,end_time,source_result,ext_params) else: print("数据库连接失败！！") def res_in_quere(self,bres,instr,reslut,start_time,end_time,th_DBM,source_result,ext_params): ''' 执行结果入队列，若批量执行的指令都完成，则提交LLM生成下一步指令 :param bres: :param instr: :param reslut: :return: ''' matched = False bover = False #入数据库 -- bres True和False 都入数据库2025-3-10 self.res_in_db(bres,instr,reslut,start_time,end_time,th_DBM,source_result,ext_params) #结果入队列 if bres: res = {'instr':instr,'reslut':reslut} self.res_queue.put(res) #入队列 else: #对于不需要再提交给LLM的结果，如何处理待定。有执行false、未知工具 pass #判断批次指令是否都执行完 for prefix in self.long_time_instr: if instr.startswith(prefix): matched = True if not matched: with self.lock: self.batch_num -= 1 #非耗时指令-1 print(f"&&&&&&当前batch_num:{self.batch_num}") if self.batch_num ==0: #只会轮询到最后一个线程才会是0 bover = True else: #这里会有个问题，若耗时指令执行完，结果入队列后，但LLM没有进一步的指令下发，会造成结果不会提交 #需要在生成“生成报告”的指令时进行研判--该点很重要 with self.lock: self.long_instr_num -=1 #耗时指令数量-1 if bover: #整合结果提交 -- 需要确保只有一个线程会执行 self.batch_res_post(th_DBM) def batch_res_post(self,th_DBM): ''' 组合批量指令的结果，一起提交到LLM，生成下一步具体指令 :return: ''' post_string = "" while not self.res_queue.empty(): res = self.res_queue.get() str = f"执行指令：{res['instr']}的结果是：{res['reslut']}。" post_string = post_string + "\n" post_string = post_string + str if post_string: post_string = post_string + "\n请根据这些结果生成下一步具体的指令。" print(f"***************\n{post_string}") with open("res","w",encoding="utf-8") as f: f.write(post_string) #*****测试时中断下一步指令的获取 # 提交提示词，得到下一步指令 # instr_list = self.LLMM.get_llm_instruction(post_string,th_DBM) # if instr_list: # if instr_list[0] == "生成报告": #“生成报告”要特殊处理#？ # self.brun = False # print("生成报告--退出工作线程") # else: # 继续工作 # self.instr_in_quere(instr_list) def do_worker_th(self): #线程的dbm需要一个线程一个 th_DBM = DBManager() th_DBM.connect() while self.brun: 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 = "无" #? #启动--初始化指令 prompt = self.LLMM.build_initial_prompt(target,know_info) instr_list = self.LLMM.get_llm_instruction(prompt,self.DBM) self.instr_in_quere(instr_list) #指令入队列 #创建工作线程 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.LLMM.init_data() #清空一些全局变量 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) TM.LLMM.test_old_message(strMsg) #先设置message的值 test_type = 3 if test_type == 1: #测试执行指令 # instrS = ['gobuster dir -u https://58.216.217.70 -w /usr/share/wordlists/dirbuster/directory-list-2.3-medium.txt -k -t 20 --timeout 10s', # 'searchsploit SoftEther VPN', # 'curl -kv -X POST -d "username=admin&password=admin" https://58.216.217.70/vpn/index.html --connect-timeout 10', # 'gobuster dir -u http://58.216.217.70:10001 -w /usr/share/wordlists/dirbuster/directory-list-2.3-medium.txt -t 20 --timeout 10s'] #instrS = ['nmap -sV -sC -p- -Pn 192.168.204.137'] with open("test", "r", encoding="utf-8") as f: messages = json.load(f) text = messages[-1]["content"] #list = TM.LLMM.fetch_instruction(text)[9:10] list = TM.LLMM.fetch_instruction(text) #print(list) TM.instr_in_quere(list) #创建线程执行指令 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("res", "r", encoding="utf-8") as f: prompt = f.read() if prompt: instr_list = TM.LLMM.get_llm_instruction(prompt,TM.DBM) TM.instr_in_quere(instr_list) # 指令入队列问题不大，任务执行线程没有起 elif test_type ==3: #新目标测试 prompt = TM.LLMM.build_initial_prompt("192.168.204.137","无") if prompt: instr_list = TM.LLMM.get_llm_instruction(prompt,TM.DBM) TM.instr_in_quere(instr_list) # 指令入队列问题不大，任务执行线程没有起 else: #完整过程测试---要设定终止条件 pass