verify/backend/internal/middleware/crypto.go

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
			}

			// 空请求体（如 GET 请求），跳过解析，直接放行
			if len(body) > 0 {
				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.EncryptKey)
						}
					}
				}

				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)
}