1  # Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  #
  3  # Licensed under the Apache License, Version 2.0 (the "License");
  4  # you may not use this file except in compliance with the License.
  5  # You may obtain a copy of the License at
  6  #
  7  #     http://www.apache.org/licenses/LICENSE-2.0
  8  #
  9  # Unless required by applicable law or agreed to in writing, software
 10  # distributed under the License is distributed on an "AS IS" BASIS,
 11  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  # See the License for the specific language governing permissions and
 13  # limitations under the License.
 14  #
 15  # Requirement: Any integration or derivative work must explicitly attribute
 16  # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  # documentation or user interface, as detailed in the NOTICE file.
 18  
 19  from typing import List
 20  from deepteam.plugin_system import PluginManager
 21  from .mappings import TECHNIQUE_CLASS_MAP, SCENARIO_CLASS_MAP, METRIC_CLASS_MAP
 22  
 23  
 24  def load_plugins_from_args(plugin_paths: List[str], plugin_manager: PluginManager) -> None:
 25      """从命令行参数加载插件"""
 26      if not plugin_paths:
 27          return
 28      
 29      print("正在加载自定义插件...")
 30      for plugin_path in plugin_paths:
 31          result = plugin_manager.load_plugin(plugin_path)
 32          if result['success']:
 33              print(f"✓ 成功加载插件: {plugin_path}")
 34              if result['warnings']:
 35                  for warning in result['warnings']:
 36                      print(f"  警告: {warning}")
 37          else:
 38              print(f"✗ 加载插件失败: {plugin_path}")
 39              for error in result['errors']:
 40                  print(f"  错误: {error}")
 41  
 42  
 43  def list_plugins(plugin_manager: PluginManager) -> None:
 44      """列出所有可用插件"""
 45      print("\n=== 内置攻击插件 ===")
 46      for name in TECHNIQUE_CLASS_MAP.keys():
 47          print(f"  {name}")
 48      
 49      print("\n=== 内置漏洞场景 ===")
 50      for name in SCENARIO_CLASS_MAP.keys():
 51          print(f"  {name}")
 52      
 53      print("\n=== 内置指标 ===")
 54      for name in METRIC_CLASS_MAP.keys():
 55          print(f"  {name}")
 56      
 57      # 显示自定义插件
 58      custom_plugins = plugin_manager.get_loaded_plugins()
 59      if custom_plugins['attacks']:
 60          print("\n=== 自定义攻击插件 ===")
 61          for name in custom_plugins['attacks']:
 62              print(f"  {name}")
 63      
 64      if custom_plugins['vulnerabilities']:
 65          print("\n=== 自定义漏洞插件 ===")
 66          for name in custom_plugins['vulnerabilities']:
 67              print(f"  {name}")
 68      
 69      if custom_plugins['metrics']:
 70          print("\n=== 自定义指标插件 ===")
 71          for name in custom_plugins['metrics']:
 72              print(f"  {name}")
 73  
 74  
 75  def show_plugin_template(plugin_type: str, plugin_manager: PluginManager) -> None:
 76      """显示插件模板"""
 77      template = plugin_manager.get_plugin_template(plugin_type)
 78      print(f"\n=== {plugin_type.title()} 插件模板 ===")
 79      print(template)
 80  
 81  
 82  def validate_plugin(plugin_path: str, plugin_manager: PluginManager) -> None:
 83      """验证插件"""
 84      result = plugin_manager.validate_plugin(plugin_path)
 85      if result['valid']:
 86          print(f"✓ 插件验证通过: {plugin_path}")
 87          print(f"  类型: {result['plugin_type']}")
 88          print(f"  类名: {result['class_name']}")
 89          if result['warnings']:
 90              for warning in result['warnings']:
 91                  print(f"  警告: {warning}")
 92      else:
 93          print(f"✗ 插件验证失败: {plugin_path}")
 94          for error in result['errors']:
 95              print(f"  错误: {error}")
 96  
 97  
 98  def auto_discover_plugins(plugin_manager: PluginManager) -> None:
 99      """自动发现插件"""
100      print("自动发现插件...")
101      result = plugin_manager.auto_discover_plugins()
102      if result['success']:
103          print(f"✓ 自动发现并加载了 {len(result['loaded_plugins'])} 个插件")
104          for plugin in result['loaded_plugins']:
105              print(f"  - {plugin['class']} ({plugin['type']})")
106      else:
107          print("没有发现任何插件")