#对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()