Initial commit: 网络验证平台

backend/internal/router/app/response_encryption.go

@@ -0,0 +1,73 @@
+package app
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"verification-platform-backend/pkg/crypto"
+
+	"github.com/gin-gonic/gin"
+)
+
+type bodyLogWriter struct {
+	gin.ResponseWriter
+	body *bytes.Buffer
+}
+
+func (w *bodyLogWriter) Write(b []byte) (int, error) {
+	return w.body.Write(b)
+}
+
+func (w *bodyLogWriter) WriteHeader(statusCode int) {
+	w.ResponseWriter.WriteHeader(statusCode)
+}
+
+type EncryptedResponse struct {
+	Data string `json:"data"`
+}
+
+func ResponseEncryption() 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("[ResponseEncryption] ShouldEncrypt: %v\n", shouldEncrypt)
+
+		blw.ResponseWriter.Header().Set("Content-Type", "application/json; charset=utf-8")
+
+		if shouldEncrypt {
+			var response map[string]interface{}
+			if err := json.Unmarshal(blw.body.Bytes(), &response); err == nil {
+				responseJSON, _ := json.Marshal(response)
+				fmt.Printf("[ResponseEncryption] Response to encrypt: %s\n", string(responseJSON))
+
+				var cryptoManager *crypto.CryptoManager
+				if cm, ok := c.Get("crypto_manager"); ok {
+					cryptoManager = cm.(*crypto.CryptoManager)
+				}
+
+				if cryptoManager != nil {
+					encrypted, err := cryptoManager.Encrypt(string(responseJSON))
+					if err == nil {
+						encryptedResponse := EncryptedResponse{Data: encrypted}
+						encryptedJSON, _ := json.Marshal(encryptedResponse)
+						blw.ResponseWriter.Write(encryptedJSON)
+						fmt.Printf("[ResponseEncryption] Encrypted response sent\n")
+						return
+					} else {
+						fmt.Printf("[ResponseEncryption] Encryption error: %v\n", err)
+					}
+				} else {
+					fmt.Printf("[ResponseEncryption] CryptoManager is nil\n")
+				}
+			} else {
+				fmt.Printf("[ResponseEncryption] JSON unmarshal error: %v\n", err)
+			}
+		}
+
+		blw.ResponseWriter.Write(blw.body.Bytes())
+	}
+}