zfsafe/TaskManager.py

'''
渗透测试任务管理类  一次任务的闭合性要检查2025-3-10  一次任务后要清理LLM和InstrM的数据
'''
from TargetManager import TargetManager  # 从模块导入类
#from LLMManager import LLMManager        # 同理修正其他导入
from mycode 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 = queue.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':instr,'reslut':reslut}
        work_node.add_res(res)  #入节点结果队列

    def do_worker_th(self):
        #线程的dbm需要一个线程一个
        th_DBM = DBManager()
        th_DBM.connect()
        while self.brun:
            try:
                with self.lock:
                    self.node_num += 1

                #获取测试节点
                work_node = self.node_queue.get(block=False)
                #开始执行指令
                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 self.node_queue.empty():
                        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)
            except queue.Empty:
                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)
            for node in node_list:
                self.node_queue.put(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)
        with open("node_list", 'rb') as f:
            TM.node_queue = pickle.load(f)

        #创建线程执行指令
        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 not node.res_quere.empty():
                node_list = TM.CCM.get_llm_instruction(node)
                if not node_list:#该节点测试完成
                    continue
                for do_node in node_list:
                    TM.node_queue.put(do_node)
    elif test_type ==3: #新目标测试
        # 启动--初始化指令
        node_list = TM.CCM.start_do("192.168.204.137", 0)
        for node in node_list:
            TM.node_queue.put(node)
        #暂存状态--
        with open("attack_tree",'wb') as f:
            pickle.dump(TM.CCM.attack_tree,f)
        with open("node_list",'wb') as f:
            pickle.dump(TM.node_queue,f)

    else:
        #完整过程测试---要设定终止条件
        pass