feat: encrypt agent cmd content

Author: ZhengYa-0110

server/controller/common/aes.go

@@ -22,12 +22,14 @@ import (
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/md5"
+	"crypto/sha256"
 	"encoding/base64"
 	"errors"
 	"fmt"
 	"net"
 	"os"
 
+	"golang.org/x/crypto/pbkdf2"
 	"google.golang.org/grpc"
 
 	"github.com/deepflowio/deepflow/message/controller"
@@ -164,3 +166,8 @@ func getCAMD5() string {
 func GetCAMD5() string {
 	return CAMD5
 }
+
+func DerivePBKDF2Key(userID, orgID int) []byte {
+	salt := fmt.Sprintf("%d", orgID) + "_deepflow"
+	return pbkdf2.Key([]byte(fmt.Sprintf("%d", userID)), []byte(salt), 10000, 32, sha256.New)
+}

server/controller/common/aes_test.go

@@ -0,0 +1,51 @@
+package common
+
+import (
+	"fmt"
+	"testing"
+)
+
+func TestAesDecrypt(t *testing.T) {
+	// Test data
+	origData := "Hello, World!"
+	key := "1234567890123456" // 16 bytes for AES-128
+
+	// Encrypt first
+	encrypted, err := AesEncrypt(origData, key)
+	if err != nil {
+		t.Fatalf("AesEncrypt failed: %v", err)
+	}
+	fmt.Printf("Encrypted data: %s\n", encrypted)
+
+	// Now decrypt
+	decrypted, err := AesDecrypt(encrypted, key)
+	if err != nil {
+		t.Fatalf("AesDecrypt failed: %v", err)
+	}
+	fmt.Printf("Decrypted data: %s\n", decrypted)
+
+	if decrypted != origData {
+		t.Errorf("Expected %s, got %s", origData, decrypted)
+	}
+}
+
+func TestDerivePBKDF2Key(t *testing.T) {
+	userID := 1
+	orgID := 1
+
+	key := DerivePBKDF2Key(userID, orgID)
+	fmt.Printf("Derived key: %x\n", key)
+
+	// Check length
+	if len(key) != 32 {
+		t.Errorf("Expected key length 32, got %d", len(key))
+	}
+
+	// Test with different inputs
+	key2 := DerivePBKDF2Key(2, 1)
+	fmt.Printf("Derived key2: %x\n", key2)
+
+	if string(key) == string(key2) {
+		t.Errorf("Keys should be different for different userIDs")
+	}
+}

server/controller/http/router/agent/agent_cmd.go

@@ -18,6 +18,7 @@ package agent
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -27,7 +28,6 @@ import (
 	"strconv"
 
 	"github.com/gin-gonic/gin"
-	"github.com/gin-gonic/gin/binding"
 
 	grpcapi "github.com/deepflowio/deepflow/message/agent"
 	"github.com/deepflowio/deepflow/server/controller/common"
@@ -263,8 +263,32 @@ func getAgentID(c *gin.Context, db *metadb.DB) (int, error) {
 
 func (a *AgentCMD) cmdRunHandler() gin.HandlerFunc {
 	return func(c *gin.Context) {
+		userID, ok := c.Get(common.HEADER_KEY_X_USER_ID)
+		if !ok {
+			response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(fmt.Errorf("missing header %s", common.HEADER_KEY_X_USER_ID)))
+			return
+		}
+		orgID, ok := c.Get(common.HEADER_KEY_X_ORG_ID)
+		if !ok {
+			response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(fmt.Errorf("missing header %s", common.HEADER_KEY_X_ORG_ID)))
+			return
+		}
+
+		cipherKey := string(common.DerivePBKDF2Key(userID.(int), orgID.(int)))
+		rawPayload, err := io.ReadAll(c.Request.Body)
+		if err != nil {
+			response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(err))
+			return
+		}
+
+		decryptedPayload, err := common.AesDecrypt(string(rawPayload), cipherKey)
+		if err != nil {
+			response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(err))
+			return
+		}
+
 		req := service.RemoteExecReq{}
-		if err := c.ShouldBindBodyWith(&req, binding.JSON); err != nil {
+		if err := json.Unmarshal([]byte(decryptedPayload), &req); err != nil {
 			response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(err))
 			return
 		}
@@ -287,7 +311,6 @@ func (a *AgentCMD) cmdRunHandler() gin.HandlerFunc {
 			Params:       req.Params,
 		}
 
-		orgID, _ := c.Get(common.HEADER_KEY_X_ORG_ID)
 		db, err := metadb.GetDB(orgID.(int))
 		if err != nil {
 			response.JSON(c, response.SetError(err))
@@ -299,16 +322,21 @@ func (a *AgentCMD) cmdRunHandler() gin.HandlerFunc {
 			return
 		}
 		content, err := service.RunAgentCMD(a.cfg.AgentCommandTimeout, orgID.(int), agentID, &agentReq, req.CMD)
+		encryptedContent, encryptErr := common.AesEncrypt(content, cipherKey)
+		if encryptErr != nil {
+			response.JSON(c, response.SetData(content), response.SetOptStatus(httpcommon.SERVER_ERROR), response.SetError(encryptErr))
+			return
+		}
 		if err != nil {
-			response.JSON(c, response.SetData(content), response.SetOptStatus(httpcommon.SERVER_ERROR), response.SetError(err))
+			response.JSON(c, response.SetData(encryptedContent), response.SetOptStatus(httpcommon.SERVER_ERROR), response.SetError(err))
 			return
 		}
 
 		if req.OutputFormat.String() == grpcapi.OutputFormat_TEXT.String() {
-			response.JSON(c, response.SetData(content))
+			response.JSON(c, response.SetData(encryptedContent))
 			return
 		}
-		sendAsFile(c, req.OutputFilename, bytes.NewBuffer([]byte(content)))
+		sendAsFile(c, req.OutputFilename, bytes.NewBuffer([]byte(encryptedContent)))
 	}
 }
 

server/go.mod

@@ -138,6 +138,7 @@ require (
 	go.opentelemetry.io/otel/sdk v1.43.0
 	go.opentelemetry.io/otel/trace v1.43.0
 	go.opentelemetry.io/proto/otlp v1.1.0
+	golang.org/x/crypto v0.42.0
 	golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b
 	golang.org/x/net v0.43.0
 	golang.org/x/sync v0.17.0
@@ -305,7 +306,6 @@ require (
 	go4.org/intern v0.0.0-20211027215823-ae77deb06f29 // indirect
 	go4.org/unsafe/assume-no-moving-gc v0.0.0-20231121144256-b99613f794b6 // indirect
 	golang.org/x/arch v0.20.0 // indirect
-	golang.org/x/crypto v0.42.0 // indirect
 	golang.org/x/mod v0.27.0 // indirect
 	golang.org/x/oauth2 v0.30.0 // indirect
 	golang.org/x/term v0.35.0 // indirect