zfsafe/mycode/AttackMap.py

import queue
import copy
import re
import threading

#渗透测试树结构维护类
class AttackTree:
    def __init__(self,root_node):
        #针对根节点处理
        self.root = root_node
        self.root.path = f"目标系统->{root_node.name}"


    def set_root(self,root_node):
        self.root = root_node

    def add_node(self,parent_name,new_node):
        """根据父节点名称添加新节点"""
        parent_node = self.find_node_by_name(parent_name)
        if parent_node:
            parent_node.add_child(new_node)
            return True
        return False

    def traverse_bfs(self):
        """广度优先遍历"""
        if not self.root:
            return []

        queue = [self.root]
        result = []
        while queue:
            current = queue.pop(0)
            result.append(current)
            queue.extend(current.children)
        return result

    def traverse_dfs(self, node=None, result=None):
        """深度优先遍历（前序遍历）"""
        if result is None:
            result = []
        if node is None:
            node = self.root
            if not node:
                return []

        result.append(node)
        for child in node.children:
            self.traverse_dfs(child, result)
        return result

    #生成节点树字典数据
    def node_to_dict(self,node):
        return {
            "node_name":node.name,
            "node_path":node.path,
            "node_status":node.status,
            "node_bwork":node.bwork,
            "node_vultype":node.vul_type,
            "node_vulgrade":node.vul_grade,
            "node_workstatus":node.get_work_status(),
            "children":[self.node_to_dict(child) for child in node.children] if node.children else []
        }

    #树简化列表，用户传输到前端
    def get_node_dict(self):
        node_dict = self.node_to_dict(self.root)    #递归生成
        return node_dict

    def find_node_by_name(self, name):
        """根据名称查找节点（广度优先）"""
        nodes = self.traverse_bfs()
        for node in nodes:
            if node.name == name:
                return node
        return None

    def find_node_by_nodepath_parent(self,node_path,node):
        node_names = node_path.split('->')
        node_name = node_names[-1]
        if node_name == node.name:#当前节点
            return node
        else:
            if node_names[-2] == node.name: #父节点是当前节点
                for child_node in node.children:
                    if child_node.name == node_name:
                        return child_node
                #走到这说明没有匹配到-则新建一个节点
                newNode = TreeNode(node_name,node.task_id)
                node.add_child(newNode)
                return newNode
            else:
                return None #约束：不处理

    def find_node_by_nodepath(self,node_path):
        '''基于节点路径查找节点，只返回找到的第一个节点，若有节点名称路径重复的情况暂不处理'''
        current_node = self.root #从根节点开始
        node_names = node_path.split('->')
        layer_num = 0
        for node_name in node_names:
            if node_name == "目标系统":
                layer_num +=1
                continue
            if node_name == current_node.name:#根节点开始
                layer_num += 1
                continue
            else:
                bfound = False
                for child_node in current_node.children:
                    if child_node.name == node_name:    #约束同一父节点下的子节点名称不能相同
                        current_node = child_node
                        layer_num += 1
                        bfound = True
                        break
                if not bfound:  #如果遍历子节点都没有符合的，说明路径有问题的，不处理中间一段路径情况
                    return  None
        return current_node

    #更新节点的bwork状态
    def update_node_bwork(self,node_path):
        node = self.find_node_by_nodepath(node_path)
        if not node:
            return False,False
        if node.bwork:
            node.bwork = False
        else:
            node.bwork = True
        return True,node.bwork

    def find_nodes_by_status(self, status):
        """根据状态查找所有匹配节点"""
        return [node for node in self.traverse_bfs() if node.status == status]

    def find_nodes_by_vul_type(self, vul_type):
        """根据漏洞类型查找所有匹配节点"""
        return [node for node in self.traverse_bfs() if node.vul_type == vul_type]

    #考虑要不要用tree封装节点的操作--待定
    def update_node_status(self, node_name, new_status):
        """修改节点状态"""
        node = self.find_node_by_name(node_name)
        if node:
            node.status = new_status
            return True
        return False

    def update_node_vul_type(self,node_name,vul_type):
        """修改节点漏洞类型"""
        node = self.find_node_by_name(node_name)
        if node:
            node.vul_type = vul_type
            return True
        return False

    def print_tree(self, node=None, level=0):
        """可视化打印树结构"""
        if node is None:
            node = self.root
        prefix = "    " * level + "|-- " if level > 0 else ""
        print(f"{prefix}{node.name} [{node.status}, {node.vul_type}]")
        for child in node.children:
            self.print_tree(child, level + 1)


