Trae AI
144 lines
3.5 KiB
Go
144 lines
3.5 KiB
Go
package app
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"io"
|
|
"verification-platform-backend/internal/database"
|
|
"verification-platform-backend/internal/model"
|
|
"verification-platform-backend/pkg/crypto"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func AppCryptoMiddleware() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
appKey := c.Param("appKey")
|
|
|
|
if appKey == "" {
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
if database.DB == nil {
|
|
c.JSON(500, gin.H{"code": 500, "message": "数据库未初始化"})
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
var app model.Application
|
|
if err := database.DB.Where("app_key = ?", appKey).First(&app).Error; err != nil {
|
|
c.JSON(404, gin.H{"code": 404, "message": "应用不存在"})
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
var encryptType crypto.EncryptType
|
|
switch app.EncryptType {
|
|
case "aes":
|
|
encryptType = crypto.EncryptTypeAES
|
|
case "rc4":
|
|
encryptType = crypto.EncryptTypeRC4
|
|
default:
|
|
encryptType = crypto.EncryptTypeNone
|
|
}
|
|
|
|
shouldEncrypt := encryptType != crypto.EncryptTypeNone
|
|
cm := crypto.NewCryptoManager(encryptType, app.EncryptKey)
|
|
|
|
// 解密请求体
|
|
if shouldEncrypt {
|
|
contentType := c.GetHeader("Content-Type")
|
|
if contentType != "" {
|
|
// 先尝试从 X-Encrypted-Data Header 解密
|
|
encryptedHeader := c.GetHeader("X-Encrypted-Data")
|
|
if encryptedHeader != "" {
|
|
decrypted, err := cm.Decrypt(encryptedHeader)
|
|
if err != nil {
|
|
c.JSON(400, gin.H{"code": 400, "message": "解密失败: " + err.Error()})
|
|
c.Abort()
|
|
return
|
|
}
|
|
c.Request.Body = io.NopCloser(bytes.NewBufferString(decrypted))
|
|
c.Request.ContentLength = int64(len(decrypted))
|
|
} else {
|
|
// 从请求体 JSON 的 data 字段解密
|
|
body, err := io.ReadAll(c.Request.Body)
|
|
if err == nil && len(body) > 0 {
|
|
var req map[string]interface{}
|
|
if json.Unmarshal(body, &req) == nil {
|
|
if dataStr, ok := req["data"].(string); ok {
|
|
decrypted, err := cm.Decrypt(dataStr)
|
|
if err != nil {
|
|
c.JSON(400, gin.H{"code": 400, "message": "解密失败: " + err.Error()})
|
|
c.Abort()
|
|
return
|
|
}
|
|
c.Request.Body = io.NopCloser(bytes.NewBufferString(decrypted))
|
|
c.Request.ContentLength = int64(len(decrypted))
|
|
} else {
|
|
c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
|
|
}
|
|
} else {
|
|
c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
c.Set("should_encrypt_response", shouldEncrypt)
|
|
c.Set("crypto_manager", cm)
|
|
c.Set("app", &app)
|
|
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func getAppCryptoManager(appKey string) (*crypto.CryptoManager, error) {
|
|
var app model.Application
|
|
if err := database.DB.Where("app_key = ?", appKey).First(&app).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var encryptType crypto.EncryptType
|
|
switch app.EncryptType {
|
|
case "aes":
|
|
encryptType = crypto.EncryptTypeAES
|
|
case "rc4":
|
|
encryptType = crypto.EncryptTypeRC4
|
|
default:
|
|
encryptType = crypto.EncryptTypeNone
|
|
}
|
|
|
|
return crypto.NewCryptoManager(encryptType, app.EncryptKey), nil
|
|
}
|
|
|
|
func decryptRequest(c *gin.Context, appKey string) ([]byte, error) {
|
|
encryptedData := c.GetHeader("X-Encrypted-Data")
|
|
if encryptedData == "" {
|
|
return nil, nil
|
|
}
|
|
|
|
cryptoManager, err := getAppCryptoManager(appKey)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
decrypted, err := cryptoManager.Decrypt(encryptedData)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return []byte(decrypted), nil
|
|
}
|
|
|
|
func encryptResponse(c *gin.Context, appKey string, data []byte) (string, error) {
|
|
cryptoManager, err := getAppCryptoManager(appKey)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
return cryptoManager.Encrypt(string(data))
|
|
}
|