Skip to content

Commit a29576f

Browse files
ZhengYa-0110SongZhen0704
ZhengYa-0110
authored and
SongZhen0704
committed
feat: encrypt agent cmd content
1 parent 61b6a57 commit a29576f

4 files changed

Lines changed: 102 additions & 16 deletions

File tree

‎server/controller/common/aes.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,14 @@ import (
2222
"crypto/aes"
2323
"crypto/cipher"
2424
"crypto/md5"
25+
"crypto/sha256"
2526
"encoding/base64"
2627
"errors"
2728
"fmt"
2829
"net"
2930
"os"
3031

32+
"golang.org/x/crypto/pbkdf2"
3133
"google.golang.org/grpc"
3234

3335
"github.com/deepflowio/deepflow/message/controller"
@@ -173,3 +175,8 @@ func getCAMD5() string {
173175
func GetCAMD5() string {
174176
return CAMD5
175177
}
178+
179+
func DerivePBKDF2Key(userID, orgID int) []byte {
180+
salt := fmt.Sprintf("%d", orgID) + "_deepflow"
181+
return pbkdf2.Key([]byte(fmt.Sprintf("%d", userID)), []byte(salt), 10000, 32, sha256.New)
182+
}

‎server/controller/common/aes_test.go‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package common
2+
3+
import (
4+
"fmt"
5+
"testing"
6+
)
7+
8+
func TestAesDecrypt(t *testing.T) {
9+
// Test data
10+
origData := "Hello, World!"
11+
key := "1234567890123456" // 16 bytes for AES-128
12+
13+
// Encrypt first
14+
encrypted, err := AesEncrypt(origData, key)
15+
if err != nil {
16+
t.Fatalf("AesEncrypt failed: %v", err)
17+
}
18+
fmt.Printf("Encrypted data: %s\n", encrypted)
19+
20+
// Now decrypt
21+
decrypted, err := AesDecrypt(encrypted, key)
22+
if err != nil {
23+
t.Fatalf("AesDecrypt failed: %v", err)
24+
}
25+
fmt.Printf("Decrypted data: %s\n", decrypted)
26+
27+
if decrypted != origData {
28+
t.Errorf("Expected %s, got %s", origData, decrypted)
29+
}
30+
}
31+
32+
func TestDerivePBKDF2Key(t *testing.T) {
33+
userID := 1
34+
orgID := 1
35+
36+
key := DerivePBKDF2Key(userID, orgID)
37+
fmt.Printf("Derived key: %x\n", key)
38+
39+
// Check length
40+
if len(key) != 32 {
41+
t.Errorf("Expected key length 32, got %d", len(key))
42+
}
43+
44+
// Test with different inputs
45+
key2 := DerivePBKDF2Key(2, 1)
46+
fmt.Printf("Derived key2: %x\n", key2)
47+
48+
if string(key) == string(key2) {
49+
t.Errorf("Keys should be different for different userIDs")
50+
}
51+
}

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

Lines changed: 43 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ package agent
1818

1919
import (
2020
"bytes"
21+
"encoding/json"
2122
"errors"
2223
"fmt"
2324
"io"
@@ -27,13 +28,12 @@ import (
2728
"strconv"
2829

2930
"github.com/gin-gonic/gin"
30-
"github.com/gin-gonic/gin/binding"
3131

3232
grpcapi "github.com/deepflowio/deepflow/message/agent"
3333
"github.com/deepflowio/deepflow/server/controller/common"
3434
"github.com/deepflowio/deepflow/server/controller/config"
35-
mysql "github.com/deepflowio/deepflow/server/controller/db/metadb"
36-
mysqlmodel "github.com/deepflowio/deepflow/server/controller/db/metadb/model"
35+
metadb "github.com/deepflowio/deepflow/server/controller/db/metadb"
36+
metadbmodel "github.com/deepflowio/deepflow/server/controller/db/metadb/model"
3737
httpcommon "github.com/deepflowio/deepflow/server/controller/http/common"
3838
"github.com/deepflowio/deepflow/server/controller/http/common/response"
3939
service "github.com/deepflowio/deepflow/server/controller/http/service/agent"
@@ -91,7 +91,7 @@ func (a *AgentCMD) RegisterTo(e *gin.Engine) {
9191
func forwardToServerConnectedByAgent() gin.HandlerFunc {
9292
return func(c *gin.Context) {
9393
orgID, _ := c.Get(common.HEADER_KEY_X_ORG_ID)
94-
db, err := mysql.GetDB(orgID.(int))
94+
db, err := metadb.GetDB(orgID.(int))
9595
if err != nil {
9696
log.Error(err, db.LogPrefixORGID)
9797
response.JSON(c, response.SetOptStatus(httpcommon.SERVER_ERROR), response.SetError(err))
@@ -105,7 +105,7 @@ func forwardToServerConnectedByAgent() gin.HandlerFunc {
105105
c.Abort()
106106
return
107107
}
108-
var agent *mysqlmodel.VTap
108+
var agent *metadbmodel.VTap
109109
if err = db.Where("id = ?", agentID).First(&agent).Error; err != nil {
110110
log.Error(err, db.LogPrefixORGID)
111111
response.JSON(c, response.SetOptStatus(httpcommon.SERVER_ERROR), response.SetError(err))
@@ -188,7 +188,7 @@ func forwardToServerConnectedByAgent() gin.HandlerFunc {
188188
func (a *AgentCMD) getCMDAndNamespaceHandler() gin.HandlerFunc {
189189
return func(c *gin.Context) {
190190
orgID, _ := c.Get(common.HEADER_KEY_X_ORG_ID)
191-
db, err := mysql.GetDB(orgID.(int))
191+
db, err := metadb.GetDB(orgID.(int))
192192
if err != nil {
193193
response.JSON(c, response.SetError(err))
194194
return
@@ -198,7 +198,7 @@ func (a *AgentCMD) getCMDAndNamespaceHandler() gin.HandlerFunc {
198198
response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(err))
199199
return
200200
}
201-
var agent *mysqlmodel.VTap
201+
var agent *metadbmodel.VTap
202202
if err = db.Where("id = ?", agentID).First(&agent).Error; err != nil {
203203
response.JSON(c, response.SetError(err))
204204
return
@@ -241,14 +241,14 @@ func (a *AgentCMD) getCMDAndNamespaceHandler() gin.HandlerFunc {
241241
}
242242
}
243243

244-
func getAgentID(c *gin.Context, db *mysql.DB) (int, error) {
244+
func getAgentID(c *gin.Context, db *metadb.DB) (int, error) {
245245
agentIDentStr := c.Param("id-or-name")
246246
if agentIDentStr == "" {
247247
return 0, errors.New("ident can not be empty")
248248
}
249249
agentID, err := strconv.Atoi(agentIDentStr)
250250
if err != nil {
251-
var agent mysqlmodel.VTap
251+
var agent metadbmodel.VTap
252252
if err := db.Where("name = ?", agentIDentStr).First(&agent).Error; err != nil {
253253
return 0, fmt.Errorf("failed to get agent by name(%s), error: %s", err.Error())
254254
}
@@ -259,8 +259,32 @@ func getAgentID(c *gin.Context, db *mysql.DB) (int, error) {
259259

260260
func (a *AgentCMD) cmdRunHandler() gin.HandlerFunc {
261261
return func(c *gin.Context) {
262+
userID, ok := c.Get(common.HEADER_KEY_X_USER_ID)
263+
if !ok {
264+
response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(fmt.Errorf("missing header %s", common.HEADER_KEY_X_USER_ID)))
265+
return
266+
}
267+
orgID, ok := c.Get(common.HEADER_KEY_X_ORG_ID)
268+
if !ok {
269+
response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(fmt.Errorf("missing header %s", common.HEADER_KEY_X_ORG_ID)))
270+
return
271+
}
272+
273+
cipherKey := string(common.DerivePBKDF2Key(userID.(int), orgID.(int)))
274+
rawPayload, err := io.ReadAll(c.Request.Body)
275+
if err != nil {
276+
response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(err))
277+
return
278+
}
279+
280+
decryptedPayload, err := common.AesDecrypt(string(rawPayload), cipherKey)
281+
if err != nil {
282+
response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(err))
283+
return
284+
}
285+
262286
req := service.RemoteExecReq{}
263-
if err := c.ShouldBindBodyWith(&req, binding.JSON); err != nil {
287+
if err := json.Unmarshal([]byte(decryptedPayload), &req); err != nil {
264288
response.JSON(c, response.SetOptStatus(httpcommon.INVALID_PARAMETERS), response.SetError(err))
265289
return
266290
}
@@ -283,8 +307,7 @@ func (a *AgentCMD) cmdRunHandler() gin.HandlerFunc {
283307
Params: req.Params,
284308
}
285309

286-
orgID, _ := c.Get(common.HEADER_KEY_X_ORG_ID)
287-
db, err := mysql.GetDB(orgID.(int))
310+
db, err := metadb.GetDB(orgID.(int))
288311
if err != nil {
289312
response.JSON(c, response.SetError(err))
290313
return
@@ -295,16 +318,21 @@ func (a *AgentCMD) cmdRunHandler() gin.HandlerFunc {
295318
return
296319
}
297320
content, err := service.RunAgentCMD(a.cfg.AgentCommandTimeout, orgID.(int), agentID, &agentReq, req.CMD)
321+
encryptedContent, encryptErr := common.AesEncrypt(content, cipherKey)
322+
if encryptErr != nil {
323+
response.JSON(c, response.SetData(content), response.SetOptStatus(httpcommon.SERVER_ERROR), response.SetError(encryptErr))
324+
return
325+
}
298326
if err != nil {
299-
response.JSON(c, response.SetData(content), response.SetOptStatus(httpcommon.SERVER_ERROR), response.SetError(err))
327+
response.JSON(c, response.SetData(encryptedContent), response.SetOptStatus(httpcommon.SERVER_ERROR), response.SetError(err))
300328
return
301329
}
302330

303331
if req.OutputFormat.String() == grpcapi.OutputFormat_TEXT.String() {
304-
response.JSON(c, response.SetData(content))
332+
response.JSON(c, response.SetData(encryptedContent))
305333
return
306334
}
307-
sendAsFile(c, req.OutputFilename, bytes.NewBuffer([]byte(content)))
335+
sendAsFile(c, req.OutputFilename, bytes.NewBuffer([]byte(encryptedContent)))
308336
}
309337
}
310338

‎server/go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -294,7 +294,7 @@ require (
294294
go.uber.org/goleak v1.3.0 // indirect
295295
go4.org/intern v0.0.0-20211027215823-ae77deb06f29 // indirect
296296
go4.org/unsafe/assume-no-moving-gc v0.0.0-20231121144256-b99613f794b6 // indirect
297-
golang.org/x/crypto v0.42.0 // indirect
297+
golang.org/x/crypto v0.42.0
298298
golang.org/x/oauth2 v0.30.0 // indirect
299299
golang.org/x/sync v0.17.0
300300
golang.org/x/term v0.35.0 // indirect

0 commit comments

Comments
 (0)