zfsafe/TaskManager.py

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