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  	"strings"
23  
24  	"github.com/Tencent/AI-Infra-Guard/common/websocket"
25  	"github.com/Tencent/AI-Infra-Guard/internal/gologger"
26  	"github.com/Tencent/AI-Infra-Guard/internal/options"
27  	"github.com/spf13/cobra"
28  )
29  
30  // 为webserverCmd定义标志变量
31  var (
32  	webServerAddr string
33  )
34  
35  // webserverCmd 表示webserver子命令
36  var webserverCmd = &cobra.Command{
37  	Use:   "webserver",
38  	Short: "启动Web服务器",
39  	Long:  `启动Web服务器功能，提供Web界面进行扫描。`,
40  	Run: func(cmd *cobra.Command, args []string) {
41  		if !strings.Contains(webServerAddr, "127.0.0.1") {
42  			gologger.Infoln("请注意，Web服务器监听地址为本地IP,外部用户可访问，可能会导致安全风险，请确保在安全的网络环境下运行。")
43  		}
44  		// 创建Options对象
45  		webOptions := &options.Options{
46  			TimeOut:       10,
47  			RateLimit:     200,
48  			FPTemplates:   "data/fingerprints",
49  			AdvTemplates:  "data/vuln",
50  			WebServer:     true,
51  			WebServerAddr: webServerAddr,
52  		}
53  		// 设置日志级别
54  		websocket.RunWebServer(webOptions)
55  	},
56  }
57  
58  func init() {
59  	rootCmd.AddCommand(webserverCmd)
60  
61  	// 设置webserver子命令的标志
62  	webserverCmd.Flags().StringVar(&webServerAddr, "server", "127.0.0.1:8088", "WebSocket服务器地址")
63  }