feat:protect content

#145

api/waf_gpt.go

@@ -4,6 +4,7 @@ import (
 	"SamWaf/global"
 	"SamWaf/model"
 	"SamWaf/model/request"
+	"SamWaf/wafsec"
 	"bufio"
 	"bytes"
 	"encoding/json"
@@ -53,10 +54,10 @@ func SendDeltaMessage(messageChan chan<- string, content string, role ...string)
 	if len(role) > 0 {
 		r = role[0]
 	}
-
+	encryptStr, _ := wafsec.AesEncrypt([]byte(content), global.GWAF_COMMUNICATION_KEY)
 	// 创建消息结构
 	msg := Delta{
-		Content: content,
+		Content: encryptStr,
 		Role:    &r,
 	}
 

middleware/sec_api_middleware.go

@@ -35,6 +35,9 @@ func SecApi() gin.HandlerFunc {
 			//fmt.Println("Raw body解密:", string(deBytes))
 			// Store the modified body back in the request
 			c.Request.Body = ioutil.NopCloser(bytes.NewBuffer(decryptBytes))
+		} else if c.Request.Header.Get("accept") == ", text/event-stream" {
+			decryptBytes, _ := wafsec.AesDecrypt(string(bodyBytes), global.GWAF_COMMUNICATION_KEY)
+			c.Request.Body = ioutil.NopCloser(bytes.NewBuffer(decryptBytes))
 		}
 		c.Next()
 	}

router/waf_gpt.go

@@ -2,6 +2,7 @@ package router
 
 import (
 	"SamWaf/api"
+	"SamWaf/middleware"
 	"github.com/gin-gonic/gin"
 )
 
@@ -11,5 +12,6 @@ type WafGPTRouter struct {
 func (receiver *WafGPTRouter) InitGPTRouter(group *gin.RouterGroup) {
 	api := api.APIGroupAPP.WafGPTApi
 	router := group.Group("")
+	router.Use(middleware.StreamMiddleware())
 	router.POST("/samwaf/gpt/chat", api.ChatApi)
 }

wafmangeweb/localserver.go

@@ -63,10 +63,7 @@ func (web *WafWebManager) initRouter(r *gin.Engine) {
 		router.ApiGroupApp.InitWafHttpAuthBaseRouter(RouterGroup)
 		router.ApiGroupApp.InitWafTaskRouter(RouterGroup)
 		router.ApiGroupApp.InitWafBlockingPageRouter(RouterGroup)
-
-		gptRouterGroup := r.Group("")
-		gptRouterGroup.Use(middleware.StreamMiddleware())
-		router.ApiGroupApp.InitGPTRouter(gptRouterGroup)
+		router.ApiGroupApp.InitGPTRouter(RouterGroup)
 	}
 
 	if global.GWAF_RELEASE == "true" {