From 1b1c15995da7dc251c3643d262a3ebc75b250da9 Mon Sep 17 00:00:00 2001
From: TsMask <340112800@qq.com>
Date: Wed, 14 Aug 2024 10:42:13 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E4=B8=AD=E9=97=B4=E4=BB=B6-=E6=8E=A5?=
 =?UTF-8?q?=E5=8F=A3=E5=8A=A0=E8=A7=A3=E5=AF=86AES-CBC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/framework/constants/result/result.go |   9 +-
 src/framework/middleware/crypto_api.go   | 134 +++++++++++++++++++++++
 2 files changed, 141 insertions(+), 2 deletions(-)
 create mode 100644 src/framework/middleware/crypto_api.go

diff --git a/src/framework/constants/result/result.go b/src/framework/constants/result/result.go
index dbd86302..19761ce4 100644
--- a/src/framework/constants/result/result.go
+++ b/src/framework/constants/result/result.go
@@ -8,8 +8,13 @@ const (
 	// 响应-msg错误失败
 	MSG_ERROR = "error"
 
-	// 响应-msg正常成功
-	CODE_SUCCESS = 1
 	// 响应-code正常成功
+	CODE_SUCCESS = 1
+	// 响应-msg正常成功
 	MSG_SUCCESS = "success"
+
+	// 响应-code加密数据
+	CODE_ENCRYPT = 2
+	// 响应-msg加密数据
+	MSG_ENCRYPT = "encrypt"
 )
diff --git a/src/framework/middleware/crypto_api.go b/src/framework/middleware/crypto_api.go
new file mode 100644
index 00000000..6f425cbb
--- /dev/null
+++ b/src/framework/middleware/crypto_api.go
@@ -0,0 +1,134 @@
+package middleware
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+
+	"be.ems/src/framework/config"
+	constResult "be.ems/src/framework/constants/result"
+	"be.ems/src/framework/logger"
+	"be.ems/src/framework/utils/crypto"
+	"be.ems/src/framework/utils/parse"
+	"github.com/gin-gonic/gin"
+)
+
+// CryptoApi 接口加解密
+//
+// 示例参数：middleware.CryptoApi(true, true)
+//
+// 参数表示：对请求解密，对响应加密
+//
+// 请将中间件放在最前置，对请求优先处理
+func CryptoApi(requestDecrypt, responseEncrypt bool) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		// 请求解密时对请求data注入
+		if requestDecrypt {
+			contentDe := ""
+			if c.Request.Method == "GET" {
+				contentDe = c.Query("data")
+			} else if c.ContentType() == gin.MIMEJSON {
+				var body struct {
+					Data string `json:"data" binding:"required"`
+				}
+				if err := c.ShouldBindJSON(&body); err == nil {
+					contentDe = body.Data
+				}
+			}
+			// 是否存在data字段数据
+			if contentDe == "" {
+				c.JSON(400, map[string]any{
+					"code": constResult.CODE_ERROR,
+					"msg":  "decrypt not found field data",
+				})
+				c.Abort() // 停止执行后续的处理函数
+				return
+			}
+			// 解密
+			apiKey := config.Get("aes.apiKey").(string)
+			dataBodyStr, err := crypto.AESDecryptBase64(contentDe, apiKey)
+			if err != nil {
+				logger.Errorf("CryptoApi decrypt err => %v", err)
+				c.JSON(400, map[string]any{
+					"code": constResult.CODE_ERROR,
+					"msg":  "decrypted data could not be parsed",
+				})
+				c.Abort() // 停止执行后续的处理函数
+				return
+			}
+			// 分配回请求体
+			c.Request.Body = io.NopCloser(bytes.NewBuffer([]byte(dataBodyStr)))
+		}
+
+		// 响应加密时替换原有的响应体
+		var rbw *replaceBodyWriter
+		if responseEncrypt {
+			rbw = &replaceBodyWriter{
+				body:           &bytes.Buffer{},
+				ResponseWriter: c.Writer,
+			}
+			c.Writer = rbw
+		}
+
+		// 调用下一个处理程序
+		c.Next()
+
+		// 响应加密时对响应data数据进行加密
+		if responseEncrypt {
+			// 满足成功并带数据的响应进行加密
+			if c.Writer.Status() == 200 {
+				var resBody map[string]any
+				json.Unmarshal(rbw.body.Bytes(), &resBody)
+				codeV, codeOk := resBody["code"]
+				dataV, dataOk := resBody["data"]
+				if codeOk && dataOk {
+					if parse.Number(codeV) == constResult.CODE_SUCCESS {
+						byteBodyData, _ := json.Marshal(dataV)
+						// 加密
+						apiKey := config.Get("aes.apiKey").(string)
+						contentEn, err := crypto.AESEncryptBase64(string(byteBodyData), apiKey)
+						if err != nil {
+							logger.Errorf("CryptoApi encrypt err => %v", err)
+							rbw.ReplaceWrite([]byte(fmt.Sprintf(`{"code":"%d","msg":"encrypt err"}`, constResult.CODE_ERROR)))
+						} else {
+							// 响应加密
+							byteBody, _ := json.Marshal(map[string]any{
+								"code": constResult.CODE_ENCRYPT,
+								"msg":  constResult.MSG_ENCRYPT,
+								"data": contentEn,
+							})
+							rbw.ReplaceWrite(byteBody)
+						}
+					}
+				} else {
+					rbw.ReplaceWrite(nil)
+				}
+			} else {
+				rbw.ReplaceWrite(nil)
+			}
+		}
+		//
+	}
+}
+
+// replaceBodyWriter 替换默认的响应体
+type replaceBodyWriter struct {
+	gin.ResponseWriter
+	body *bytes.Buffer
+}
+
+// Write 写入响应体
+func (r replaceBodyWriter) Write(b []byte) (int, error) {
+	return r.body.Write(b)
+}
+
+// ReplaceWrite 替换响应体
+func (r *replaceBodyWriter) ReplaceWrite(b []byte) (int, error) {
+	if b == nil {
+		return r.ResponseWriter.Write(r.body.Bytes())
+	}
+	r.body = &bytes.Buffer{}
+	r.body.Write(b)
+	return r.ResponseWriter.Write(r.body.Bytes())
+}