Initial commit: 网络验证平台

backend/internal/middleware/crypto.go

@@ -0,0 +1,156 @@
+package middleware
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"strings"
+	"verification-platform-backend/internal/database"
+	"verification-platform-backend/internal/model"
+	"verification-platform-backend/pkg/crypto"
+
+	"github.com/gin-gonic/gin"
+)
+
+type CryptoMiddleware struct {
+	cryptoManager *crypto.CryptoManager
+	shouldEncrypt bool
+}
+
+func NewCryptoMiddleware(encryptType crypto.EncryptType, secretKey string) *CryptoMiddleware {
+	shouldEncrypt := encryptType != crypto.EncryptTypeNone
+	return &CryptoMiddleware{
+		cryptoManager: crypto.NewCryptoManager(encryptType, secretKey),
+		shouldEncrypt: shouldEncrypt,
+	}
+}
+
+type EncryptedRequest struct {
+	Data string `json:"data" binding:"required"`
+}
+
+type EncryptedResponse struct {
+	Data string `json:"data"`
+}
+
+func (cm *CryptoMiddleware) ProcessRequest() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		contentType := c.GetHeader("Content-Type")
+
+		if strings.Contains(contentType, "application/json") {
+			body, err := io.ReadAll(c.Request.Body)
+			if err != nil {
+				c.JSON(400, gin.H{"code": 400, "message": "读取请求体失败"})
+				c.Abort()
+				return
+			}
+
+			var req interface{}
+			if err := json.Unmarshal(body, &req); err != nil {
+				c.JSON(400, gin.H{"code": 400, "message": "解析请求体失败"})
+				c.Abort()
+				return
+			}
+
+			reqMap, ok := req.(map[string]interface{})
+			if ok {
+				if encryptedData, exists := reqMap["data"]; exists {
+					if dataStr, ok := encryptedData.(string); ok {
+						decrypted, err := cm.cryptoManager.Decrypt(dataStr)
+						if err != nil {
+							c.JSON(400, gin.H{"code": 400, "message": "解密失败: " + err.Error()})
+							c.Abort()
+							return
+						}
+
+						var decryptedData interface{}
+						if err := json.Unmarshal([]byte(decrypted), &decryptedData); err != nil {
+							c.JSON(400, gin.H{"code": 400, "message": "解析解密数据失败"})
+							c.Abort()
+							return
+						}
+
+						c.Set("decrypted_data", decryptedData)
+						c.Set("is_encrypted", true)
+					}
+				}
+			}
+
+			c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
+		}
+
+		c.Set("should_encrypt_response", cm.shouldEncrypt)
+		c.Set("crypto_manager", cm.cryptoManager)
+		c.Next()
+	}
+}
+
+func (cm *CryptoMiddleware) ProcessResponse() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		blw := &bodyLogWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer}
+		c.Writer = blw
+
+		c.Next()
+
+		shouldEncrypt := c.GetBool("should_encrypt_response")
+		fmt.Printf("[ProcessResponse] ShouldEncrypt: %v\n", shouldEncrypt)
+
+		if shouldEncrypt {
+			var response map[string]interface{}
+			if err := json.Unmarshal(blw.body.Bytes(), &response); err == nil {
+				responseJSON, _ := json.Marshal(response)
+				fmt.Printf("[ProcessResponse] Response to encrypt: %s\n", string(responseJSON))
+
+				var cryptoManager *crypto.CryptoManager
+				if cm, ok := c.Get("crypto_manager"); ok {
+					cryptoManager = cm.(*crypto.CryptoManager)
+				} else {
+					appKey := c.Param("appKey")
+					if appKey != "" {
+						var app model.Application
+						if err := database.DB.Where("app_key = ?", appKey).First(&app).Error; err == nil {
+							var encryptType crypto.EncryptType
+							switch app.EncryptType {
+							case "aes":
+								encryptType = crypto.EncryptTypeAES
+							case "rc4":
+								encryptType = crypto.EncryptTypeRC4
+							default:
+								encryptType = crypto.EncryptTypeNone
+							}
+							cryptoManager = crypto.NewCryptoManager(encryptType, app.SecretKey)
+						}
+					}
+				}
+
+				if cryptoManager != nil {
+					encrypted, err := cryptoManager.Encrypt(string(responseJSON))
+					if err == nil {
+						encryptedResponse := EncryptedResponse{Data: encrypted}
+						encryptedJSON, _ := json.Marshal(encryptedResponse)
+						c.Writer.Header().Set("Content-Type", "application/json")
+						c.Writer.Write(encryptedJSON)
+						fmt.Printf("[ProcessResponse] Encrypted response sent\n")
+						return
+					} else {
+						fmt.Printf("[ProcessResponse] Encryption error: %v\n", err)
+					}
+				} else {
+					fmt.Printf("[ProcessResponse] CryptoManager is nil\n")
+				}
+			} else {
+				fmt.Printf("[ProcessResponse] JSON unmarshal error: %v\n", err)
+			}
+		}
+	}
+}
+
+type bodyLogWriter struct {
+	gin.ResponseWriter
+	body *bytes.Buffer
+}
+
+func (w *bodyLogWriter) Write(b []byte) (int, error) {
+	return w.body.Write(b)
+}