zfsafe/mycode/ControlCenter.py

#自动化测试逻辑规则控制
#统一控制规则  和 渗透测试树的维护
import json
import re
import queue
import time
import threading
from mycode.AttackMap import AttackTree
from mycode.AttackMap import TreeNode
from myutils.MyLogger_logger import LogHandler
from mycode.DBManager import DBManager

class ControlCenter:
    def __init__(self):
        self.logger = LogHandler().get_logger("ControlCenter")

    def __del__(self):
        self.brun =False
        self.task_id = None
        self.target = None
        self.attack_tree = None

    def init_cc_data(self):
        #一次任务一次数据
        pass

    def get_user_init_info(self):
        '''开始任务初，获取用户设定的基础信息，初始信息可以分为两块：
        1.提交llm的补充信息 和 保留在本地的信息（如工具补充参数等，cookie）
        2.用户可以设置全局，和指定节点（端口）
        3.补充测试节点
        '''
        # ?包括是否对目标进行初始化的信息收集
        return {"已知信息":"无"}


    def restore_one_llm_work(self,node,llm_type,res_list):
        node.llm_type = llm_type
        node.res_quere = res_list

    #需要用户确认的信息--待完善
    def need_user_know(self,strinfo,node):
        pass

    #待修改
    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.append(instr)

    def stop_do(self):
        #清空数据
        self.task_id = None
        self.target = None
        self.attack_tree = None
        #停止llm处理线程
        self.brun =False
        for th in self.llmth_list:
            th.jion()