zfsafe/TaskManager.py

'''
渗透测试任务管理类  一次任务的闭合性要检查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