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  package cmd
 20  
 21  import (
 22  	"github.com/Tencent/AI-Infra-Guard/common/runner"
 23  	"github.com/Tencent/AI-Infra-Guard/internal/gologger"
 24  	"github.com/Tencent/AI-Infra-Guard/internal/options"
 25  	"github.com/spf13/cobra"
 26  )
 27  
 28  // 为scanCmd定义标志变量
 29  var (
 30  	scanTargets         []string
 31  	scanTargetFile      string
 32  	scanOutputFile      string
 33  	scanProxyURL        string
 34  	scanTimeOut         int
 35  	scanRateLimit       int
 36  	scanFpTemplates     string
 37  	scanAdvTemplates    string
 38  	scanListVulTemplate bool
 39  	scanCheckVulTargets bool
 40  	scanLocalScan       bool
 41  	scanAIAnalysis      bool
 42  	scanAIHunyuanToken  string
 43  	scanAIDeepSeekToken string
 44  	scanHeaders         []string
 45  	scanLanguage        string
 46  )
 47  
 48  // scanCmd 表示scan子命令
 49  var scanCmd = &cobra.Command{
 50  	Use:   "scan",
 51  	Short: "执行原始扫描功能",
 52  	Long:  `执行原始的扫描功能，扫描指定目标的安全漏洞。`,
 53  	Run: func(cmd *cobra.Command, args []string) {
 54  		// 创建Options对象
 55  		scanOptions := &options.Options{
 56  			Target:          scanTargets,
 57  			TargetFile:      scanTargetFile,
 58  			Output:          scanOutputFile,
 59  			ProxyURL:        scanProxyURL,
 60  			TimeOut:         scanTimeOut,
 61  			RateLimit:       scanRateLimit,
 62  			FPTemplates:     scanFpTemplates,
 63  			AdvTemplates:    scanAdvTemplates,
 64  			ListVulTemplate: scanListVulTemplate,
 65  			CheckVulTargets: scanCheckVulTargets,
 66  			LocalScan:       scanLocalScan,
 67  			Headers:         scanHeaders,
 68  			Language:        scanLanguage,
 69  		}
 70  
 71  		// 设置日志级别
 72  		gologger.StdLogger.Logrus().SetLevel(gologger.InfoLevel)
 73  
 74  		r, err := runner.New(scanOptions)
 75  		if err != nil {
 76  			gologger.Fatalf("Could not create runner: %s\n", err)
 77  		}
 78  		defer r.Close()
 79  
 80  		if scanOptions.CheckVulTargets {
 81  			return
 82  		}
 83  		if scanOptions.ListVulTemplate {
 84  			r.ShowFpAndVulList(scanOptions.ListVulTemplate)
 85  			return
 86  		}
 87  
 88  		r.RunEnumeration()
 89  	},
 90  }
 91  
 92  func init() {
 93  	rootCmd.AddCommand(scanCmd)
 94  
 95  	// 设置scan子命令的标志
 96  	scanCmd.Flags().StringArrayVarP(&scanTargets, "target", "t", []string{}, "目标URL，可以指定多个目标，例如: --target xx.com --target aa.com")
 97  	scanCmd.Flags().StringVarP(&scanTargetFile, "file", "f", "", "包含目标URL的文件路径")
 98  	scanCmd.Flags().StringVarP(&scanOutputFile, "output", "o", "", "输出文件路径")
 99  	scanCmd.Flags().IntVar(&scanTimeOut, "timeout", 5, "请求超时时间(秒)")
100  	scanCmd.Flags().StringVar(&scanProxyURL, "proxy-url", "", "代理服务器URL")
101  	scanCmd.Flags().StringArrayVar(&scanHeaders, "header", []string{}, "HTTP请求头，可以指定多个，例如: --header \"key:value\" --header \"key:value\"")
102  	scanCmd.Flags().IntVar(&scanRateLimit, "limit", 200, "每秒最大请求数")
103  	scanCmd.Flags().StringVar(&scanFpTemplates, "fps", "data/fingerprints", "指纹模板文件或目录")
104  	scanCmd.Flags().StringVar(&scanAdvTemplates, "vul", "data/vuln", "漏洞数据库目录")
105  	scanCmd.Flags().BoolVar(&scanListVulTemplate, "list-vul", false, "列出漏洞模板")
106  	scanCmd.Flags().BoolVar(&scanCheckVulTargets, "check-vul", false, "验证漏洞模板")
107  	scanCmd.Flags().BoolVar(&scanLocalScan, "localscan", false, "一键本地扫描")
108  	scanCmd.Flags().BoolVar(&scanAIAnalysis, "ai", false, "启用AI分析")
109  	scanCmd.Flags().StringVar(&scanAIHunyuanToken, "hunyuan-token", "", "混元API令牌")
110  	scanCmd.Flags().StringVar(&scanAIDeepSeekToken, "deepseek-token", "", "DeepSeek API令牌")
111  	scanCmd.Flags().StringVar(&scanLanguage, "lang", "zh", "响应语言 zh/en")
112  }