class TreeNode:
    def __init__(self, name,task_id,status="未完成", vul_type="未发现"):
        self.task_id = task_id  #任务id
        self.name = name  # 节点名称
        #self.node_lock = threading.Lock()    #线程锁
        self.bwork = True  # 当前节点是否工作，默认True --停止/启动
        self.status = status  # 节点测试状态 -- 由llm返回指令触发更新
        #work_status需要跟两个list统一管理：初始0，入instr_queue为1，入instr_node_mq为2，入res_queue为3，入llm_node_mq为4，llm处理完0或1
        self._work_status = 0    #0-无任务，1-待执行测试指令，2-执行指令中，3-待提交Llm,4-提交llm中， 2025-4-6新增，用来动态显示节点的工作细节。
        self.vul_type = vul_type  # 漏洞类型--目前赋值时没拆json
        self.vul_name = ""
        self.vul_grade = ""
        self.vul_info = ""
        self.children = []  # 子节点列表
        self.parent = None  # 父节点引用
        self.path = ""      #当前节点的路径
        self.messages = []  # 针对当前节点积累的messages -- 针对不同节点提交不同的messages

        self.llm_type = 0   #llm提交类型 0--初始状态无任务状态，1--指令结果反馈，2--llm错误反馈
        self.llm_sn = 0     #针对该节点llm提交次数
        self._llm_quere = [] #待提交llm的数据

        self.do_sn = 0      #针对该节点instr执行次数
        self._instr_queue = []  #针对当前节点的待执行指令----重要约束：一个节点只能有一个线程在执行指令

        self.his_instr = []   #保留执行指令的记录｛“instr”:***,"result":***｝
        #用户补充信息
        self.cookie = ""
        self.ext_info = ""

    #设置用户信息
    def set_user_info(self,cookie,ext_info):
        self.cookie = cookie
        self.ext_info = ext_info

    def copy_messages(self,childe_node):
        '''
        子节点继承父节点的messages，目前规则保留上两层节点的message信息
        :param childe_node:
        :return:
        '''
        tmp_messages = copy.deepcopy(self.messages)
        if not self.parent:
            childe_node.messages = tmp_messages
        else:
            parent_path = self.parent.path
            bfind = False
            for msg in tmp_messages:
                if msg["role"] == "system":
                    childe_node.messages.append(msg)
                elif msg["role"] == "user":
                    if not bfind:
                        #获取user的node_path
                        content = msg["content"]
                        pattern = r"当前分支路径：(.+?)\n"
                        match = re.search(pattern, content)
                        if match:
                            path = match.group(1)
                            if parent_path in path:#当前节点的父节点路径在存入子节点messages
                                childe_node.messages.append(msg)
                                bfind = True    #后续messages都保留
                        else:
                            print("提交的用户提示词结构有问题！")
                    else:
                        childe_node.messages.append(msg)
                elif msg["role"] == "assistant":
                    if bfind:
                        childe_node.messages.append(msg)
                else:
                    print("非法的信息体类型！")

    #添加子节点
    def add_child(self, child_node):
        child_node.parent = self
        child_node.path = self.path + f"->{child_node.name}"    #子节点的路径赋值
        #child_node.messages = copy.deepcopy(self.messages)      #传递messages #给什么时候的messages待验证#？
        self.copy_messages(child_node)  #传递messages--只保留两层

        self.children.append(child_node)

    #修改节点的执行状态--return bchange
    def update_work_status(self,work_status):
        if self._work_status == work_status:
            return False
        self._work_status = work_status
        return True

    def get_work_status(self):
        #加锁有没有意义--待定
        return self._work_status

    #-------后期扩充逻辑，目前wokr_status的修改交给上层类对象-------
    def add_instr(self,instr):
        self._instr_queue.append(instr)

    def get_instr(self):
        return self._instr_queue.pop(0) if self._instr_queue else None

    def get_instr_user(self):
        return self._instr_queue

    def del_instr(self,instr):
        if instr in self._instr_queue:
            self._instr_queue.remove(instr)
            #指令删除后要判断是否清空指令了
            if not self._instr_queue:
                self._work_status = 0  #状态调整为没有带执行指令
            return True,""
        else:
            return False,"该指令不在队列中！"



    def add_res(self,str_res): #结构化结果字串
        self._llm_quere.append(str_res)

    def get_res(self):
        return self._llm_quere.pop(0) if self._llm_quere else None

    def get_res_user(self):
        return self._llm_quere

    def get_work_status(self):
        return self._work_status

    def updatemsg(self,newtype,newcontent,index):
        newmsg = {"llm_type":int(newtype),"result":newcontent}
        if self._llm_quere:
            self._llm_quere[0] = newmsg
        else:#新增消息
            self._llm_quere.append(newmsg)
            #更新节点状态
            self._work_status = 3   #待提交
        return True,""

    def is_instr_empty(self):
        if self._instr_queue:
            return False
        return True

    def is_llm_empty(self):
        if self._llm_quere:
            return False
        return True

    def __repr__(self):
        return f"TreeNode({self.name}, {self.status}, {self.vul_type})"


if __name__ == "__main__":
    pass