zfsafe/mycode/CommandVerify.py

#对llm返回的指令进行校验
import re
class CommandVerify:
    def __init__(self):
        pass

    #验证节点指令的结构完整性--主要是判断JSON元素是否完整
    def verify_node_cmds(self,node_cmds):
        '''
- 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1，节点2\"}；
- 未生成指令节点列表：{\"action\": \"no_instruction\", \"nodes\": \"节点1,节点2\"}；
- 漏洞验证成功：{\"action\": \"find_vul\", \"node\": \"节点\",\"vulnerability\": {\"name\":\"漏洞名称\",\"risk\":\"风险等级(低危/中危/高危)\",\"info\":\"补充信息（没有可为空）\"}}；
- 完成测试：{\"action\": \"end_work\", \"node\": \"节点\"}；
        '''
        strerror = ""
        for node_json in node_cmds:
            if "action" not in node_json:
                self.logger.error(f"缺少action节点：{node_json}")
                strerror = {"节点指令错误":f"{node_json}缺少action节点，不符合格式要求!"}
                break

            action = node_json["action"]
            if action == "add_node":
                if "parent" not in node_json or "nodes" not in node_json:
                    strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"}
                    break
            elif action == "end_work":
                if "node" not in node_json:
                    strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"}
                    break
            elif action =="no_instruction":
                if "nodes" not in node_json:
                    strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"}
                    break
            elif action =="find_vul":
                if "node" not in node_json or "vulnerability" not in node_json:
                    strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"}
                    break
            else:
                strerror = {"节点指令错误": f"{node_json}不可识别的action值！"}
                break
        if not strerror:
            return True,strerror
        return False,strerror

    # 验证节点数据的合规性
    def verify_node_data(self,node_cmds):
        add_nodes = []
        no_instr_nodes = []
        for node_cmd in node_cmds:
            do_type = node_cmd["action"]
            if do_type == "add_node":
                nodes = node_cmd["nodes"].split(",")
                add_nodes.extend(nodes)
            elif do_type == "no_instruction":
                nodes = node_cmd["nodes"].split(",")
                no_instr_nodes.extend(nodes)
            else:# 其他类型暂时不验证
                pass
        #核对指令是否有缺失
        had_inst_nodes = self._difference_a_simple(add_nodes,no_instr_nodes) #在新增节点，但不在没有指令列表，就是应该要有指令的节点数据
        no_add_nodes = self._difference_a_simple(no_instr_nodes,add_nodes)  #在未新增指令的节点，但不在新增节点，就是没有add的节点，需要新增
        return had_inst_nodes,no_add_nodes

    #--------------辅助函数-----------------
    def get_path_from_command(self,command):
        pass

    def _difference_a(self,list_a: list, list_b: list) -> list:
        """获取 list_a 中存在但 list_b 中不存在的元素（去重版）"""
        set_b = set(list_b)
        return [x for x in list_a if x not in set_b]

    def _difference_b(self,list_a: list, list_b: list) -> list:
        """获取 list_b 中存在但 list_a 中不存在的元素（去重版）"""
        set_a = set(list_a)
        return [x for x in list_b if x not in set_a]

    def _difference_a_keep_duplicates(self,list_a: list, list_b: list) -> list:
        """获取 list_a 中存在但 list_b 中不存在的元素（保留所有重复项和顺序）"""
        set_b = set(list_b)
        return [x for x in list_a if x not in set_b]

    def _difference_b_keep_duplicates(self,list_a: list, list_b: list) -> list:
        """获取 list_b 中存在但 list_a 中不存在的元素（保留所有重复项和顺序）"""
        set_a = set(list_a)
        return [x for x in list_b if x not in set_a]

    def _difference_a_simple(self,list_a: list, list_b: list) -> list:
        """集合差集：list_a - list_b"""
        return list(set(list_a) - set(list_b))

    def _difference_b_simple(self,list_a: list, list_b: list) -> list:
        """集合差集：list_b - list_a"""
        return list(set(list_b) - set(list_a))

g_CV = CommandVerify()