# 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 time from pathlib import Path import argparse from cli.aig_logger import logger from cli.aig_logger import ( newPlanStep, statusUpdate, toolUsed, actionLog, resultUpdate ) from deepteam.plugin_system import PluginManager from cli.models import create_model from cli.plugin_commands import list_plugins, load_plugins_from_args, show_plugin_template, validate_plugin, auto_discover_plugins from cli.red_team_runner import RedTeamRunner from cli.tool_scanner_cli import handle_tool_scanning # logger config logger.add(f"logs/red_team_{{time:YYYY-MM-DD_HH-mm-ss}}.log", level="DEBUG", enqueue=True, retention="7 days") # 全局插件管理器 plugin_manager = PluginManager() def cleanup_expired_files(log_path: str = "logs", max_age_seconds: int = 86400*30, pattern: str = "attachment_*.csv"): now = time.time() for file in Path(log_path).glob(pattern): if (now - file.stat().st_mtime) > max_age_seconds: file.unlink() def main(): """主函数""" parser = argparse.ArgumentParser(description="Red Team CLI Runner") # 工具扫描相关参数（放在最前面，优先级最高） parser.add_argument("--scan-tools", type=str, choices=["all", "techniques", "metrics", "scenarios"], help="Scan and display all available tools and their parameters") parser.add_argument("--show-tool-params", type=str, help="Show detailed parameter information for a specific tool") # 插件相关参数 parser.add_argument("--plugins", type=str, nargs='+', help="Custom plugin files or directories to load") parser.add_argument("--list-plugins", action="store_true", help="List all available plugins") parser.add_argument("--show-template", type=str, choices=["attack", "metric", "vulnerability"], help="Show plugin template") parser.add_argument("--validate-plugin", type=str, help="Validate a plugin file or directory") parser.add_argument("--auto-discover", action="store_true", help="Auto-discover plugins from default directories") # 红队测试相关参数 parser.add_argument("--base_url", type=str, action='append', help="Base URL for ChatOpenAI") parser.add_argument("--api_key", type=str, nargs=1, action='append', help="API Key for ChatOpenAI") parser.add_argument("--model", type=str, action='append', help="Model name for ChatOpenAI") parser.add_argument("--max_concurrent", type=int, action='append', help="Max concurrent") parser.add_argument("--sim_base_url", type=str, help="Base URL for a simulator model") parser.add_argument("--sim_api_key", type=str, nargs=1, help="API Key for a simulator model") parser.add_argument("--simulator_model", type=str, help="Model name for a simulator model") parser.add_argument("--sim_max_concurrent", type=int, default=10, help="Max concurrent") parser.add_argument("--eval_base_url", type=str, help="Base URL for a evaluate model") parser.add_argument("--eval_api_key", type=str, nargs=1, help="API Key for a evaluate model") parser.add_argument("--evaluate_model", type=str, help="Model name for a evaluate model") parser.add_argument("--eval_max_concurrent", type=int, default=10, help="Max concurrent") parser.add_argument("--scenarios", type=str, nargs='+', help="Scenarios to test") parser.add_argument("--techniques", type=str, nargs='+', help="Techniques to test") parser.add_argument("--async_mode", action='store_true', help="Enable async mode") parser.add_argument("--choice", type=str, default="random", choices=["random", "serial", "parallel"], help="Technique selection strategy: 'random' (default) or 'serial' (nested techniques) or 'parallel'") parser.add_argument("--metric", type=str, help="Metric class name (e.g., 'RandomMetric')") parser.add_argument("--report", type=str, default="logs/report.md", help="Path to save the risk assessment report (default: logs/report.md)") parser.add_argument("--progress-log", type=str, default="logs/red_team_progress.jsonl", help="Path to write finished test cases to while the run is still going, so an interrupted run keeps its results (default: logs/red_team_progress.jsonl, pass an empty value to disable)") parser.add_argument("--resume", action='store_true', help="Reuse the saved attacks checkpoint and skip the test cases a previous run already finished (default: off)") parser.add_argument("--lang", type=str, default="zh_CN", help="Report language") args = parser.parse_args() logger.set_language(lang=args.lang) # 处理工具扫描相关命令（优先级最高） if args.scan_tools or args.show_tool_params: if handle_tool_scanning(args): exit(0) # 处理插件相关命令 if args.show_template: show_plugin_template(args.show_template, plugin_manager) exit(0) if args.validate_plugin: validate_plugin(args.validate_plugin, plugin_manager) exit(0) # 加载插件（在list_plugins之前） if args.auto_discover: auto_discover_plugins(plugin_manager) if args.plugins: load_plugins_from_args(args.plugins, plugin_manager) if args.list_plugins: list_plugins(plugin_manager) exit(0) # 初始化模型 models = [] lengths = list(map(len, (args.base_url, args.api_key, args.model, args.max_concurrent))) if len(set(lengths)) != 1: raise ValueError("base_url, api_key, model, max_concurrent must have same number of parameters") for base_url, api_key, model_name, max_concurrent in zip(args.base_url, args.api_key, args.model, args.max_concurrent): model = create_model(model_name, base_url, api_key[0], max_concurrent) models.append(model) if any(param is None for param in (args.evaluate_model, args.eval_base_url, args.eval_api_key, args.eval_max_concurrent)): evaluate_model = models[0] else: evaluate_model = create_model(args.evaluate_model, args.eval_base_url, args.eval_api_key[0], args.eval_max_concurrent) if any(param is None for param in (args.simulator_model, args.sim_base_url, args.sim_api_key, args.sim_max_concurrent)): simulator_model = evaluate_model else: simulator_model = create_model(args.simulator_model, args.sim_base_url, args.sim_api_key[0], args.sim_max_concurrent) # 创建红队运行器 runner = RedTeamRunner(plugin_manager) # 运行红队测试 runner.run_red_team( models=models, simulator_model=simulator_model, evaluate_model=evaluate_model, scenarios=args.scenarios, techniques=args.techniques, async_mode=args.async_mode, choice=args.choice, metric=args.metric, report_path=args.report, progress_log_path=args.progress_log or None, resume=args.resume, ) if __name__ == "__main__": try: main() # 清理过期文件 cleanup_expired_files() except Exception as e: logger.error(e) logger.critical_issue(content=logger.translated_msg("Something went wrong. Please try again in a few moments."))