/ AIG-PromptSecurity / cli / tool_scanner_cli.py
tool_scanner_cli.py
  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  """
 20  工具扫描CLI模块
 21  处理工具扫描相关的命令行功能
 22  """
 23  
 24  from deepteam.plugin_system.tool_scanner import ToolScanner
 25  
 26  
 27  def handle_tool_scanning(args):
 28      """处理工具扫描相关命令"""
 29      scanner = ToolScanner()
 30      
 31      # 添加插件路径
 32      if args.plugins:
 33          for plugin_path in args.plugins:
 34              scanner.add_plugin_path(plugin_path)
 35      
 36      # 扫描所有工具
 37      tools_info = scanner.scan_all_tools()
 38      
 39      if args.scan_tools:
 40          print("=== 可用工具列表 ===")
 41          
 42          # 根据参数选择要显示的工具类型
 43          tool_types = []
 44          if args.scan_tools == "all":
 45              tool_types = ['attack', 'metric', 'vulnerability']
 46          elif args.scan_tools == "techniques":
 47              tool_types = ['attack']
 48          elif args.scan_tools == "metrics":
 49              tool_types = ['metric']
 50          elif args.scan_tools == "scenarios":
 51              tool_types = ['vulnerability']
 52          
 53          for tool_type in tool_types:
 54              # 显示用户友好的类型名称
 55              display_name = {
 56                  'attack': 'TECHNIQUES (攻击技术)',
 57                  'metric': 'METRICS (评估指标)',
 58                  'vulnerability': 'SCENARIOS (测试场景)'
 59              }.get(tool_type, tool_type.upper())
 60              
 61              print(f"\n## {display_name}:")
 62              tools_of_type = scanner.get_tools_by_type(tool_type)
 63              if tools_of_type:
 64                  for tool_name, tool_info in tools_of_type.items():
 65                      print(f"  - {tool_name}")
 66                      if tool_info['parameters']:
 67                          for param_name, param_info in tool_info['parameters'].items():
 68                              required = "必需" if param_info['required'] else "可选"
 69                              default_str = f" (默认: {param_info['default']})" if param_info['default'] is not None else ""
 70                              print(f"    * {param_name} ({required}){default_str}")
 71                              if param_info['description']:
 72                                  print(f"      {param_info['description']}")
 73              else:
 74                  print("  (无可用工具)")
 75          
 76          # # 显示验证警告
 77          # warnings = scanner.validate_tool_completeness()
 78          # if warnings:
 79          #     print("\n=== 验证警告 ===")
 80          #     for warning in warnings:
 81          #         print(warning)
 82          
 83          return True
 84      
 85      if args.show_tool_params:
 86          tool_name = args.show_tool_params
 87          tool_info = scanner.get_tool_info(tool_name)
 88          
 89          if tool_info:
 90              print(f"=== {tool_name} 详细信息 ===")
 91              print(f"类型: {tool_info['type']}")
 92              print(f"文件: {tool_info['file']}")
 93              if tool_info['description']:
 94                  print(f"描述: {tool_info['description']}")
 95              
 96              if tool_info['parameters']:
 97                  print("\n参数:")
 98                  for param_name, param_info in tool_info['parameters'].items():
 99                      required = "必需" if param_info['required'] else "可选"
100                      default_str = f" (默认值: {param_info['default']})" if param_info['default'] is not None else ""
101                      print(f"  {param_name} ({required}){default_str}")
102                      if param_info['description']:
103                          print(f"    描述: {param_info['description']}")
104              else:
105                  print("\n参数: (无参数)")
106          else:
107              print(f"错误: 找不到工具 '{tool_name}'")
108              print("可用工具:")
109              for tool_name in tools_info.keys():
110                  print(f"  - {tool_name}")
111          
112          return True
113      
114      return False