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 websocket 实现WebSocket服务器功能
20  package websocket
21  
22  import (
23  	"net/http"
24  
25  	"github.com/gin-gonic/gin"
26  )
27  
28  const (
29  	// WSMsgTypeLog 日志消息类型
30  	WSMsgTypeLog = "log"
31  	// WSMsgTypeScanResult 扫描结果消息类型
32  	WSMsgTypeScanResult = "result"
33  	// WSMsgTypeProcessInfo 进度消息类型
34  	WSMsgTypeProcessInfo = "processing"
35  	// WSMsgTypeReportInfo 报告消息类型
36  	WSMsgTypeReportInfo = "report"
37  	// WSMsgTypeScanRet 扫描状态返回
38  	WSMsgTypeScanRet = "scan_ret"
39  )
40  
41  const (
42  	WSLogLevelInfo  = "info"
43  	WSLogLevelDebug = "debug"
44  	WSLogLevelError = "error"
45  )
46  
47  // ScanRequest 扫描请求结构
48  type ScanRequest struct {
49  	ScanType string            `json:"scan_type"`
50  	Target   []string          `json:"target,omitempty"`
51  	Headers  map[string]string `json:"headers,omitempty"`
52  	Lang     string            `json:"lang,omitempty"`
53  }
54  
55  // Response 基础响应结构
56  type Response struct {
57  	Status  int         `json:"status"`
58  	Message string      `json:"message"`
59  	Data    interface{} `json:"data"`
60  }
61  
62  // WSMessage WebSocket消息结构
63  type WSMessage struct {
64  	Type    string      `json:"type"`
65  	Content interface{} `json:"content"`
66  }
67  
68  // ReportInfo 报告信息结构
69  type ReportInfo struct {
70  	SecScore   int `json:"sec_score"`
71  	HighRisk   int `json:"high_risk"`
72  	MiddleRisk int `json:"middle_risk"`
73  	LowRisk    int `json:"low_risk"`
74  }
75  
76  // ScanRet 扫描状态返回
77  type ScanRet struct {
78  	Status int    `json:"status"`
79  	Msg    string `json:"msg"`
80  }
81  
82  // Log 日志信息结构
83  type Log struct {
84  	Message string `json:"message"`
85  	Level   string `json:"level"`
86  }
87  
88  func SuccessResponse(c *gin.Context, message interface{}) {
89  	c.JSON(http.StatusOK, gin.H{"status": 0, "message": message})
90  }
91  
92  func ErrorResponse(c *gin.Context, message interface{}) {
93  	c.JSON(http.StatusOK, gin.H{"status": 1, "message": message})
94  }
95  
96  func ErrorResponseWithStatus(c *gin.Context, message interface{}, status int) {
97  	c.JSON(http.StatusOK, gin.H{"status": status, "message": message})
98  }