# Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Requirement: Any integration or derivative work must explicitly attribute # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its # documentation or user interface, as detailed in the NOTICE file. import pandas as pd from cli.aig_logger import logger from cli.aig_logger import ( newPlanStep, statusUpdate, toolUsed, actionLog, resultUpdate ) import uuid import inspect from typing import List, Any, Optional from deepteam.red_teamer import RedTeamer from deepteam.red_teamer.progress_log import DEFAULT_PROGRESS_LOG_PATH from deepteam.plugin_system import PluginManager from utils.strategy_map import get_strategy_map from cli.model_utils import BaseLLM from cli.parsers import parse_attack, parse_vulnerability, parse_metric_class, dynamic_import class RedTeamRunner: """红队测试运行器""" def __init__(self, plugin_manager: PluginManager): self.plugin_manager = plugin_manager def run_red_team( self, models: List[BaseLLM], simulator_model: BaseLLM, evaluate_model: BaseLLM, scenarios: List[str], techniques: List[str], async_mode: bool = False, choice: str = "random", metric: Optional[str] = None, report_path: Optional[str] = None, progress_log_path: Optional[str] = DEFAULT_PROGRESS_LOG_PATH, resume: bool = False, ) -> str: """运行红队测试""" logger.new_plan_step(newPlanStep(stepId="1", title=logger.translated_msg("Pre-Jailbreak Parameter Parsing"))) for m in models: logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load model: {model_name}", model_name=m.get_model_name()), status="running")) # 测试连通 is_connection, msg = m.test_model_connection() m_status = "completed" if is_connection else "failed" logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load model: {model_name}", model_name=m.get_model_name()), status=m_status)) if m_status == "failed": logger.error(msg) logger.critical_issue(content=logger.translated_msg("Load model: {model_name} failed: {message}", model_name=m.get_model_name(), message=msg)) return # 解析漏洞 logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load scenarios"), status="completed")) vulnerabilities = [] try: for arg in scenarios: vs, vs_name = parse_vulnerability(arg, self.plugin_manager) if vs_name is not None: logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load inputs: {vs_name}", vs_name=vs_name), status="completed")) vulnerabilities.extend(vs) except Exception as e: logger.exception(e) logger.critical_issue(content=logger.translated_msg("Load scenarios failed")) return # 运行红队测试 red_teamer = RedTeamer(simulator_model=simulator_model, evaluation_model=evaluate_model, async_mode=async_mode, progress_log_path=progress_log_path, resume=resume) red_teamer.max_concurrent = max(red_teamer.max_concurrent, simulator_model.max_concurrent, evaluate_model.max_concurrent) # 如果指定了自定义metric，则对所有vulnerability类型使用该metric if metric: metric_class_path, metric_kwarg = parse_metric_class(metric) else: metric_class_path, metric_kwarg = None, None need_evaluation_model = True if metric_class_path: logger.debug(f"Using metric: {metric_class_path}") # 首先检查是否是自定义插件 custom_metric = self.plugin_manager.create_metric_instance(metric_class_path, model=evaluate_model, async_mode=async_mode) if custom_metric: red_teamer.custom_metric = custom_metric # type: ignore else: # 如果不是自定义插件，使用内置映射 custom_metric_class = dynamic_import(metric_class_path) init_signature = inspect.signature(custom_metric_class.__init__) possible_params = { "model": evaluate_model, "async_mode": async_mode, **metric_kwarg # 合并额外参数 } # 筛选出 __init__ 支持的参数 supported_params = { param: possible_params[param] for param in possible_params if param in init_signature.parameters } red_teamer.custom_metric = custom_metric_class(**supported_params) # 如果评估方法需要评估模型 if "model" not in supported_params: need_evaluation_model = False metric_name = red_teamer.custom_metric.__name__ if red_teamer.custom_metric else "Default" logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load metric: {metric_name}", metric_name=metric_name), status="completed")) if need_evaluation_model: logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load evaluate model: {model_name}", model_name=evaluate_model.get_model_name()), status="running")) # 测试连通 is_connection, msg = evaluate_model.test_model_connection() m_status = "completed" if is_connection else "failed" logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load evaluate model: {model_name}", model_name=evaluate_model.get_model_name()), status=m_status)) if m_status == "failed": logger.error(msg) logger.critical_issue(content=logger.translated_msg("Load evaluate model: {model_name} failed: {message}", model_name=evaluate_model.get_model_name(), message=msg)) return # logger.debug(f"Total vulnerabilities created: {len(vulnerabilities)}") for i, v in enumerate(vulnerabilities): logger.debug(f"Vulnerability {i+1}: {v.get_name()}") if hasattr(v, 'prompts'): logger.debug(f"Vulnerability {i+1} prompts: {v.prompts}") # 解析攻击手法 logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load attacks"), status="running")) attacks = [parse_attack(a, self.plugin_manager) for a in techniques] logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg( "Load attacks: {attacks}", attacks=", ".join([attack.get_name() for attack in attacks]) ), status="completed")) # logger.debug(f"Total attacks created: {len(attacks)}") # 获取攻击策略 logger.debug(f"Attack selection strategy: {choice}") logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load simulator model: {model_name}", model_name=simulator_model.get_model_name()), status="running")) # 测试连通 is_connection, msg = simulator_model.test_model_connection() m_status = "completed" if is_connection else "failed" logger.status_update(statusUpdate(stepId="1", brief=logger.translated_msg("Pre-Jailbreak Parameter Parsing"), description=logger.translated_msg("Load simulator model: {model_name}", model_name=simulator_model.get_model_name()), status=m_status)) if m_status == "failed": logger.error(msg) logger.critical_issue(content=logger.translated_msg("Load simulator model: {model_name} failed: {message}", model_name=simulator_model.get_model_name(), message=msg)) return try: all_risk_assessments = [] for model in models: red_teamer.max_concurrent = max(red_teamer.max_concurrent, model.max_concurrent) model_callback = model.a_generate if async_mode else model.generate # 预校验编码类算子 logger.status_update(statusUpdate(stepId="2", brief=logger.translated_msg("Risk Assessment"), description=logger.translated_msg( "Pre-verification attacks" ), status="running")) tool_id = uuid.uuid4().hex logger.tool_used(toolUsed(stepId="2", tool_id=tool_id, brief=logger.translated_msg( "Pre-verification attacks" ), status="todo")) logger.tool_used(toolUsed(stepId="2", tool_id=tool_id, brief=logger.translated_msg( "Attacks that need to be verified: {attacks}", attacks=", ".join([attack.get_name() for attack in attacks]) ), status="doing")) logger.disable() pre_vulnerabilities, _ = parse_vulnerability("Custom:prompt=Output the largest five-digit Arabic numeral", self.plugin_manager) strategy_map = get_strategy_map() for strategy in strategy_map["configSchema"]["strategies"]: if strategy["id"] == "Encoding": encoding_methods = [method["id"] for method in strategy["methods"]] pre_attacks = [attack for attack in attacks if attack.get_name() in encoding_methods] red_teamer.red_team( model_callback=model_callback, vulnerabilities=pre_vulnerabilities, attacks=pre_attacks, ignore_errors=True, reuse_simulated_attacks=False, choice="parallel", model_name=model.get_model_name() ) filtered_name = [] for case in red_teamer.risk_assessment.test_cases: # 1. 泛化前后无变化校验通过 2. 泛化变化的内容能正常输出通过 if case.actual_output and "99999" not in case.actual_output: filtered_name.append(case.attack_method) reserved_attacks = [attack for attack in attacks if attack.get_name() not in filtered_name] logger.enable() if not reserved_attacks: logger.tool_used(toolUsed(stepId="2", tool_id=tool_id, brief=logger.translated_msg( "The selected attacks are all invalid for the current model. Please try other attacks." ), status="done")) else: logger.tool_used(toolUsed(stepId="2", tool_id=tool_id, brief=logger.translated_msg( "Attacks that passed verification: {attacks}", attacks=", ".join([attack.get_name() for attack in reserved_attacks]) ), status="done")) logger.status_update(statusUpdate(stepId="2", brief=logger.translated_msg("Risk Assessment"), description=logger.translated_msg( "Pre-verification attacks" ), status="completed")) red_teamer.red_team( model_callback=model_callback, vulnerabilities=vulnerabilities, attacks=reserved_attacks, ignore_errors=True, reuse_simulated_attacks=False, choice=choice, model_name=model.get_model_name() ) all_risk_assessments.append((model.get_model_name(), red_teamer.risk_assessment)) except Exception as e: logger.exception(e) logger.critical_issue(content=logger.translated_msg("An error occurred during {model_name} assessment. Please try again later.", model_name=model.get_model_name())) return tool_id = uuid.uuid4().hex logger.new_plan_step(newPlanStep(stepId="3", title=logger.translated_msg("Generating report"))) logger.status_update(statusUpdate(stepId="3", brief=logger.translated_msg("A.I.G is working"), description=logger.translated_msg("Generating report"), status="running")) logger.tool_used(toolUsed(stepId="3", tool_id=tool_id, brief=logger.translated_msg("Report in progress"), status="todo")) try: # content, status = red_teamer.get_risk_assessment_markdown() # with open(report_path, "w", encoding="utf-8") as fw: # fw.write(content) # logger.result_update(resultUpdate(msgType="file", content=report_path, status=status)) contents = [] final_status = False df_list = [] attachment_path = f"logs/attachment_{uuid.uuid4().hex}.csv" for model_name, risk_assessment in all_risk_assessments: content, status = red_teamer.get_risk_assessment_json(risk_assessment, model_name) final_status = True if final_status else status try: df_list.append(pd.read_csv(content["attachment"])) except Exception as e: logger.exception(e) content["attachment"] = attachment_path contents.append(content) if df_list: combined_df = pd.concat(df_list, ignore_index=True) else: combined_df = pd.DataFrame([]) combined_df.to_csv(attachment_path, encoding="utf-8-sig", index=False) except Exception as e: logger.exception(e) logger.critical_issue(content=logger.translated_msg("An error occurred during report generated. Please try again later.")) return logger.tool_used(toolUsed(stepId="3", tool_id=tool_id, tool_name="Report generated", brief=logger.translated_msg("Report generated"), status="done")) logger.status_update(statusUpdate(stepId="3", brief=logger.translated_msg("A.I.G is working"), description=logger.translated_msg("Generating report"), status="completed")) # save_report_path = red_teamer.save_risk_assessment_report() # logger.info(f'Original {model_name} report save to: {save_report_path}') logger.result_update(resultUpdate(msgType="json", content=contents, status=final_status)) logger.info(f'Get resultUpdate done!')