#自动化测试逻辑规则控制
#统一控制规则  和 渗透测试树的维护
import json
import re
import queue
from mycode.AttackMap import AttackTree
from mycode.AttackMap import TreeNode
from mycode.LLMManager import LLMManager
from myutils.ConfigManager import myCongif  #单一实例
from myutils.MyLogger_logger import LogHandler

class ControlCenter:
    def __init__(self,DBM):
        self.logger = LogHandler().get_logger("ControlCenter")
        self.task_id = None
        self.target = None
        self.attack_tree = None
        self.DBM = DBM
        #LLM对象
        self.LLM = LLMManager(myCongif.get_data("LLM_type"))

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

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

    def get_user_init_info(self):
        '''开始任务初，获取用户设定的基础信息'''
        # ?包括是否对目标进行初始化的信息收集
        return "无"

    def start_do(self,target,task_id):
        '''一个新任务的开始'''
        self.task_id = task_id
        self.target = target
        #创建测试树
        if self.attack_tree:
            self.attack_tree = None #释放
        root_node = TreeNode(target)
        self.attack_tree = AttackTree(root_node)#创建测试树，同时更新根节点相关内容
        #获取初始指令
        know_info = self.get_user_init_info()
        prompt = self.LLM.build_initial_prompt(target,know_info,root_node)
        node_cmds, commands = self.LLM.get_llm_instruction(prompt,self.DBM,root_node)
        # 更新tree
        self.tree_manager(node_cmds,root_node)
        # 分析指令入对应节点
        node_list = self.instr_in_node(commands,root_node)
        return node_list

    #约束1：一个节点只能同时提交一次，未测试的节点不要重复
    def get_llm_instruction(self,node):
        #拼接结果字符串--由于测试需要queue改成了list 2025-3-19
        # json_strings = []
        # if node.res_quere:
        #     for item in node.res_quere:
        #         json_str = json.dumps(item, ensure_ascii=False)
        #         json_strings.append(json_str)
        # res_str = ','.join(json_strings)

        res_str = json.dumps(node.res_quere,ensure_ascii=False)

        #构造本次提交的prompt
        user_Prompt = f'''
            当前分支路径：{node.path}
            当前节点信息：
            - 节点名称：{node.name}
            - 节点状态：{node.status}
            - 漏洞类型：{node.vul_type}
            上一步结果：{res_str}
            任务：生成下一步渗透测试指令或结束该节点渗透测试。
            '''
        node_cmds,commands = self.LLM.get_llm_instruction(user_Prompt,self.DBM,node)
        #更新tree
        self.tree_manager(node_cmds,node)
        #分析指令入对应节点
        node_list = self.instr_in_node(commands,node)
        return node_list

    def tree_manager(self,node_cmds,node):
        '''更新渗透测试树'''
        if not node_cmds or len(node_cmds)==0:
            return
        try:
            node_jsons = json.loads(node_cmds)
            for node_json in node_jsons:
                if node_json["action"] == "add_node":   #新增节点
                    parent_node_name = node_json["parent"]
                    node_name = node_json["node"]
                    status = node_json["status"]
                    #新增节点原则上应该都是当前节点增加子节点
                    if node.name == parent_node_name:
                        new_node = TreeNode(node_name,status)
                        node.add_child(new_node) #message的传递待验证
                    else:
                        self.logger.error("添加子节点时，遇到父节点名称不一致的，需要介入！！")
                elif node_json["action"] == "update_status":
                    node_name = node_json["node"]
                    status = node_json["status"]
                    vul_type = node_json["vulnerability"]
                    if node.name == node_name:
                        node.status = status
                        node.vul_type = vul_type
                    else:
                        self.logger.error("遇到不是修改本节点状态的，需要介入！！")
                else:
                    self.logger.error("node更新JSON遇到异常参数！")
        except json.JSONDecodeError as e:
            self.logger.error(f"JSON 解析失败：{e.msg}")
            self.logger.error(f"错误位置：第{e.lineno}行，列{e.colno}")
        except TypeError as e:
            self.logger.error(f"输入类型错误：{str(e)}")

    def instr_in_node(self,commands,node):
        node_list = []
        for command in commands:
            # 使用正则匹配方括号中的node_path（非贪婪模式）
            match = re.search(r'\[(.*?)\]', command)
            if match:
                node_path = match.group(1)
                find_node = self.attack_tree.find_node_by_nodepath(node_path)
                if find_node:
                    instruction = re.sub(r'\[.*?\]', "", command,count=1,flags=re.DOTALL)
                    find_node.instr_queue.append(instruction)
                    #入输出队列
                    if find_node not in node_list:
                        node_list.append(find_node)
                else:
                    self.logger.error(f"基于节点路径没有找到对应的节点{node_path}")
            else:
                self.logger.error(f"得到的指令格式不符合规范：{command}")
        return node_list

    #待修改
    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